4D 并行 DP/TP/PP/SP

大模型训练把参数量、梯度、激活、序列切到多卡上的四种并行维度。① DP 数据并行: 每卡一份完整模型, 只切数据; 每步迭代结束做一次全量梯度 AllReduce, 环形实现下每卡通信量 $2\frac{N-1}{N}S \approx 2S$($S$ = 单卡梯度总大小, 随参数量线性增长——70B 每卡梯度 ≈ 140GB, 因此大模型 DP 同步很贵); 通信只发生在 step 之间, 可被 bucket 异步重叠隐藏。② TP 张量并行: 层内矩阵按列/行切分(详见 megatron-3d-bubble), 每个 transformer 层 2 次 AllReduce, 每次消息量仅 ~$b \times s \times h$ 的激活(与参数量无关, 但发生频率高、延迟敏感)——必须用节点内 NVLink(900GB/s+)才能承受。③ PP 流水线并行: 层按 stage 切分, 每 stage 持有若干层; 通信只在相邻 stage 边界, 每个 microbatch 传一次激活与一次梯度(~$b \times s \times h$), 量大频率低、延迟线性叠加——适合跨节点 IB 链路; 代价是气泡(bubble, 气泡率 $\frac{p-1}{m+p-1}$, 见 megatron-3d-bubble)。④ SP 序列并行: 把序列维度切成 N 段分到 N 卡, 解决长序列下 KV/激活显存随 seq 线性膨胀; 两种实现: Megatron-SP(与 TP 结合, 把 LayerNorm/Dropout 的激活也切分, 削减 TP 的激活通信)与 Ring-Attention(ring 传递 KV, 每卡显存 O(1))。3D 并行组合 DP × TP × PP: 以 GPT-3 175B 为例, 1024 张 A100 = DP 64 × TP 8 × PP 2; 分配原则: TP 先切满节点内(吃 NVLink 带宽)、PP 跨节点(通信最少)、DP 最外层做梯度同步——通信代价自内向外递减, 可扩展性自内向外递增。

训练 Infra 面试必考: “DP/TP/PP 各通信多少”“为什么 TP 不能跨节点”“SP 与 TP 什么关系”“给 512 张卡怎么配 3D 并行”。

单卡显存放不下(权重+优化器+激活)与单卡算力/带宽打不满——TP 切权重、PP 切层、DP 切数据、SP 切序列, 四者叠加让任意大的模型在任意多的卡上以接近线性的效率训练, 且每层通信量都能落在合适的互联(节点内 NVLink / 跨节点 IB)上。

手算 DP 通信量: 每步环形 AllReduce $2\frac{N-1}{N}S \approx 2S$($S$ = 单卡梯度), 与卡数 N 无关; 解释为什么大模型 DP 仍然贵——S 随参数量线性增长(70B 梯度 ≈ 140GB)。

TP 通信特征: 每层 2 次 AllReduce、每次消息量 ~$b \times s \times h$ 与参数量无关; 推导为什么 TP 只能在节点内(NVLink 900GB/s、低延迟), 跨节点会怎样。

PP 通信: 为什么 PP 通信量小(每 microbatch 相邻 stage 一次激活 + 一次梯度, ~$b \times s \times h$)且适合跨节点; 它的代价是什么(气泡率 $\frac{p-1}{m+p-1}$, 见对应卡片)。

SP 序列并行: 切什么(序列维度)、解决什么(KV/激活显存随 seq 线性)、两种实现 Megatron-SP 与 Ring-Attention 各自的通信代价与适用场景。

3D 并行配置题: 给定 GPU 总数与模型规模, 按「TP 节点内 → PP 节点间 → DP 最外层」分配; 复述 GPT-3 175B 的 1024 卡 = DP 64 × TP 8 × PP 2 并解释每层选择的理由。