Support multi-process/multi-node sharding for S3IterableDataset

[jamesbornholt]

We currently don't have a built in way to do sharding for S3IterableDataset, so every worker process in a DataLoader will see the same stream of objects. We should have a way to do this.

In the meantime, something like this from torchdata will work as a workaround:

from s3torchconnector import S3IterableDataset
from torch.utils.data import DataLoader
from torch.utils.data.datapipes.iter import IterableWrapper

dataset = S3IterableDataset.from_prefix("s3://doc-example-bucket/", region="us-west-2")
dataset = IterableWrapper(dataset, deepcopy=False)
dataset = dataset.sharding_filter() # Use torchdata's sharding for iterable datasets
loader = DataLoader(dataset, num_workers=2)

[cfregly]

Related pull request for Megatron: NVIDIA/Megatron-LM#729

[jamesbornholt]

The torchdata IterableWrapper is being deprecated in a future release, but it will still be present in PyTorch core. I've updated the code example above to point to that instead.

[IsaevIlya]

Changes was released as part of v1.3.0

[noepionentrust]

👋 folks, you may want to edit the documentation at https://github.com/awslabs/s3-connector-for-pytorch/blob/main/examples/Getting%20started%20with%20the%20Amazon%20S3%20Connector%20for%20PyTorch.ipynb

Important: When combining multi-process data loading with S3IterableDataset, by default each worker process will get its own replica of the dataset, and so each training sample will be duplicated num_workers times by the DataLoader. This is very likely not the behavior you want. We're currently [tracking this issue on GitHub](https://github.com/awslabs/s3-connector-for-pytorch/issues/53). In the meantime, as a workaround, you can use torchdata to wrap an S3IterableDataset and then shard it across workers:

[matthieu-d4r]

Hi @noepionentrust, good catch: we will update the documentation in an upcoming revision. Thanks!