文章目录
- 引言
- 一、多卡并行的基本范式
- 1. 数据并行(Data Parallelism)
- 2. 模型并行(Model Parallelism)
- 二、PyTorch 多卡方案对比
- 三、DDP 核心组件详解
- 1. 启动方式:必须用 `torchrun`
- 2. 数据分片:`DistributedSampler`
- 3. 模型包裹:`DistributedDataParallel`
- 4. 全局操作:`all_gather` / `reduce`
- 四、常见误区与最佳实践
- ❌ 误区 1:以为 `DistributedSampler` 就够了
- ❌ 误区 2:在 `__init__` 外调用 `dist.get_rank()`
- ✅ 最佳实践 1:使用 `with torch.no_grad()` 做评估
- ✅ 最佳实践 2:保存/加载模型时用 `module.state_dict()`
- ✅ 最佳实践 3:设置随机种子保证可复现
- 五、大模型训练的显存优化
- 1. FSDP(Fully Sharded Data Parallel)
- 2. 梯度检查点(Gradient Checkpointing)
- 六、性能调优建议
- 结语
- 参考资料
引言
随着深度学习模型规模的爆炸式增长,单张 GPU 的计算能力和显存已难以满足训练需求。无论是训练百亿参数的大语言模型,还是高分辨率的视觉 Transformer,多 GPU 并行训练已成为现代深度学习系统的标配。
一、多卡并行的基本范式
多 GPU 训练主要分为两类:
1. 数据并行(Data Parallelism)
- 核心思想:每个 GPU 持有完整的模型副本,但处理不同的数据子集。
- 梯度同步:前向/反向计算后,通过All-Gather操作同步梯度,再统一更新参数。
- 适用场景:绝大多数 CV/NLP 模型(如 ResNet、BERT、ViT)。
💡 优势:实现简单,扩展性好;
❗ 挑战:每个 GPU 需能容纳完整模型。
2. 模型并行(Model Parallelism)
- 核心思想:将模型本身切分到不同 GPU 上(如按层切分)。
- <