PyTorch torchrun（AWS SageMaker）DDP单实例多GPU数据下载疑问

• 在单实例多GPU的torchrun DDP模式下，每个local_rank对应的进程确实会独立执行数据下载、Dataset和DataLoader的创建流程。DDP启动时会为每个GPU生成一个独立进程，每个进程拥有独立的Python运行环境，代码里的下载、数据集初始化逻辑默认会被所有进程重复执行。

• 关于数据是否加载到每个GPU：

  • 每个进程都会把Pandas DataFrame加载到自身的CPU内存空间中，这部分是进程隔离的，会占用多份CPU内存；
  • 只有当你执行to(device)或DataLoader将batch数据送入GPU时，对应进程的当前batch数据才会加载到它绑定的GPU显存里，整个数据集不会一次性全部塞进GPU显存。

• 是否需要确保数据适配GPU内存：

  • 不需要让完整数据集适配GPU显存，因为DataLoader是按批次加载数据的，每次仅会将一个batch的数据送入GPU；
  • 但要注意CPU内存占用问题：如果数据集较大，多进程重复加载完整DataFrame可能耗尽实例的CPU内存，可通过以下方式优化：
    • 提前将数据下载到实例的共享存储（如SageMaker的EBS卷），让所有进程直接从本地读取，避免重复下载；
    • 使用DistributedSampler自动为每个进程分配数据分片，减少单个进程加载的数据量；
    • 采用内存映射格式（如Parquet、Feather）存储数据，降低Pandas DataLoader加载时的内存开销。

内容的提问来源于stack exchange，提问作者B_Miner