Support setup for torch DataPipes

[awaelchli]

Description & Motivation

When working with torchdata's DataPipes in distributed settings, the user can use DataLoader2 to run the pipe with a specific reading service that applies sharding of the data pipe correctly. For example:

# multiprocessing in a single node
from torchdata.dataloader2 import DataLoader2
from torchdata.dataloader2 import PrototypeMultiProcessingReadingService

service = PrototypeMultiProcessingReadingService(num_workers=5)
dataloader = DataLoader2(datapipe, reading_service=service)

# distributed/multi-gpu/multi-node
from torchdata.dataloader2 import DistributedReadingService

service = DistributedReadingService()
dataloader = DataLoader2(datapipe, reading_service=service)

However, this does not mix well with Lightning as the user would have to change the reading service when switching from one strategy to another.

Pitch

Similar to what we do with the injection of the DistributedSampler for the regular DataLoader, add the reading service for the datapipe automatically for the user.

fabric = Fabric(...)

dataloader = DataLoader2(datapipe)

# sets up the reading service for distributed training
dataloader = fabric.setup_dataloader(dataloader)

Alternatives

No response

Additional context

If you are interested in learning more about torchdata, here is a good YouTube video by PyTorch that introduces the main concepts an values.

cc @Borda @justusschock @awaelchli @carmocca

[NivekT]

I am not super familiar with Lightning, but what happens if users do:

mp_rs = MultiProcessingReadingService(num_workers=2)
dist_rs = DistributedReadingService()
rs = SequentialReadingService(dist_rs, mp_rs)  # Execute both distributed and multiprocessing
dataloader = DataLoader2(datapipe, reading_service=rs)

This would normally work with standalone torchdata (the bottom example here).

Would it still work if dataloader (as defined above) is passed into Lightning?

[awaelchli]

In Lightning Trainer yes, because it supports arbitrary iterables. But we don't have any correctness tests for DataLoader2 specifically, so no guarantees. In Fabric, the setup_dataloaders() method does not support DataLoader2 atm.

[tensorcopy]

Is there any update or plan to support Dataloader2? Thanks!

[awaelchli]

The future of torchdata is very unclear. The development has paused, here is the official statement: pytorch/data#1196

I suggest we wait until we know about future plans. For now, users can have datapipes/dataloader2 with Lightning by configuring them manually.