GPU多卡训练实战指南
文章目录
- 引言
- 一、多卡并行的基本范式
- 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 上（如按层切分）。
- <