fix(train): show resumed position in epoch progress bar

src/speculators/train/trainer.py

@@ -382,39 +382,41 @@ def _schedulers_step(self):
         for scheduler in self.schedulers:
             scheduler.step()
 
-    def _prepare_resume_skip(self, epoch: int) -> int:
-        """Prepare fast-skip state for mid-epoch resume and return skipped steps."""
-        skip_steps = 0
+    def _prepare_resume_skip(self, epoch: int) -> tuple[int, int]:
+        """Prepare resume state and return its local offset and fast-skip count."""
+        resume_local_step = 0
         if epoch == getattr(self, "current_epoch", epoch):
-            skip_steps = getattr(self, "_resume_local_step", 0)
+            resume_local_step = getattr(self, "_resume_local_step", 0)
             # Only skip once — clear after use.
             self._resume_local_step = 0
 
         # Fast-skip: slice the sampler's pre-generated batch list so we never
         # call __getitem__ (and thus never call vLLM) for skipped batches.
+        fast_skipped_steps = 0
         sampler = self.train_loader.batch_sampler
         has_fast_skip_api = hasattr(sampler, "_generate_batches") and hasattr(
             sampler, "_cached_generated_batches"
         )
-        if skip_steps > 0 and has_fast_skip_api:
+        if resume_local_step > 0 and has_fast_skip_api:
             all_batches = sampler._generate_batches(epoch)  # type: ignore[union-attr]  # noqa: SLF001
-            remaining = all_batches[skip_steps:]
+            remaining = all_batches[resume_local_step:]
+            fast_skipped_steps = len(all_batches) - len(remaining)
             # Temporarily override the sampler cache with the sliced list.
             sampler._cached_generated_batches = (  # type: ignore[union-attr]  # noqa: SLF001
                 epoch,
                 remaining,
             )
             root_logger.info(
-                f"Fast-skipping {skip_steps} batches via sampler slice "
+                f"Fast-skipping {fast_skipped_steps} batches via sampler slice "
                 f"(no vLLM calls for skipped batches). "
                 f"epoch={epoch}, global_step={self.global_step}."
             )
-        elif skip_steps > 0:
+        elif resume_local_step > 0:
             root_logger.warning(
                 "Sampler lacks fast-skip API; resume will replay "
-                f"{skip_steps} batches from the start of the epoch."
+                f"{resume_local_step} batches from the start of the epoch."
             )
-        return skip_steps
+        return resume_local_step, fast_skipped_steps
 
     def train_epoch(self, epoch: int):
         self.model.train()
@@ -425,11 +427,16 @@ def train_epoch(self, epoch: int):
         num_steps = len(self.train_loader)
 
         # Determine how many batches to skip for mid-epoch resume.
-        skip_steps = self._prepare_resume_skip(epoch)
+        resume_local_step, fast_skipped_steps = self._prepare_resume_skip(epoch)
 
         train_loader = self.train_loader
         if self.rank == 0:
-            train_loader = tqdm(train_loader, desc=f"Epoch {epoch}")  # type: ignore[assignment]
+            train_loader = tqdm(  # type: ignore[assignment]
+                train_loader,
+                desc=f"Epoch {epoch}",
+                total=num_steps,
+                initial=fast_skipped_steps,
+            )
 
         step_interval = (
             max(1, round(num_steps * self.config.checkpoint_freq))
@@ -440,7 +447,7 @@ def train_epoch(self, epoch: int):
         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
+            local_step = local_step_rel + resume_local_step
             timer.reset(self.global_step % self.config.log_freq == 0)
 
             timer.mark_value("start", t_before_fetch)

tests/unit/train/test_mid_epoch_resume.py

@@ -348,6 +348,61 @@ def __iter__(self):
         yield from self._generate_batches(self.current_epoch)
 
 
+class _ReplayLoader:
+    """Empty loader with a known epoch length and no fast-skip sampler API."""
+
+    batch_sampler = object()
+
+    def __init__(self, num_steps: int):
+        self.num_steps = num_steps
+
+    def __len__(self) -> int:
+        return self.num_steps
+
+    def __iter__(self):
+        return iter(())
+
+
+def _capture_progress_kwargs(monkeypatch: pytest.MonkeyPatch) -> dict[str, object]:
+    captured: dict[str, object] = {}
+
+    class _Progress:
+        def __init__(self, iterable, **kwargs) -> None:
+            captured.update(kwargs)
+            self.iterable = iterable
+
+        def __iter__(self):
+            return iter(self.iterable)
+
+    monkeypatch.setattr("speculators.train.trainer.tqdm", _Progress)
+    return captured
+
+
+def _make_progress_test_trainer(
+    loader: object,
+    tmp_path: Path,
+    resume_local_step: int,
+) -> Trainer:
+    config = TrainerConfig(
+        save_path=str(tmp_path),
+        num_epochs=1,
+        lr=1e-4,
+        resume_from_checkpoint=False,
+        checkpoint_freq=0.3,
+        log_freq=1,
+        scheduler_type="none",
+    )
+    trainer = Trainer.__new__(Trainer)
+    trainer.model = _dummy_model()
+    trainer.train_loader = loader  # type: ignore[assignment]
+    trainer.rank = 0
+    trainer.current_epoch = 0
+    trainer._resume_local_step = resume_local_step
+    trainer.global_step = resume_local_step
+    trainer.config = config
+    return trainer
+
+
 class _FastSkipMockTrainer(_MockTrainer):
     def train_epoch(self, epoch: int) -> None:
         if hasattr(self.train_loader.batch_sampler, "set_epoch"):
@@ -403,3 +458,34 @@ def test_fast_skip_sampler_slice_avoids_skipped_getitem(
         assert sampler.generated_for_epoch == 0
         assert sampler._cached_generated_batches == (0, sampler.all_batches[3:])
         assert dataset.seen_indices == list(range(3, 10))
+
+
+def test_fast_skip_progress_bar_starts_at_resume_step(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    """Rank-zero progress reflects already skipped mid-epoch batches."""
+    dataset = _CountingDataset(n_items=3)
+    sampler = _FastSkipBatchSampler(n_items=3)
+    loader = DataLoader(dataset, batch_sampler=sampler)
+    trainer = _make_progress_test_trainer(loader, tmp_path, resume_local_step=3)
+    captured = _capture_progress_kwargs(monkeypatch)
+
+    trainer.train_epoch(0)
+
+    assert captured == {"desc": "Epoch 0", "total": 3, "initial": 3}
+
+
+def test_replay_progress_bar_starts_at_zero(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    """Fallback replay must not count unskipped batches as completed."""
+    trainer = _make_progress_test_trainer(
+        _ReplayLoader(num_steps=3), tmp_path, resume_local_step=3
+    )
+    captured = _capture_progress_kwargs(monkeypatch)
+
+    trainer.train_epoch(0)
+
+    assert captured == {"desc": "Epoch 0", "total": 3, "initial": 0}