From 7b1c0b15eae2836135486f4325462bd3822aa711 Mon Sep 17 00:00:00 2001 From: Alphaxiaoteng <230277249+Alphaxiaoteng@users.noreply.github.com> Date: Mon, 31 Aug 2026 06:26:50 +0800 Subject: [PATCH] Preserve pretraining epoch across resume --- litgpt/pretrain.py | 8 +++- litgpt/utils.py | 26 ++++++++++- tests/test_pretrain.py | 98 +++++++++++++++++++++++++++++++++++++++++- tests/test_utils.py | 9 ++++ 4 files changed, 135 insertions(+), 6 deletions(-) diff --git a/litgpt/pretrain.py b/litgpt/pretrain.py index 35cf0bcb9a..01beeaf061 100644 --- a/litgpt/pretrain.py +++ b/litgpt/pretrain.py @@ -230,11 +230,15 @@ def main( "iter_num": 0, "step_count": 0, } + train_iterator = CycleIterator(train_dataloader) resume = find_resume_path(resume, out_dir) if resume: fabric.print(f"Resuming training from {resume}") - fabric.load(resume, state) + checkpoint = fabric.load(resume, state) + if "train_iterator" in checkpoint: + train_iterator.load_state_dict(checkpoint["train_iterator"]) + state["train_iterator"] = train_iterator train_time = time.perf_counter() @@ -323,7 +327,7 @@ def fit( max_iters = max_tokens_per_device // tokens_per_iter log_iter_interval = train.log_interval * train.gradient_accumulation_iters(devices, num_nodes) initial_iter = state["iter_num"] - train_iterator = CycleIterator(train_dataloader) + train_iterator = state.get("train_iterator") or CycleIterator(train_dataloader) running_loss = RunningMean(window=train.gradient_accumulation_iters(devices, num_nodes), sync_on_compute=False).to( fabric.device diff --git a/litgpt/utils.py b/litgpt/utils.py index ec79fa0764..a924f0393b 100644 --- a/litgpt/utils.py +++ b/litgpt/utils.py @@ -495,20 +495,42 @@ def __init__(self, iterable: Iterable) -> None: self.iterable = iterable self.epoch = 0 self._iterator = None + self._num_yielded = 0 def __next__(self) -> Any: if self._iterator is None: self._iterator = iter(self.iterable) try: - return next(self._iterator) + item = next(self._iterator) except StopIteration: self._iterator = iter(self.iterable) self.epoch += 1 - return next(self._iterator) + self._num_yielded = 0 + item = next(self._iterator) + self._num_yielded += 1 + return item def __iter__(self) -> Self: return self + def state_dict(self) -> dict[str, int]: + epoch = self.epoch + try: + iterable_length = len(self.iterable) + is_epoch_complete = ( + iterable_length > 0 and self._iterator is not None and self._num_yielded == iterable_length + ) + except TypeError: + is_epoch_complete = False + if is_epoch_complete: + epoch += 1 + return {"epoch": epoch} + + def load_state_dict(self, state_dict: dict[str, int]) -> None: + self.epoch = state_dict["epoch"] + self._iterator = None + self._num_yielded = 0 + def copy_config_files(source_dir: Path, out_dir: Path) -> None: """Copies the specified configuration and tokenizer files into the output directory.""" diff --git a/tests/test_pretrain.py b/tests/test_pretrain.py index 2461fb1836..99ee873fbf 100644 --- a/tests/test_pretrain.py +++ b/tests/test_pretrain.py @@ -1,8 +1,9 @@ # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. import os -from contextlib import redirect_stdout +from contextlib import nullcontext, redirect_stdout from io import StringIO +from types import SimpleNamespace from unittest import mock from unittest.mock import ANY, Mock @@ -15,7 +16,11 @@ from litgpt.args import EvalArgs, TrainArgs from litgpt.config import Config from litgpt.pretrain import initialize_weights -from litgpt.utils import _RunIf +from litgpt.utils import CycleIterator, _RunIf + + +def _identity(item): + return item @_RunIf(min_cuda_gpus=1, standalone=True) @@ -129,3 +134,92 @@ def __init__(self): initialize_weights(fabric_mock, model, n_layer=2, n_embd=8) assert model.reset_parameters.call_count == int(expected) assert model.child.reset_parameters.call_count == int(expected) + + +@pytest.mark.parametrize(("num_consumed", "num_resumed"), [(20, 12), (8, 9)]) +def test_stateful_dataloader_resume_preserves_cursor_and_cycle_epoch(tmp_path, num_consumed, num_resumed): + from lightning import Fabric + from litdata import optimize + from litdata.streaming import StreamingDataLoader, StreamingDataset + + optimize(fn=_identity, inputs=list(range(8)), output_dir=tmp_path, num_workers=1, chunk_size=1) + + def create_dataloader(): + dataset = StreamingDataset(input_dir=tmp_path, shuffle=False, drop_last=False) + return StreamingDataLoader(dataset, batch_size=1, num_workers=0) + + fabric = Fabric(accelerator="cpu", devices=1) + train_dataloader = fabric.setup_dataloaders(create_dataloader()) + train_iterator = CycleIterator(train_dataloader) + + assert [next(train_iterator).item() for _ in range(num_consumed)] == [i % 8 for i in range(num_consumed)] + assert train_iterator.epoch == (num_consumed - 1) // 8 + assert train_iterator.state_dict() == {"epoch": num_consumed // 8} + + checkpoint_path = tmp_path / "checkpoint.pt" + fabric.save( + checkpoint_path, + {"train_dataloader": train_dataloader, "train_iterator": train_iterator}, + ) + + resumed_dataloader = fabric.setup_dataloaders(create_dataloader()) + resumed_iterator = CycleIterator(resumed_dataloader) + checkpoint = fabric.load(checkpoint_path, {"train_dataloader": resumed_dataloader}) + resumed_iterator.load_state_dict(checkpoint["train_iterator"]) + + assert [next(resumed_iterator).item() for _ in range(num_resumed)] == [ + (num_consumed + i) % 8 for i in range(num_resumed) + ] + assert resumed_iterator.epoch == (num_consumed + num_resumed - 1) // 8 + + +@pytest.mark.parametrize(("checkpoint", "expected_epoch"), [({}, 0), ({"train_iterator": {"epoch": 2}}, 2)]) +def test_main_restores_train_iterator(monkeypatch, tmp_path, checkpoint, expected_epoch): + model = torch.nn.Linear(1, 1) + model.config = SimpleNamespace(n_layer=1, n_embd=1) + model.max_seq_length = 2 + dataloader = DataLoader(torch.tensor([[0, 1, 2]])) + + fabric = mock.Mock() + fabric.device = torch.device("cpu") + fabric.global_rank = 0 + fabric.world_size = 1 + fabric.init_module.return_value = nullcontext() + fabric.setup.side_effect = lambda module: module + fabric.setup_optimizers.side_effect = lambda optimizer: optimizer + fabric.setup_dataloaders.side_effect = lambda *dataloaders: dataloaders + + def load(_, state): + assert "train_iterator" not in state + state["iter_num"] = 1 + return checkpoint + + fabric.load.side_effect = load + fit = mock.Mock() + save_checkpoint = mock.Mock() + monkeypatch.setattr(pretrain, "GPT", lambda _: model) + monkeypatch.setattr(pretrain, "fit", fit) + monkeypatch.setattr(pretrain, "get_dataloaders", mock.Mock(return_value=(dataloader, dataloader))) + monkeypatch.setattr(pretrain, "initialize_weights", mock.Mock()) + monkeypatch.setattr(pretrain, "save_checkpoint", save_checkpoint) + monkeypatch.setattr(pretrain.torch, "compile", lambda module: module) + + pretrain.main( + fabric=fabric, + devices=1, + seed=42, + initial_checkpoint_dir=None, + resume=tmp_path / "checkpoint.pt", + config=model.config, + data=mock.Mock(), + out_dir=tmp_path / "out", + tokenizer_dir=None, + tokenizer=None, + train=TrainArgs(global_batch_size=1, max_tokens=2, micro_batch_size=1, max_norm=1.0), + eval=EvalArgs(final_validation=False), + optimizer="SGD", + ) + + train_iterator = fit.call_args.kwargs["state"]["train_iterator"] + assert train_iterator.epoch == expected_epoch + assert save_checkpoint.call_args.args[1]["train_iterator"] is train_iterator diff --git a/tests/test_utils.py b/tests/test_utils.py index ae10780826..81f776fed7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -194,6 +194,7 @@ def test_cycle_iterator(): iterator = CycleIterator([]) with pytest.raises(StopIteration): next(iterator) + assert iterator.state_dict() == {"epoch": 1} iterator = CycleIterator(range(3)) assert iterator.epoch == 0 @@ -203,6 +204,14 @@ def test_cycle_iterator(): assert iterator.epoch == 0 assert next(iterator) == 2 assert iterator.epoch == 0 + assert iterator.state_dict() == {"epoch": 1} + + resumed_iterator = CycleIterator(range(3)) + assert next(resumed_iterator) == 0 + resumed_iterator.load_state_dict(iterator.state_dict()) + assert resumed_iterator.epoch == 1 + assert next(resumed_iterator) == 0 + assert next(iterator) == 0 assert iterator.epoch == 1