Add step-level performance profiling to the trainer (#722)

<!-- markdownlint-disable -->

<!-- PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS
(AT THE BOTTOM) HAVE BEEN CONSIDERED. -->

## Purpose
Add step-level performance profiling to the trainer (resolves #717,
trainer portion).
Emits 7 new metrics through the existing `speculators.metrics` logger:
`profile/fetch_ms`, `profile/fwd_ms`, `profile/bwd_ms`,
`profile/opt_ms`, `profile/step_ms`, `profile/tokens_per_s`,
`profile/fetch_frac`.
## Tests
Add unit tests, can be run with:
```bash
python3 -m pytest tests/unit/train/test_step_timer.py -v
```
## Checklist

I have filled in:

- [x] The purpose of the PR, such as "Fix some issue (link existing
issues this PR will resolve)".
- [x] The test plan/results, such as providing test command and pasting
the results.
- [ ] (Optional) The necessary documentation update.
- [x] I (a human) have written or reviewed the code in this pr to the
best of my ability.

Signed-off-by: Omer Aplatony <omerap12@gmail.com>

Author: omerap12

src/speculators/train/trainer.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import time
 import warnings
 from pathlib import Path
 from typing import Literal, NamedTuple
@@ -37,6 +38,57 @@
 root_logger = logging.getLogger("speculators")
 metric_logger = logging.getLogger("speculators.metrics")
 
+
+class _StepTimer:
+    # Each mark()/now() forces a cuda.synchronize to capture true GPU time.
+    # This serialises the CUDA pipeline, so profiled steps are slower; keep
+    # log_freq > 1 in perf-sensitive runs.
+    def __init__(self, enabled: bool = False):
+        self.enabled = enabled
+        self._marks: dict[str, float] = {}
+
+    def reset(self, enabled: bool) -> None:
+        self.enabled = enabled
+        self._marks.clear()
+
+    def mark(self, name: str) -> None:
+        if self.enabled:
+            torch.cuda.synchronize()
+            self._marks[name] = time.perf_counter()
+
+    def mark_value(self, name: str, value: float) -> None:
+        if self.enabled:
+            self._marks[name] = value
+
+    def now(self) -> float | None:
+        if not self.enabled:
+            return None
+        torch.cuda.synchronize()
+        return time.perf_counter()
+
+    def profile(self, num_tokens: int) -> dict[str, float] | None:
+        if not self.enabled:
+            return None
+        m = self._marks
+        has_start = "start" in m
+        fwd_ms = (m["fwd"] - m["fetch"]) * 1000
+        bwd_ms = (m["bwd"] - m["fwd"]) * 1000
+        opt_ms = (m["opt"] - m["bwd"]) * 1000
+        fetch_ms = (m["fetch"] - m["start"]) * 1000 if has_start else 0.0
+        step_ms = (m["opt"] - m["start"]) * 1000 if has_start else 0.0
+        tokens_per_s = num_tokens / (step_ms / 1000) if step_ms > 0 else 0.0
+        fetch_frac = fetch_ms / step_ms if step_ms > 0 else 0.0
+        return {
+            "fetch_ms": fetch_ms,
+            "fwd_ms": fwd_ms,
+            "bwd_ms": bwd_ms,
+            "opt_ms": opt_ms,
+            "step_ms": step_ms,
+            "tokens_per_s": tokens_per_s,
+            "fetch_frac": fetch_frac,
+        }
+
+
 warnings.filterwarnings("ignore", category=TqdmExperimentalWarning)
 MIN_STEP_PCT = 0.25
 
@@ -326,31 +378,45 @@ def train_epoch(self, epoch: int):
             if self.config.checkpoint_freq < 1
             else None
         )
+        t_before_fetch = time.perf_counter()
+        timer = _StepTimer()
         for local_step_rel, batch in enumerate(train_loader, 1):
             # local_step is 1-based index into the *full* epoch (not the slice).
             local_step = local_step_rel + skip_steps
+            timer.reset(self.global_step % self.config.log_freq == 0)
+
+            timer.mark_value("start", t_before_fetch)
             gpu_batch = {
                 k: v.to(self.local_rank, non_blocking=True)
                 if isinstance(v, torch.Tensor)
                 else v
                 for k, v in batch.items()
             }
 
+            timer.mark("fetch")
             _draft_tokens, loss, metrics = self.model(
                 **gpu_batch, **self.config.train_call_kwargs
             )
 
+            timer.mark("fwd")
             self._optimizers_zero_grad()
             loss.backward()
             torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+
+            timer.mark("bwd")
             self._optimizers_step()
 
             current_lrs = {
                 type(opt).__name__: opt.param_groups[0]["lr"] for opt in self.optimizers
             }
             self._schedulers_step()
+            timer.mark("opt")
+            t_before_fetch = timer.now() or time.perf_counter()
 
-            if self.global_step % self.config.log_freq == 0:
+            profile = None
+            if timer.enabled:
+                num_tokens = int((gpu_batch["document_ids"] != -1).sum().item())
+                profile = timer.profile(num_tokens)
                 if self.is_distributed:
                     for v in metrics.values():
                         dist.reduce(v, dst=0, op=dist.ReduceOp.SUM)
