最近学习 MIT 6.S191 Lecture 9,陆续把 DDP、ZeRO、FSDP2、Tensor Parallelism、Pipeline Parallelism、Sequence Parallelism、Context Parallelism,以及 Expert Parallelism 过了一遍。
上一篇把单个 MoE Layer 的执行流走了一遍:Router 选 Expert、Token Dispatch 按 Expert 重排、Grouped GEMM 统一调度、Combine 加权合回原始 Token 顺序。整条链路假设所有 Expert 都在同一张 GPU 上。
上一篇 TP 写完之后,我对 Megatron 的 Column Parallel + Row Parallel 已经比较清楚了:前一层按输出特征切分,后一层按输入特征切分,两层之间不需要 AllGather,只在最后做一次 AllReduce 把部分和加起来。
起因 # 上一篇搞清楚了 Pipeline Parallelism 的调度机制:GPipe 先 Forward 后 Backward,1F1B 通过交替调度降低 Activation 暂存压力,但理想条件下 Bubble Time 相同。
起因 # 上一篇把 ZeRO 的三个 Stage 走了一遍之后,我以为关于参数分片这件事已经想得差不多了——ZeRO-3 / FSDP 把参数、梯度和 Optimizer States 全部分片存储到多张 GPU 上,单卡 Model States 从 112 GB 降到约 28 GB,问题解决。
起因 # 上一篇 Tensor Parallelism 写完之后,我对 TP 的基本原理有了底:把单个算子内部的矩阵乘法拆到多张 GPU 上执行,每张卡只算部分输出,最后通过 AllReduce 聚合。TP 分片的是计算本身。
起因 # 上一篇搞清楚了 AllReduce、ReduceScatter 和 AllGather 三种通信原语,末尾提到 DDP 和 ZeRO Stage 2 选择了不同的通信原语组合,根本原因是参数更新的责任分配不同。
起因 # 之前几篇文章从模型显存组成一路算到 7B 模型的训练账本,结论很直接:BF16 mixed precision + AdamW 的 Model States 就有 112 GB,一张 80 GB A100 连放都放不下。Checkpointing + Offloading 两个一起用可以勉强塞进去,但训练速度要被 PCIe 传输和额外重计算拖累。