大模型训练的分布式优化策略：从数据并行到3D并行
训练一个千亿参数的大模型，需要数千GPU协同工作数周。如何高效利用这些GPU，是决定训练成本和速度的关键。分布式训练策略就是这个问题的核心答案。本文将系统解析各种并行策略的原理与工程实践。 一、为什么单GPU不够 1.1 显存瓶颈 一个70B参数的模型： 模型参数（FP16）：140GB 梯度（FP16）：140GB 优化器状态（Adam, FP32）：560GB 激活值：取决于batch和序列长度 总计约840GB+ ——远超单张A100的80GB显存。 1.2 计算瓶颈 训练70B模型一个epoch（1万亿token）： 单A100算力：312 TFLOPS (FP16) 训练所需FLOPS：~6×10²³ 单卡理论时间：~22000天 需要2048张A100并行才能在11天内完成。 二、数据并行（DP） 2.1 基本原理 最直观的并行——每张GPU持有完整模型副本，处理不同数据： GPU 0: 完整模型 + Batch[0:32] GPU 1: 完整模型 + Batch[32:64] GPU 2: 完整模型 + Batch[64:96] GPU 3: 完整模型 + Batch[96:128] 前向传播 → 各GPU独立计算loss 反向传播 → AllReduce同步梯度 更新 → 各GPU同步更新参数 2.2 问题：大模型装不下 数据并行要求每张GPU能装下完整模型+优化器状态+梯度。对于70B模型，需要840GB——单卡80GB根本装不下。 2.3 ZeRO优化：显存突破 DeepSpeed的ZeRO（Zero Redundancy Optimizer）通过切分训练状态来突破显存限制： ZeRO-1: 切分优化器状态 每GPU只存 1/N 的优化器状态 70B模型: 560GB/N → N=8时70GB/卡 ZeRO-2: 切分优化器状态 + 梯度 每GPU只存 1/N 的优化器状态和梯度 70B模型: (560+140)GB/N → N=8时87.5GB/卡 ZeRO-3: 切分优化器状态 + 梯度 + 参数 每GPU只存 1/N 的所有状态 70B模型: (560+140+140)GB/N → N=8时105GB/卡 参数在前向/反向时按需All-Gather # ZeRO-3的参数按需获取 class ZeRO3Layer: def forward(self, x): # 按需All-Gather本层参数 full_weight = self.all_gather(self.shard_weight) output = F.linear(x, full_weight) # 立即释放完整参数 del full_weight return output def backward(self, grad_output): # 按需All-Gather本层参数 full_weight = self.all_gather(self.shard_weight) grad_input = grad_output @ full_weight.T # Reduce-Scatter梯度 self.shard_grad = self.reduce_scatter(grad_output.T @ self.input) return grad_input 三、张量并行（TP） 3.1 原理 将单个层的参数矩阵切分到多张GPU上： ...
