resume state of data loader and data module from mid-epoch checkpoint

[theissenhelen]

Is your feature request related to a problem? Please describe.

When resuming a run mid-epoch, we currently go through the whole dataset again making the epoch longer. This is quite expensive in rollout training. This was shown in PR when training with max_steps.

Describe the solution you'd like

Ideally, we want to restore the state of our dataloader and datamodule when loading from a checkpoint. Torchdata offers the torchdata stateful dataloader. However, it does no (yet?) support distributed training.

There are four stages we need to consider:

1. Dataloader
2. Datamodule
3. Dataset
4. Sampler

Things to brainstorm and discuss:

• We need to understand limitations for workers/ranks - might require to design the interface as a flag so if this too limitating people can switch off
• An intermediate solution could be limited to using the same number of workers/ranks.
• Is torchdata maintained!?
• Keep an eye on this issue
• How does pytorch usesa stateful dataset in its dataloader.
• How pytorch lightning connect the stateful dataloader with the stateful dataset (Is there functionality for this?)

First steps:

• create a number of toy examples to understand the workings of the stateful dataloader/data-module (@anaprietonem, @theissenhelen)

Describe alternatives you've considered

No response

Additional context

No response

Organisation

ECMWF