diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index 2c63fe312..fb72b3ec7 100644 --- a/src/speculators/train/trainer.py +++ b/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) diff --git a/tests/unit/train/test_mid_epoch_resume.py b/tests/unit/train/test_mid_epoch_resume.py index 8421f9410..a797b12b3 100644 --- a/tests/unit/train/test_mid_epoch_resume.py +++ b/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}