分布式训练
分布式训练是训练百亿、千亿参数模型的基础技术。
共 3 篇文章·阅读时间：约35分钟
01数据并行
数据并行是最常用的分布式训练方法，通过复制模型并行处理不同数据。
数据并行原理
工作流程
- 模型复制到N个GPU
- 每个GPU处理不同的数据批次
- 计算各自的梯度
- 所有GPU的梯度聚合（AllReduce）
数据并行变体
DDP
DistributedDataParallel
通信效率高，是PyTorch官方实现
FSDP
FullyShardedDataParallel
同时分片模型参数，支持超大模型
02模型并行
当单个GPU放不下模型时，需要将模型拆分到多个GPU。
模型并行原理
核心思想
将模型的不同层放到不同GPU上，每个GPU只负责部分层的计算。 需要时进行跨GPU通信。
MP/TP/PP
Tensor Parallelism (TP)
张量并行：将单层的权重矩阵拆分到多个GPU
Megatron-LM是典型实现
Pipeline Parallelism (PP)
流水线并行：将不同层放到不同GPU
需要mini-batch内部再切分micro-batch
3D Parallelism
TP + PP + DP 组合
训练万亿参数模型的标配
03ZeRO优化
ZeRO通过分片优化器状态、梯度和参数，大幅减少显存占用。
ZeRO阶段
- ZeRO-1：分片优化器状态（减少4x显存）
- ZeRO-2：分片优化器状态 + 梯度（减少8x显存）
- ZeRO-3：分片优化器状态 + 梯度 + 参数（减少Nx显存，N=GPU数）
使用指南
DeepSpeed
微软出品，支持ZeRO
提供Training Optimization库
选择建议
中小模型：ZeRO-1/2 足够
超大模型：ZeRO-3 或 ZeRO-3 + Offload
实战技巧
- ZeRO-3通信量大，适合高带宽NVLink
- 可结合CPU Offload处理超长上下文
- FSDP是PyTorch原生的ZeRO实现