StreamingDataset with DDP hangs and then crashes with NCCL timeout error

Environment

AWS Deep Learning Machine with 8xA100 and CUDA 11.8

To reproduce

Steps to reproduce the behavior:

1. Use StreamingDataset to load ImageNet from a local SSD using DDP.

Expected behavior

The data loads as expected when running on a single GPU. I expect the data to load in the same way on multiple GPUs.

Additional context

I'm using accelerate launch / torchrun to launch 8 processes. I'm loading from a local disk, not a remote file. I do this by passing the same (local) directory to both the local and remote arguments of StreamingDataset. Specifically, I have a dataset that looks like:

class CustomStreamingDataset(StreamingDataset):
    def __init__(
        self, 
        local: str, 
        remote: Optional[str] = None, 
        shuffle: bool = False, 
        batch_size: int = 1, 
        transform: Optional[Callable] = None,
    ):
        remote = local if remote is None else remote
        super().__init__(remote=remote, local=local, shuffle=shuffle, batch_size=batch_size)
        self.transform = transform

    def __getitem__(self, idx):
        item = super().__getitem__(idx)
        feats = item['features'].squeeze(0)
        label = item['class']
        if self.transform is not None:
            feats = self.transform(feats)
        return feats, label

And then I load it as follows:

if args.use_streaming_dataset:
    data_dir = f"{args.feature_path}/imagenet256_streaming"
    dataset = CustomStreamingDataset(data_dir, shuffle=True, batch_size=batch_size)
    load_kwargs = dict()
else:
    features_dir = f"{args.feature_path}/imagenet256_features"
    labels_dir = f"{args.feature_path}/imagenet256_labels"
    dataset = CustomDataset(features_dir, labels_dir)
    load_kwargs = dict(shuffle=True, pin_memory=True, drop_last=True)
loader = DataLoader(
    dataset, batch_size=batch_size, num_workers=args.num_workers, **load_kwargs
)

The code does (not) work under the following settings:

• 1 GPU, without streaming dataset: works correctly
• 8 GPUs, without streaming dataset: works correctly
• 1 GPU, with streaming dataset: works correctly
• 8 GPUs, with streaming dataset: hangs forever

Eventually the program crashes with the following error:

RuntimeError: NCCL communicator was aborted on rank {RANK}.  Original reason for failure was: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=15, OpType=BROADCAST, Timeout(ms)=1800000) ran for 1804394 milliseconds before timing out.

where {RANK} is replaced by 0, 1, ... 7 on each process.

Perhaps this is related to #293. However, since it's not exactly the same, I thought I should leave a separate issue.

Start by reproducing the difference between single-GPU and 8-GPU runs using StreamingDataset with the local ImageNet directory and accelerate launch or torchrun. Inspect the StreamingDataset and DDP data-loading path, then verify that all ranks complete loading without hanging and that the NCCL broadcast timeout does not occur.