第 6 章 流水线并行与专家并行¶
开场问题:TP 卡在 8,DP 再大也压不住 671B 的参数量。而且一开 MoE,profile 里 40% 的时间是一段叫 AllToAll 的东西。
6.1 流水线并行与气泡¶
PP 把模型按层切成若干 stage,每个 stage 放在不同 GPU 上,数据像流水线一样依次经过。第一张卡算完前 8 层,把激活传给第二张卡,自己接着算下一个 micro-batch。
它的固有代价是气泡:流水线填满之前和排空之后,总有卡在空转。气泡率约为 (P-1)/(M+P-1),其中 P 是 stage 数、M 是 micro-batch 数。4 个 stage、4 个 micro-batch 时气泡占 43%;把 micro-batch 提到 32,气泡降到 8.6%。所以 PP 要求 micro-batch 数远大于 stage 数,这也是它和大 batch 训练天然搭配的原因。
6.2 1F1B 与交错调度¶
朴素流水线先把所有 micro-batch 的前向做完再做反向,峰值时要同时保存所有 micro-batch 的激活。1F1B(one-forward-one-backward)让每个 stage 在稳态下交替执行一次前向和一次反向,激活驻留量从 M 份降到约 P 份。
交错流水线更进一步:一张卡不再负责连续的若干层,而是负责几段不连续的层(比如第 1–4 层和第 17–20 层),让流水线的"节拍"更密,进一步压缩气泡,代价是通信次数增加、调度复杂。
还有一个实际问题:stage 不能按层数平均切。第一个 stage 多了 embedding,最后一个 stage 多了 lm_head 和 loss 计算(词表 15 万时这是一个巨大的矩阵乘,还会产生巨大的 logits 张量)。按层数均分会让首尾两段明显更慢,而流水线的速度由最慢的 stage 决定。
6.3 专家并行¶
MoE 的专家天然可以分到不同卡上。一个 token 经路由决定去哪几个专家后,需要被发送到持有那些专家的卡上,算完再收回原来的位置——这就是每层两次 AllToAll。
第一册第 8 章讲了 MoE 为什么这么设计,这里讲它要付什么钱:MoE 省的是 FLOPs,花的是带宽。AllToAll 的通信量正比于被路由的 token 数乘以 hidden 维度,且每层都要来一次。这也是 MoE 训练对网络要求远高于稠密模型的原因。
6.4 Token Dispatch 的工程细节¶
高性能实现不会简单地"发出去、算、收回来"。它会把 token 按目标专家排序分组(让每个专家收到的 token 连续,才能做一次大 GEMM)、把 AllToAll 拆块与本地专家计算重叠、用专门的通信库降低小消息开销。
正确实现一个 MoE 层不难,难的是让它跑得快。朴素实现和优化实现的差距可以到数倍,这也是为什么 MoE 训练框架的门槛比稠密模型高。
6.5 负载不均是系统问题也是算法问题¶
即使路由在整个训练集上是均衡的,单个 batch 内仍可能严重偏斜。某张卡分到的 token 是别人的三倍,其他卡就得等它——同步训练下这是直接的时间损失。
第一册 8.4 节讲的负载均衡损失和免辅助损失偏置调节,从系统视角看就是为了防止这件事。容量因子(每个专家最多接多少 token)和超出后的丢弃策略,则是在算法效果和系统确定性之间做的取舍:设死容量能保证不出现慢卡,代价是部分 token 没被处理。
核心问答¶
Q:PP 最大的问题是通信量吗? 不是。PP 的通信量其实很小(只传 stage 边界的激活),它的问题是气泡和 stage 不均衡。
Q:为什么 micro-batch 数要远大于 stage 数? 填充和排空的成本是固定的,micro-batch 越多,这部分固定成本被摊得越薄。公式 (P-1)/(M+P-1) 直接给出了这个关系。
Q:EP 为什么特别吃网络? 每层两次 token 级 AllToAll,且是全连接式的通信模式,跨机时对带宽和延迟都敏感。
Q:已经用了 EP,还需要 TP 吗? 可能需要。EP 只切专家,注意力部分和单个专家内部的矩阵如果还是太大,仍要靠 TP。大模型上常见 EP 和 TP 组合使用。