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.