NCCL Timeouts from Adding More Datasets + Proportion Sampling

Environment

• OS: Ubuntu 22.04.5 LTS (GNU/Linux 5.15.0-1048-oracle x86_64)
• Hardware (GPU, or instance type): 4 Nodes, 8xH100
• PyTorch & CUDA Versions: PyTorch 2.5.1, CUDA 12.4
• Streaming Version: 0.9.0 w/ this fix
• NCCL Timeout: 10 minutes

To reproduce

Train model with config like --

stream_configs:
    - remote: <s3://path_to_r2_dataset>
       local: /scratch/<path_to_local_dataset_on_nvme>
       proportion: <proportion_value>
    - remote: <s3://path_to_r2_dataset>
       local: /scratch/<path_to_local_dataset_on_nvme>
       proportion: <proportion_value>
    ...
predownload: 3072 # trying to provide enough of a buffer so that we don't stall out, waiting for data during training
cache_limit: XX TB # ~60% of the entire dataset can be stored in cache on local NVME drives (without running out of space)
batch_size: 128
batching_method: `device_per_stream` # each dataset has samples with different sizes; so they have to be batched per stream
shuffle: True

Expected behavior - No NCCL Timeouts

When we hit "breaks" in training (i.e., evaluation every K batches OR the end of the epoch), we often get NCCL timeout errors. Sometimes we can get through multiple epochs of the dataset -- before we hit a NCCL timeout at the end of an epoch. Other times, we get NCCL timeouts every time we have an evaluation period. When we restart training, we are able to get deeper into training, so the underlying data does not seem to be corrupted.

If we train with shuffle: False, we never get NCCL timeouts -- so we're fairly confident that the NCCL timeouts are tied to shuffling logic. (see comments below)

My best hunch (pure guesswork) is that it's taking a long time to shuffle the data (and download new shards) at these "breaks" in training -- and that drives the NCCL timeout. For example, if we are at the end of an epoch - the shuffling may require the download of A LOT more data, since we cannot hold the entire dataset in cache. Perhaps, the GPUs are not waiting (in the same way) when we restart training -- so they're able to spend a lot of time downloading and shuffling the data, without the NCCL timeouts.

Additional context

We're not using the latest version of StreamingDataset because of the errors discussed in this thread, from January.

No source files or tests are named. Start by reproducing the provided multi-dataset proportion-sampling configuration on the listed Streaming version, then inspect shuffle, shard downloading, and evaluation or epoch-boundary behavior; done means training completes those boundaries without NCCL timeouts.