Introduce a mechanism to set the epoch on the sampler in LightningLite

[awaelchli]

🚀 Feature

We need a mechanism to set the epoch on the distributed sampler via .set_epoch().

Motivation

To correctly handle shuffling with the DistributedSampler in DDP, the PyTorch user would normally call

sampler.set_epoch(epoch)

in their training loop. PL handles this automatically for the user, but in Lite, the training loop is owned by the user and hence should be signaled in a different way.

Pitch

Provide a strategy-agnostic API to set the epoch. The idea here is that you don't need to change the code or add boilerplate conditional logic to handle the sampler when you switch from DDP to single device or vice versa.

Before

....
train_dataloader, val_dataloader = self.setup_dataloaders(train_dataloader, val_dataloader)

for epoch in range(num_epochs):

    # Beginning of a new epoch

    # Boilerplate code
    if isinstance(train_dataloader.sampler, DistributedSampler):
        train_dataloader.sampler.set_epoch(epoch)


    for idx, data in enumerate(train_dataloader):
        ...

Now:

....
train_dataloader, val_dataloader = self.setup_dataloaders(train_dataloader, val_dataloader)

for epoch in range(num_epochs):

    # Beginning of a new epoch

    # This is a no-op if not using distributed sampler (DDP)
    self.set_epoch(epoch, train_dataloader)

    for idx, data in enumerate(train_dataloader):
        ...

Alternatives

Don't introduce this. It is left to the user to handle this boilerplate code.

Additional context

PyTorch docs for set_epoch

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging PyTorch Lightning, Transformers, and Hydra.

cc @Borda @carmocca @justusschock @awaelchli

[awaelchli]

There are two other alternatives:

1. Since we are wrapping the dataloader as _LiteDataloader in setup, we could expose dataloader.set_epoch(int)
2. We could call set_epoch internally in _LiteDataloader whenever __iter__() is called, and keep a count.

[carmocca]

I like the option of doing it internally as long as there's an easy opt-out (maybe as an attribute in the dataloader) so that the user can customize this.

[awaelchli]

After my comment I investigated more and found that we are actually already doing it indirectly in our sampler wrapper here:
https://github.com/Lightning-AI/lightning/blob/38acba08fc009caac32b6b0ab8984d8935569db9/src/lightning_lite/utilities/distributed.py#L303

It's easy to forget (and I did) that we are wrapping with our subclassed version of DistributedSampler.
By calling dataset.reset(), we are re-drawing all samples from the underlying sampler. For the two standard samplers in torch, SequentialSampler and RandomSampler, this works fine and is essentially equivalent to set_epoch().

Since we always use this wrapper class for every sampler, we don't need the set_epoch as far as I can tell. Maybe we should document this though, because it's not very obvious from reading the code that the code path in DistributedSampler.__iter__ that does the shuffling of indices is not necessary for our purposes.

[carmocca]

this works fine and is essentially equivalent to set_epoch()

How so? set_epoch just changes the epoch value, whereas our reset will just re-create the iterator, but the same old epoch value gets used.

[awaelchli]

So my explanation wasn't 100% correct. The thing is, we can't just add the set_epoch call to our wrapper class there. It wouldn't have any meaningful effect because the shuffling of data is completely controlled by the sampler we wrap around. For custom samplers, there is nothing we can do about that.

In #16101, I fix this problem for the two standard samplers that handle shuffle=True|False in the DataLoader, by not wrapping them. And then we only wrap the custom samplers to make them "distributed". I added a note to the wrapper class to clarify that shuffle and seed have no impact when using custom samplers:

https://github.com/Lightning-AI/lightning/pull/16101/files#diff-b49e7b638737b063bc51ee0cc6727c82a3ff52bf72a40f6a49292818b2834910