Implement DistribtuedReadingService

[ejguan]

Add DistributedReadingService

• Single process
• Share shuffle seeds across distributed process
• Automatically distributed sharding

Add tests for both DataLoader2 and DataLoader.

• Spawn processes
• Elastic training

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[VitalyFedyunin]

Just to be clear, does the first version assumes 1 process per rank?

[NivekT]

Question about the distributed design in general:

1. Are we expecting users to always use something like torch.multiprocessing.spawn to start distributed training? And that will properly start/clean up all the processes?
2. To what extent is the optimization from Second prototype to do pre-sharding work in single process #555 compatible with this? Maybe it can work for every node?

[ejguan]

Just to be clear, does the first version assumes 1 process per rank?

Yeah. I think our next step is to support mixed reading service (distributed reading service + multiprocessing reading service)

[ejguan]

1. Are we expecting users to always use something like torch.multiprocessing.spawn to start distributed training? And that will properly start/clean up all the processes?

Nope. I can add a different test for elastic training.

[ejguan]

This PR should be ready to review. cc: @NivekT @VitalyFedyunin
In terms of multiprocessing + distributed, I will figure out the design of mixed reading services first then add such feature.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[ejguan]

Actually, there is one thing left, which is doc and tutorial for DistributedReadingService. I will do a separate PR to document DataLoader2 and DistributedReadingService.

[ejguan]

I will wait until pytorch/pytorch#83741 is landed and released into nightly because it will also use the updated API.

[ejguan]

Failing DataPipe tests because the nightly binaries for mac haven't been updated yet.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[NivekT]

LGTM!

nit comment:
IIUC, the difference between this and _test_distributed_rs is DL1 vs DL2?

• If so, we can potential remove the duplicate code (not urgent).
• Alternatively, we can just label _test_distributed_dl and _test_distributed_dataloader (and elastic_dl/elastic_training) so it will be obvious from the first glance that those two are mostly the same but testing different version of DL.

[ejguan]

IIUC, the difference between this and _test_distributed_rs is DL1 vs DL2?

• If so, we can potential remove the duplicate code (not urgent).
• Alternatively, we can just label _test_distributed_dl and _test_distributed_dataloader (and elastic_dl/elastic_training) so it will be obvious from the first glance that those two are mostly the same but testing different version of DL.

Fixed

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[ejguan]

I will land this PR until PyTorch nightly is updated.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.