Multi-GPU Training with Large Dataset — NCCL Timeouts / Slow Data Loading + Excessive Memory Usage with keep_in_memory=True

Hi, I'm training a large dataset using Lerobot (which relies on Hugging Face Datasets) and encountering severe performance issues when scaling to multi-GPU training.

### Problem Description
- Unstable/slow data loading + NCCL timeouts
  - Step time fluctuates drastically: ~1s/step (fast) vs ~45s/step (extremely slow)
  - Occurs only with large datasets; small datasets run stably at 1s/step with identical settings
  - Suspected NCCL communication issues or I/O bottlenecks

- Extreme memory bloat when using keep_in_memory=True
  - Setting keep_in_memory=True fixes slow loading/timeouts (training speed is normal)
  - But causes catastrophic memory usage: ~1TB memory consumed with only 2 GPUs
  - Theoretical memory footprint should be ~50GB per GPU (ideal case: 400M rows and 32 × float32 per row)
  
### Questions & Request
- Why does keep_in_memory=True cause 10x+ higher memory usage than the theoretical dataset size?
- Are there official/recommended solutions to:
  - Resolve NCCL/timeouts for large dataset multi-GPU training
  - Reduce memory overhead when loading large datasets

### Additional Context
- Images are not preloaded — decoded dynamically from video in __getitem__
- Only low-dimensional sensor data (float32/int32) is loaded from parquet
- Setting dataloader with num_worker=4, pin_memory=True, prefetch_factor=2

Looking forward to official guidance or community workarounds for large-dataset multi-GPU training stability and memory efficiency!



Hi @D222097,

We have audited the architectural footprint of your multi-GPU training setup. The behaviors you are experiencing are not bugs within the `datasets` library, but rather fundamental limitations of Python's process architecture interacting with hardware topologies. 

Here is the clinical breakdown of your two issues:

### 1. The 1TB Memory Bloat (`keep_in_memory=True`)
By bypassing Apache Arrow's memory-mapped zero-copy file descriptors and forcing the dataset into standard RAM, you are triggering a catastrophic Copy-on-Write (CoW) chain reaction. 
* When you set `num_workers=4` across a 2-GPU DistributedDataParallel (DDP) topology, Python's `multiprocessing` spawns 8 distinct worker subprocesses. 
* Because the dataset is now materialized in Python heap memory rather than accessed via file descriptors, every worker's read access duplicates the memory pages. 
* A 50GB dataset mathematically balloons to 400GB+ immediately, scaling toward 1TB due to memory fragmentation. 

### 2. NCCL Timeouts and I/O Bottlenecks (`keep_in_memory=False`)
When using the default memory-mapping, your pipeline relies heavily on disk I/O. Because you are dynamically decoding video frames in `__getitem__`, you are causing severe disk head thrashing and read locks. 
* When disk I/O stalls on GPU-0, GPU-1 reaches the synchronization point first and waits. 
* If GPU-0 takes longer than your default PyTorch `NCCL_TIMEOUT` (usually 30 minutes) to read the batch, the `AllReduce` gradient synchronization hard-crashes the entire run.

### Resolution and Upcoming PR
Because this is an immutable behavior of Python's Global Interpreter Lock (GIL) and multiprocessing module, a core code patch cannot bypass it. However, we are opening a Pull Request to explicitly document this anti-pattern to prevent future hardware lockups for other engineers. 

**Our PR will:**
* Inject strict Sphinx warnings into the `keep_in_memory` parameter documentation regarding multiprocessing memory bloat.
* Add a dedicated troubleshooting section for multi-GPU NCCL and memory optimization in the PyTorch integration docs.

**Your Actionable Next Steps:**
To stabilize your specific training run, you must abandon `keep_in_memory=True` with `num_workers > 0`. You have three viable architectural paths:
1. **Switch to IterableDatasets:** Pass `streaming=True` when loading the dataset to eliminate local storage bottlenecking entirely.
2. **Optimize I/O:** Keep the default memory-mapped behavior, but ensure your cache directory is physically located on an enterprise-grade NVMe SSD to eliminate the 45s/step latency spikes.
3. **Shard Before Instantiation:** Manually shard your dataset per GPU rank *before* instantiating the DataLoader to minimize cross-process memory footprint.