How can I resume the state of a stateful dataloader in pytorch lightning ?

[hiyyg]

I want to use the stateful dataloader in pytorch lightning.
https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader

What should I do to make it work?

[michaelghorgan]

Just drop-in replace the standard torch DataLoader in your LightningDataModule with the StatefulDataLoader from torchdata:

def train_dataloader(self):
        return StatefulDataLoader(
            self.trainset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
        )

Note that if you use a progress bar, you might need to set the total number of training batches explicitly (I get a ? if I don't), e.g.:

class MyProgBar(TQDMProgressBar):
    def init_train_tqdm(self):
        bar = super().init_train_tqdm()
        bar.total = self.trainer.num_training_batches
        return bar

[amorehead]

@nrxszvo, I tried training a model with PyTorch Lightning as you've described, but when I try reloading my model's latest checkpoint, PyTorch Lightning raises the following error concerning a null dataset iterator.

# Line 358 of `lightning/pytorch/utilities/combined_loader.py`
if self._iterator is None:
     raise RuntimeError("Please call `iter(combined_loader)` first.") # <-- This is the error raised.

My call stack for Lightning's source code looks like the following.

My full call stack is as follows.

__len__ (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/utilities/combined_loader.py:358)
sized_len (/opt/conda/lib/python3.11/site-packages/lightning/fabric/utilities/data.py:52)
reset (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/loops/fetchers.py:76)
reset (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/loops/fetchers.py:142)
teardown (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/loops/fetchers.py:80)
teardown (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py:502)
_teardown (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:1047)
_call_and_handle_interrupt (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py:69)
fit (/opt/conda/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:567)

Do you have any ideas what might be going wrong here? I've already tried ensuring that my StatefulDataLoader has a snapshot_every_n_steps value that is synchronized with the training step frequency at which I save a model checkpoint (i.e., 250 training steps). Even still, when I set a breakpoint in the CombinedLoader's _load_state_dicts() method, I see that the (stateful) dataloader states are seemingly new even after training for 13k training steps.

This is with the following Python package versions:

lightning                      2.5.1rc2
lightning-utilities            0.14.3
pytorch-lightning              2.5.2
torch                          2.7.1+cu118
torchdata                      0.12.0a0+088e434

[michaelghorgan]

I'm not sure, but it seems strange that the code for combined_loader is being called here. Are you using multiple loaders somewhere? Or can you see why combined_loader is being used?