Sharded data loading with DDP

[niansong1996]

Relatively new to Lightning, has been struggling with this for a few days. I am wondering what is the right way to do data reading/loading under DDP.

More details:

• I am currently only running on one node, typically with 16 GPUs;
• I have a relatively large dataset (a few GiBs), so I sharded the dataset into separate files (e.g., train_shard_0.jsonl, train_shard_1.jsonl, etc), while # shards = # gpus;
• What I am doing now, is use torch.distributed.get_rank() to direct LightningDataModule on each process to read a specific shard of the dataset;
• However, it seems to me that by default, there is already some communication/protocol between the processes, so that every process will not read all of their data, but to skip according to the ranks of the process. For example, if the batchs_size=4, the first batch would have process 0 load examples [0, 1, 2, 3] and process 1 load [4, 5, 6, 7], as if they are still a copy of the same dataset, which make the sharding meaningless;
• I suppose I cound let all processes read the same original unsharded file, and since the aforementioned skipping is in place, everything will be correct. But in this way, every process reads examples that will eventually be dropped, isn't this a huge waste of time though?
• I am using map-style datasets at the moment, but I can change to iterate-style if needed

So I am wondering what is the best practice for data reading/loading under DDP, and whether I would need to define my own customized Dataloader or other classes to make it right. Thanks for help of any kind!

[tchaton]

Dear @niansong1996,

By default, Lightning assumes the same dataset is being loaded on all ranks and a DistributedSampler is injected within the DataLoader to assign each rank a sub-part of the dataset.

In your case, you have already created the shard. You could make your dataset rank aware as follow.

Using an IterableDataset is also a good solution.

class Dataset(Dataset):

    def __init__(self, global_rank = ..):
        self.list_shards = [...]
        self.current_shard = self.list_shards[global_rank]

        self.data = load(self.current_shard)

    def __getitem__(self, index):
        return len(self.data[idx])

    def __len__(self):
        return len(self.data)


class DataModule()

    def prepare_data(self):
        # create shard if they don't exist

    def train_dataloader(self):

        ds = Dataset(self.trainer.global_rank)
        return DataLoader(ds)


# prevent Lightning to inject DDSampler
Trainer(replace_sampler_ddp=False)

Here is a good reference: https://github.com/webdataset/webdataset

[amyxlu]

If using an IterableDataset as suggested, will the DistributedSampler still work? The documentation seems to suggest that it is using indices to create subsets, not sure if it is already designed to work in tandem with IterableDatasets.