Enable saving and loading stateful DataLoaders in Trainer

src/lightning/pytorch/CHANGELOG.md

@@ -27,6 +27,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added a utility function and CLI to consolidate FSDP sharded checkpoints into a single file ([#19213](https://github.com/Lightning-AI/lightning/pull/19213))
 
 
+- Added support for saving and loading stateful training DataLoaders ([#19361](https://github.com/Lightning-AI/lightning/pull/19361))
+
+
 ### Changed
 
 - `seed_everything()` without passing in a seed no longer randomly selects a seed, and now defaults to `0` ([#18846](https://github.com/Lightning-AI/lightning/pull/18846))

src/lightning/pytorch/loops/fit_loop.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 from typing_extensions import override
 
 import lightning.pytorch as pl
 from lightning.fabric.utilities.data import _set_sampler_epoch, sized_len
+from lightning.fabric.utilities.types import _Stateful
 from lightning.fabric.utilities.warnings import PossibleUserWarning
 from lightning.pytorch.loops import _Loop
 from lightning.pytorch.loops.fetchers import _DataFetcher
@@ -94,6 +95,7 @@ def __init__(
 
         self._data_source = _DataLoaderSource(None, "train_dataloader")
         self._combined_loader: Optional[CombinedLoader] = None
+        self._combined_loader_states: Optional[List[Dict[str, Any]]] = None
         self._data_fetcher: Optional[_DataFetcher] = None
         self._last_train_dl_reload_epoch = float("-inf")
 
@@ -255,6 +257,9 @@ def setup_data(self) -> None:
 
         combined_loader.limits = limits
 
+        if self.restarting:
+            self._restore_combined_loader_state()
+
         self._data_fetcher = _select_data_fetcher(trainer, RunningStage.TRAINING)
         self._data_fetcher.setup(combined_loader)
         iter(self._data_fetcher)  # creates the iterator inside the fetcher
@@ -409,9 +414,34 @@ def teardown(self) -> None:
             self._data_fetcher = None
         self.epoch_loop.teardown()
 
+    @override
+    def on_save_checkpoint(self) -> Dict:
+        state_dict = super().on_save_checkpoint()
+        loaders = self._combined_loader.flattened if self._combined_loader is not None else []
+        loader_states = [loader.state_dict() for loader in loaders if isinstance(loader, _Stateful)]
+        if loader_states:
+            state_dict["combined_loader"] = loader_states
+        return state_dict
+
+    @override
+    def on_load_checkpoint(self, state_dict: Dict) -> None:
+        self._combined_loader_states = state_dict.get("combined_loader")
+        super().on_load_checkpoint(state_dict)
+
     def _should_accumulate(self) -> bool:
         """Whether the gradients should be accumulated."""
         return self.epoch_loop._should_accumulate()
 
     def _iteration_based_training(self) -> bool:
         return self.trainer.max_steps != -1
+
+    def _restore_combined_loader_state(self) -> None:
+        if not self._combined_loader_states:
+            return
+
+        loaders = self._combined_loader.flattened if self._combined_loader is not None else []
+        stateful_loaders = [loader for loader in loaders if isinstance(loader, _Stateful)]
+        for loader, state_dict in zip(stateful_loaders, self._combined_loader_states):
+            loader.load_state_dict(state_dict)
+
+        self._combined_loader_states = None  # release memory

tests/tests_pytorch/loops/test_loops.py

@@ -24,6 +24,7 @@
 from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
 from lightning.pytorch.loops import _Loop
 from lightning.pytorch.loops.progress import _BaseProgress
+from lightning.pytorch.utilities import CombinedLoader
 from torch.utils.data.dataloader import DataLoader, _MultiProcessingDataLoaderIter
 
 from tests_pytorch.helpers.runif import RunIf
@@ -882,3 +883,94 @@ def on_validation_start(self):
     )
     trainer.fit(model)
     assert model.ran_assert
+
+
+class NotStatefulIterable:
+    def __init__(self, start=0):
+        self.index = start
+
+    def __iter__(self):
+        for i in range(self.index, len(self)):
+            self.index = i
+            yield self.index
+
+    def __len__(self):
+        return 10
+
+
+class StatefulIterable(NotStatefulIterable):
+    def state_dict(self):
+        return {"index": self.index}
+
+    def load_state_dict(self, state_dict):
+        self.index = state_dict["index"] + 1
+
+
+@pytest.mark.parametrize(
+    ("train_dataloader_factory", "has_state", "batches_before", "batches_after"),
+    [
+        # No dataloader
+        (lambda: [], False, [], []),
+        # Single stateful DataLoader
+        (lambda: StatefulIterable(), True, [0, 1], [2, 3]),
+        # Single, not stateful DataLoader
+        (lambda: CombinedLoader(NotStatefulIterable()), False, [0, 1], [0, 1]),
+        # Single stateful DataLoader
+        (lambda: CombinedLoader(StatefulIterable()), True, [0, 1], [2, 3]),
+        # Multiple stateful DataLoaders
+        (lambda: CombinedLoader([StatefulIterable(3), StatefulIterable(1)]), True, [[3, 1], [4, 2]], [[5, 3], [6, 4]]),
+        # Mix of stateful and not stateful DataLoaders
+        (
+            lambda: CombinedLoader([NotStatefulIterable(3), StatefulIterable(1), NotStatefulIterable(2)]),
+            True,
+            [[3, 1, 2], [4, 2, 3]],
+            [[3, 3, 2], [4, 4, 3]],
+        ),
+    ],
+)
+def test_fit_loop_save_and_restore_dataloaders(
+    train_dataloader_factory, has_state, batches_before, batches_after, tmp_path
+):
+    """Test that the CheckpointConnector saves the state of stateful dataloaders."""
+
+    class DummyModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.seen_data = []
+
+        def training_step(self, batch, batch_idx):
+            self.seen_data.append(batch)
+            print(batch)
+
+        def train_dataloader(self):
+            return train_dataloader_factory()
+
+    trainer_kwargs = {
+        "default_root_dir": tmp_path,
+        "accelerator": "cpu",
+        "enable_checkpointing": False,
+        "enable_model_summary": False,
+        "enable_progress_bar": False,
+        "logger": False,
+        "num_sanity_val_steps": 0,
+    }
+
+    # Train for 2 steps
+    model = DummyModel()
+    trainer = Trainer(**trainer_kwargs, max_steps=2)
+    trainer.fit(model)
+    assert model.seen_data == batches_before
+
+    # Save a checkpoint
+    trainer.save_checkpoint(tmp_path / "checkpoint.ckpt")
+    checkpoint = torch.load(tmp_path / "checkpoint.ckpt")
+    if has_state:
+        assert checkpoint["loops"]["fit_loop"]["state_dict"]["combined_loader"]
+    else:
+        assert "combined_loader" not in checkpoint["loops"]["fit_loop"]["state_dict"]
+
+    # Restore training from step 2 and continue 2 more steps
+    model = DummyModel()
+    trainer = Trainer(**trainer_kwargs, max_steps=4)
+    trainer.fit(model, ckpt_path=(tmp_path / "checkpoint.ckpt"))
+    assert model.seen_data == batches_after