@@ -366,6 +432,7 @@ def train_epoch(self, epoch: int):
                 metric_logger.info(
                     {
                         "train": metrics,
+                        "profile": profile,
                         "epoch": epoch,
                         "lr": lr_info,
                         "global_step": self.global_step,

tests/unit/train/test_step_timer.py

@@ -0,0 +1,119 @@
+from unittest.mock import patch
+
+import pytest
+
+from speculators.train.trainer import _StepTimer
+
+
+def test_disabled_timer_returns_none():
+    timer = _StepTimer(enabled=False)
+    timer.mark_value("start", 0.0)
+    timer.mark("fetch")
+    timer.mark("fwd")
+    timer.mark("bwd")
+    timer.mark("opt")
+    assert timer.now() is None
+    assert timer.profile(num_tokens=1024) is None
+
+
+@patch("speculators.train.trainer.torch.cuda.synchronize")
+def test_enabled_timer_returns_profile(mock_sync):
+    timer = _StepTimer(enabled=True)
+
+    with patch(
+        "speculators.train.trainer.time.perf_counter",
+        side_effect=[
+            0.1,
+            0.3,
+            0.5,
+            0.6,
+            0.6,
+        ],
+    ):
+        timer.mark_value("start", 0.0)
+        timer.mark("fetch")
+        timer.mark("fwd")
+        timer.mark("bwd")
+        timer.mark("opt")
+        t_next = timer.now()
+
+    assert mock_sync.call_count == 5
+    assert t_next == 0.6
+
+    profile = timer.profile(num_tokens=4096)
+    assert profile is not None
+    assert profile["fetch_ms"] == (0.1 - 0.0) * 1000
+    assert profile["fwd_ms"] == (0.3 - 0.1) * 1000
+    assert profile["bwd_ms"] == (0.5 - 0.3) * 1000
+    assert profile["opt_ms"] == (0.6 - 0.5) * 1000
+    assert profile["step_ms"] == (0.6 - 0.0) * 1000
+    assert profile["tokens_per_s"] == 4096 / 0.6
+    assert profile["fetch_frac"] == 100 / 600
+
+
+def test_disabled_to_enabled_transition():
+    """Simulate log_freq=2: a disabled step followed by an enabled step.
+
+    The disabled step's ``timer.now()`` returns None so the training loop
+    falls back to ``time.perf_counter()`` for ``t_before_fetch``.  The
+    subsequent enabled step feeds that value into ``mark_value("start", ...)``.
+    Verify the profile is valid with a realistic ``start`` mark.
+    """
+    timer = _StepTimer()
+
+    # --- disabled step (global_step=1, log_freq=2 → 1%2 != 0) ---
+    timer.reset(enabled=False)
+    timer.mark_value("start", 1.0)
+    timer.mark("fetch")
+    timer.mark("fwd")
+    timer.mark("bwd")
+    timer.mark("opt")
+    assert timer.now() is None
+    assert timer.profile(num_tokens=512) is None
+
+    # Simulate the fallback: t_before_fetch = timer.now() or time.perf_counter()
+    t_before_fetch = 2.0  # stands in for the perf_counter() fallback
+
+    # --- enabled step (global_step=2, log_freq=2 → 2%2 == 0) ---
+    timer.reset(enabled=True)
+    timer.mark_value("start", t_before_fetch)
+
+    with (
+        patch("speculators.train.trainer.torch.cuda.synchronize"),
+        patch(
+            "speculators.train.trainer.time.perf_counter",
+            side_effect=[2.1, 2.4, 2.5, 2.6, 2.6],
+        ),
+    ):
+        timer.mark("fetch")
+        timer.mark("fwd")
+        timer.mark("bwd")
+        timer.mark("opt")
+        t_next = timer.now()
+
+    assert t_next == 2.6
+    profile = timer.profile(num_tokens=2048)
+    assert profile is not None
+    assert profile["fetch_ms"] == (2.1 - 2.0) * 1000
+    assert profile["fwd_ms"] == (2.4 - 2.1) * 1000
+    assert profile["bwd_ms"] == (2.5 - 2.4) * 1000
+    assert profile["opt_ms"] == (2.6 - 2.5) * 1000
+    assert profile["step_ms"] == (2.6 - 2.0) * 1000
+    assert profile["tokens_per_s"] == pytest.approx(2048 / 0.6)
+
+
+def test_zero_step_ms_returns_zero_throughput():
+    timer = _StepTimer(enabled=True)
+    timer.mark_value("start", 1.0)
+    with (
+        patch("speculators.train.trainer.torch.cuda.synchronize"),
+        patch("speculators.train.trainer.time.perf_counter", return_value=1.0),
+    ):
+        timer.mark("fetch")
+        timer.mark("fwd")
+        timer.mark("bwd")
+        timer.mark("opt")
+    profile = timer.profile(num_tokens=4096)
+    assert profile is not None
+    assert profile["tokens_per_s"] == 0.0
+    assert profile["fetch_frac"] == 0.0