Merge pull request #52 from KempnerInstitute/dataloader-state-persistence

Persist dataloader state across checkpoints

Author: mmshad

kempnerforge/checkpoint/manager.py

@@ -60,6 +60,10 @@ def __init__(
         self._async_ckpt = AsyncCheckpointer(mode=config.async_mode)
         self._process_group = process_group
         self._pp_rank = pp_rank
+        # Dataloader state stashed during load() when the caller cannot yet
+        # provide a dataloader object. Applied later via
+        # apply_dataloader_state() once the loader is constructed.
+        self._pending_dataloader_state: dict[str, Any] | None = None
 
     def _checkpoint_dir(self, step: int) -> Path:
         return self.base_dir / f"step_{step}"
@@ -212,6 +216,11 @@ def load(
                 train_state = object_list[0]
 
             assert train_state is not None, "train_state broadcast failed"
+            # Stash dataloader state if the caller can't yet provide the loader
+            # object. Training loops construct the dataloader after load() so
+            # apply_dataloader_state() can restore it once it exists.
+            if dataloader is None and "dataloader" in train_state:
+                self._pending_dataloader_state = train_state["dataloader"]
             step, tokens_seen, extra = restore_train_state(
                 train_state,
                 scheduler=scheduler,
@@ -222,6 +231,25 @@ def load(
 
         return 0, 0, {}
 
+    def apply_dataloader_state(self, dataloader: Any) -> None:
+        """Apply any dataloader state stashed during load().
+
+        Training loops call load() before constructing the dataloader (since
+        the dataloader depends on phase/annealing state that load() restores).
+        This method applies the stashed state once the loader exists.
+
+        No-op if no state is pending, or if the loader does not support
+        ``load_state_dict`` (e.g., plain torch DataLoader for HF streaming).
+        """
+        if self._pending_dataloader_state is None:
+            return
+        if dataloader is None or not hasattr(dataloader, "load_state_dict"):
+            self._pending_dataloader_state = None
+            return
+        dataloader.load_state_dict(self._pending_dataloader_state)
+        self._pending_dataloader_state = None
+        logger.info("Applied stashed dataloader state")
+
     def _resolve_load_path(self, path: str | None = None) -> Path | None:
         """Resolve the checkpoint path to load from."""
         if path is not None:

scripts/train.py

@@ -376,6 +376,13 @@ def main() -> None:
                 f"{config.data.hf_dataset_name} ({config.data.hf_dataset_split})"
             )
 
+    # Apply any dataloader state stashed during load(). Runs after dataloader
+    # construction because the loader's identity depends on phase scheduling
+    # that load() restores. No-op when resuming without a prior dataloader
+    # state or when the loader is not stateful (plain TorchDataLoader).
+    if dataloader is not None:
+        ckpt_mgr.apply_dataloader_state(dataloader)
+
     # --- Eval data ---
     eval_config = config.eval
     eval_dataloader = None
@@ -763,6 +770,7 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
                 step=step,
                 tokens_seen=tokens_seen,
                 scheduler=scheduler,
+                dataloader=dataloader,
                 extra=ckpt_extra,
             )
             hook_runner.on_checkpoint_save(step, config.checkpoint.dir)
@@ -774,6 +782,7 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
                 step=step,
                 tokens_seen=tokens_seen,
                 scheduler=scheduler,
+                dataloader=dataloader,
                 extra=ckpt_extra,
             )
             shutdown_handler.finish()

tests/unit/test_checkpoint.py

@@ -254,3 +254,174 @@ def test_save_waits_for_previous(self, monkeypatch, tmp_path):
         ckpt.save({"model": {}}, checkpoint_id=str(tmp_path / "step_2"))
         # First future should have been waited on before second save
         mock_future1.result.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# Dataloader state persistence (two-phase apply)
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_mgr(tmp_path, monkeypatch):
+    """Build a CheckpointManager with DCP calls mocked out (no distributed)."""
+    from unittest.mock import MagicMock
+
+    from kempnerforge.checkpoint.manager import CheckpointManager
+
+    model = torch.nn.Linear(4, 4)
+    opt = torch.optim.SGD(model.parameters(), lr=0.1)
+    config = CheckpointConfig(dir=str(tmp_path), keep_last_n=5)
+    mgr = CheckpointManager(config, model, opt)
+    monkeypatch.setattr(mgr._async_ckpt, "save", MagicMock())
+    monkeypatch.setattr("kempnerforge.checkpoint.manager.dcp.load", MagicMock())
+    return mgr
+
+
+class TestDataloaderStatePersistence:
+    """Round-trip coverage for dataloader state across save -> load -> apply.
+
+    Training loops call load() before constructing the dataloader (the loader
+    depends on phase/annealing state that load() restores). Load stashes the
+    dataloader state; apply_dataloader_state() restores it into the freshly
+    built loader.
+    """
+
+    def test_apply_no_op_when_nothing_pending(self, tmp_path, monkeypatch):
+        from unittest.mock import MagicMock
+
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+        loader = MagicMock(spec=["load_state_dict"])
+        mgr.apply_dataloader_state(loader)
+        loader.load_state_dict.assert_not_called()
+
+    def test_apply_restores_state_to_stateful_loader(self, tmp_path, monkeypatch):
+        from unittest.mock import MagicMock
+
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+        stashed = {"epoch": 3, "batches_yielded": 100, "sampler": {"epoch": 3, "skip_samples": 0}}
+        mgr._pending_dataloader_state = stashed
+
+        loader = MagicMock(spec=["load_state_dict"])
+        mgr.apply_dataloader_state(loader)
+
+        loader.load_state_dict.assert_called_once_with(stashed)
+        assert mgr._pending_dataloader_state is None
+
+    def test_apply_clears_state_for_non_stateful_loader(self, tmp_path, monkeypatch):
+        """Prevent the stashed state from leaking into a later (stateful) loader."""
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+        mgr._pending_dataloader_state = {"epoch": 1}
+
+        class PlainLoader:  # no load_state_dict method
+            pass
+
+        mgr.apply_dataloader_state(PlainLoader())
+        assert mgr._pending_dataloader_state is None
+
+    def test_apply_clears_state_for_none_loader(self, tmp_path, monkeypatch):
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+        mgr._pending_dataloader_state = {"epoch": 1}
+        mgr.apply_dataloader_state(None)
+        assert mgr._pending_dataloader_state is None
+
+    def test_save_persists_dataloader_state(self, tmp_path, monkeypatch):
+        """save() must include dataloader state when a stateful loader is passed."""
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+
+        class Loader:
+            def state_dict(self):
+                return {"epoch": 4, "batches_yielded": 200}
+
+        mgr.save(step=1, tokens_seen=64, dataloader=Loader())
+        saved = torch.load(tmp_path / "step_1" / "train_state.pt", weights_only=False)
+        assert saved["dataloader"] == {"epoch": 4, "batches_yielded": 200}
+
+    def test_load_stashes_dataloader_state_when_no_loader_provided(self, tmp_path, monkeypatch):
+        """load(dataloader=None) must stash the dataloader state for later apply."""
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+        ckpt_dir = tmp_path / "step_1"
+        ckpt_dir.mkdir()
+        saved_state = {"epoch": 2, "batches_yielded": 50}
+        torch.save(
+            {
+                "step": 1,
+                "tokens_seen": 64,
+                "rng": get_rng_state(),
+                "dataloader": saved_state,
+            },
+            ckpt_dir / "train_state.pt",
+        )
+
+        step, tokens, _ = mgr.load(path=str(ckpt_dir))
+
+        assert step == 1
+        assert tokens == 64
+        assert mgr._pending_dataloader_state == saved_state
+
+    def test_load_restores_directly_when_loader_provided(self, tmp_path, monkeypatch):
+        """load(dataloader=X) must restore directly and leave pending state empty."""
+        from unittest.mock import MagicMock
+
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+        ckpt_dir = tmp_path / "step_1"
+        ckpt_dir.mkdir()
+        saved_state = {"epoch": 2, "batches_yielded": 50}
+        torch.save(
+            {
+                "step": 1,
+                "tokens_seen": 64,
+                "rng": get_rng_state(),
+                "dataloader": saved_state,
+            },
+            ckpt_dir / "train_state.pt",
+        )
+
+        loader = MagicMock(spec=["load_state_dict"])
+        mgr.load(path=str(ckpt_dir), dataloader=loader)
+
+        loader.load_state_dict.assert_called_once_with(saved_state)
+        assert mgr._pending_dataloader_state is None
+
+    def test_load_no_stash_when_no_dataloader_key(self, tmp_path, monkeypatch):
+        """Missing dataloader key in train_state leaves pending state empty."""
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+        ckpt_dir = tmp_path / "step_1"
+        ckpt_dir.mkdir()
+        torch.save(
+            {"step": 1, "tokens_seen": 64, "rng": get_rng_state()},
+            ckpt_dir / "train_state.pt",
+        )
+
+        mgr.load(path=str(ckpt_dir))
+        assert mgr._pending_dataloader_state is None
+
+    def test_round_trip_save_load_apply(self, tmp_path, monkeypatch):
+        """Save with loader, load without loader, apply to new loader — state flows through."""
+        mgr = _make_mock_mgr(tmp_path, monkeypatch)
+
+        captured: dict[str, dict] = {}
+
+        class RecorderLoader:
+            def __init__(self, initial: dict) -> None:
+                self._state = initial
+
+            def state_dict(self) -> dict:
+                return self._state
+
+            def load_state_dict(self, state: dict) -> None:
+                captured["restored"] = state
+
+        saver = RecorderLoader({"epoch": 7, "batches_yielded": 333})
+        mgr.save(step=5, tokens_seen=128, dataloader=saver)
+
+        # Simulate a fresh process: build a new manager and load without loader.
+        mgr2 = _make_mock_mgr(tmp_path, monkeypatch)
+        step, tokens, _ = mgr2.load(path=str(tmp_path / "step_5"))
+        assert step == 5
+        assert tokens == 128
+        assert mgr2._pending_dataloader_state == {"epoch": 7, "batches_yielded": 333}
+
+        # Build loader after load() and apply the stashed state.
+        restorer = RecorderLoader({"epoch": 0, "batches_yielded": 0})
+        mgr2.apply_dataloader_state(restorer)
+        assert captured["restored"] == {"epoch": 7, "batches_yielded": 333}
+        assert mgr2._pending_dataloader_state is None