Kimi K3 模型结构（8）：训练并行，MoonEP 与四类等待

对应论文 §5.2 的开头、§5.2.1（MoonEP）、§5.2.3（视觉编码器）和附录 E。第 7 篇讲的是结构改动带来的系统工作，这一篇和后面两篇讲把 2.8T 参数训练起来、跑起来本身需要什么。三篇各用一个视角：这一篇看时间，第 9 篇看显存，第 10 篇看状态放在哪里。

先说清楚论文没给什么：§5 没有任何并行度、GPU 数、MFU 或吞吐。所以下文所有数字要么来自 config.json 的形状，要么并行度用字母代替。看到 $P$ 、 $R$ 、 $N$ 这些字母时，它们是变量，不是 K3 的真实配置。

序列 · KDA序列 · Gated MLA深度 · AttnRes宽度 · Stable LatentMoE输入 · MoonViT-V2零件

这一篇的主角是 MoE 层的 expert 并行，配角是视觉编码器怎么塞进流水线的空隙里。

并行布局：一句话

论文 §5.2 开头用一句话给出了预训练的并行组合：流水线并行（PP）带虚拟段（VP）、expert 并行（EP）、ZeRO-1 数据并行（DP）、Pipeline ZeRO-2 梯度分片（见 ZeRO 那篇）、上下文并行（CP，就是第 7 篇的 KCP）。MoE 层的 shared expert 在每个 EP rank 上都有一份副本；dispatch 和 combine 的 all-to-all 与计算重叠。没有张量并行。 这不奇怪：K3 每层的非 expert 参数只有几亿，routed expert 才是大头，而 expert 用 EP 切比用 TP 切自然得多。

这几种并行各切什么维度、各让谁等谁，是通用知识，放在「大模型笔记」专栏里：并行总览逐一讲了每种并行切什么、怎么通信，流水线并行推了气泡占比 $(P-1)/m$ 和每个 rank 在途激活 $P-r$ 份这两个后面要用的结论。这一篇默认你知道这两件事，只讲 K3 自己做的。

并行的所有问题都可以问成一句话：谁在等谁。K3 训练里的等待一共四类，下面先把它们列出来，再逐个看 §5.2.1 和 §5.2.3 怎么处理。

训练中的四类等待

把一个 PP stage 上一层 MoE 的执行顺序摊开：

一层 MoE 的执行顺序。三处等待分别藏在 dispatch（rank 等 rank）、GEMM 之前（host 等 device 拿 shape）和 GEMM 内部（SM 等 SM）。shared expert 在另一个 stream 上跑，把 all-to-all 的时间用掉。

沿着这条线数，K3 训练里的等待一共四类。§5.2.1 和 §5.2.3 的每个技巧都对应其中一类：

谁在等谁	在哪里发生	K3 的做法
rank 等 rank	EP：token 分配不均，最慢的 rank 决定这一层的时间	MoonEP 的完全平衡
	PP：warmup 和 cooldown 的气泡	ViT 的计算塞进气泡（§5.2.3）
host 等 device	每层 MoE 的 token 数每步不同，CPU 要等 GPU 算出来才能 launch expert GEMM	完全平衡之后 shape 静态已知，不再同步
SM 等 SM	一个 rank 内部，各 expert 的 token 数仍然偏斜，固定顺序的 GEMM 调度让部分 SM 先空下来	按当前 token 分布选参数的 GEMM 调度器
计算等通信	dispatch / combine 的 all-to-all；冗余 expert 的权重预取	与计算重叠；shared expert 放到另一个 stream 上

前三类是 MoonEP 一次解决的，下面从它的目标开始推。流水线气泡里塞 ViT 的部分放在本篇最后。

MoonEP：从「完全平衡」这个目标出发

常规 EP 的问题论文说得很直白：token 负载在 rank 之间不均衡，计算不均衡拖慢吞吐；而且 routed expert 的激活 shape 每步在变，造成大量显存碎片。两个问题的根源是同一个：每个 rank 收到多少 token 由 router 决定，事先不知道。

MoonEP 把目标定得很死：每个 rank 恰好收到 $S\times K$ 份 token。这里 $S$ 是 micro-batch 里的 token 数， $K = 16$ 是每个 token 选的 expert 数， $S\times K$ 是一个 rank 自己的 token 产生的「token 份数」。全部 $R$ 个 rank 一共 $R\cdot S\cdot K$ 份，平均下来每个 rank 正好 $S\times K$ 。目标是让每个 rank 算的份数等于这个平均值，而不是「差不多」。

达到这个目标的手段只有一个：冗余 expert。某个 expert 太热（分到的 token 太多），就在别的 rank 上临时放一份它的权重副本，把一部分发给它的 token 分流过去。K3 的做法是每一层、每个 micro-batch 都根据当前 router 输出在线规划冗余 expert 的放置，在 routed expert 计算之前把权重预取过去；反向时冗余副本的梯度先在本地 reduce 缓冲里暂存，算完再 reduce 回它的 home rank。

于是关键问题变成：要预留多少冗余 expert 的槽位，才能保证不管 router 怎么输出都能补到完全平衡？ 论文的回答是附录 E 的定理。

定理 1：每个 rank 最多 $E / R$ 个冗余 expert 就够

记 $E = 896$ 是 expert 总数， $R$ 是 EP 的大小，每个 rank 本地有 $E/R$ 个 expert。定理 1 说：对任意 router 输出，都存在一个方案，让每个 rank 恰好收到 $S\times K$ 份 token，并且每个 rank 上的冗余 expert 不超过 $E/R$ 个。

证明是构造性的，比想象中短。先把每个 rank 按本地 token 份数分成欠载（少于 $S\times K$ ）和过载（多于 $S\times K$ ）两类。然后重复下面这一步：

挑一个欠载的 rank $u$ 和一个过载的 rank $o$ 。
从 $o$ 搬 token 到 $u$ ，一次把 $u$ 补满到恰好 $S\times K$ 。 $o$ 搬走这些之后可能仍然过载、正好平衡，或者反过来变成欠载，把它放回对应的集合。

每一步让一个欠载 rank 变成平衡，而被补满的 rank 之后不再变动，所以最多 $R-1$ 步就结束。整个证明就靠这一点：每个 rank 最多被补一次，所以它收到的远端 token 全部来自同一个 rank $o$ 。而 $o$ 搬出去的 token 都属于它自己的本地 expert（它只在过载时往外搬，过载的 rank 从来没收过别人的 token），最多 $E/R$ 个。远端 token 涉及几个 expert， $u$ 就需要几个冗余 expert，所以不超过 $E/R$ 。证毕。

把这个过程动手走一遍：

每根柱子是一个 rank 上要算的 token 份数，颜色是 token 属于哪个 expert；实心是本地 expert，斜线是从别的 rank 搬来的、需要在本 rank 放一份冗余副本的 expert。虚线是平衡值 $S \times K = 6$ 。上限 $E / R = 2$ 。

多按几次「随机 router 输出」，冗余数从来不会超过 2。注意第二步里 $o$ 被搬到欠载的情况：它之后会被别的过载 rank 补一次，仍然只涉及一个来源。

定理 2：这个上界基本是紧的

上界会不会太松？论文构造了一个最坏例子：rank 0 的本地 expert 一份 token 都没收到，其余 $R-1$ 个 rank 的 $E(R-1)/R$ 个 expert 把全部 $R\cdot S\cdot K$ 份 token 平分，每个 expert 收到

\frac{R\cdot S\cdot K}{E(R-1)/R} = \frac{S K R^2}{E(R-1)}

份。不管怎么规划，rank 0 都要收 $S\times K$ 份远端 token，每个 expert 最多贡献上面这么多，所以至少涉及

\frac{S K}{S K R^2 / (E(R-1))} = \frac{E(R-1)}{R^2}

个 expert，取整后 rank 0 至少需要 $\lceil E(R-1)/R^2\rceil$ 个冗余 expert。 $R$ 大时这个数约等于 $E/R$ ，所以没有明显更小的通用上界。demo 里的「最坏情况」按钮就是这个构造： $E = 8$ ， $R = 4$ ，rank 0 至少要 $\lceil 8\times 3/16\rceil = 2$ 个，正好碰到上界。

代入 K3

$E = 896$ 。EP 大小论文没给，按几个可能的值算本地 expert 数，它同时也是每个 rank 要预留的冗余槽位上限：

$R$	每 rank 本地 expert 数 $= E/R$	预留的冗余槽位
32	28	28
56	16	16
112	8	8

也就是说，每个 rank 要为「最多再放一倍的 expert 权重」留出显存。这部分显存是 MoonEP 主动多花的，换来的是规划永远有解，训练永远不会因为找不到可行方案而停下。这正是它和前人的区别：

从朴素 EP 到 MoonEP，每一步只改一件事。MoonEP 的新东西是一个定理：每个 rank 预留

E / R

个冗余 expert 的槽位，就永远能规划出让所有 rank 计算量完全相同的方案。

论文点名的 ECHO 和 UltraEP 走的是另一条路，预设冗余 expert 的数量，或者给每个 rank 设一个 token 上限。上限之内没有可行方案时训练只能停下，上限本身要手调，而且调好了也还有残余的不均衡。MoonEP 用一个定理去掉了这个需要手调的参数。

在线规划、零拷贝、静态 shape

定理只保证方案存在，每一步真的去求最优解太贵。论文的做法是离线用整数线性规划算出一些代表性 case 的精确最优解当参照，然后设计一个 GPU 上的规划内核：接近最优、开销可以忽略、并且永远不超过 $E/R$ 的上界。所以实际用的是一个近似算法，定理保证的是：它找到的方案一定在 $E/R$ 这个预算之内。

完全平衡之后，通信路径也跟着简化。规划内核已经算出了每个 token 的目的地，于是 permute 和 unpermute 融合成一个算子，token 直接被送到远端 rank 上按 expert 分组好的位置，通信缓冲区直接交给计算使用，中间没有任何拷贝。这一点 DeepEP 也能做，但缓冲区要大得多：不知道每个 rank 会收到多少 token，就得按最坏情况准备。最坏情况是所有 token 都涌向同一个 rank，缓冲要 $S\times K\times R$ 份；MoonEP 每个 rank 恰好收 $S\times K$ 份，缓冲就是固定的 $S\times K$ ，差了 $R$ 倍。

完全平衡解决了「每个 rank 收多少」。通信路径上还有一处可以少发。先分清两个词。payload 是 token 向量本身，就是 expert 要乘的那 3,584 个数，BF16 下约 7 KB。weight 是 router 给一对 (token, expert) 的分数，一个 FP32，4 字节。两者都不是 expert 的权重矩阵，权重矩阵从不离开它的 home rank。

一个 token 的 16 条路由里，常有几条落在同一个目标 rank 上。目标 rank 只需要这个 token 向量一份。MoonEP 的编码是：每条路由条目存的是它在目标 rank 上的写入偏移；对同一 token 里目标 rank 已经出现过的条目，存 $-(\text{offset}) - 1$ 。发送路径遇到负值就跳过 payload，只发 4 字节分数。接收端把收到的那一行本地复制给其他 expert 的槽位。用符号位当标志，不多加一个字段。

一个 token 的 6 条路由落在 3 个 rank 上（示意，K3 是 16 选 896）。每条路由条目存它在目标 rank 上的写入偏移；同一 token 里目标 rank 已出现过的条目改存

- offset - 1

。目标 rank 收到唯一一份 token 向量后，本地复制给其他 expert 的槽位。

分数为什么不能一起省：combine 阶段每个 expert 的输出要乘各自的分数，router 给同一 token 的两个 expert 的分数不同，所以第二条仍带它的 4 字节，只是不带那 7 KB。DeepEP 和 MoonEP 都在减 dispatch 流量，走的是两条路：DeepEP 把每一行量化到 FP8 让行变小，MoonEP 让行变少且保持全精度。

顺带一提：MoonEP 的 C++ 只有 478 行，全是内存管理；所有 GPU 内核包括规划器都用 CUTLASS 的 CuTe DSL 以 Python 写成。

还有一个好处：CPU 不用再每层等 GPU。常规 MoE 实现里每一层每个 expert 收到多少 token 都在变，CPU 必须等 GPU 算出实际数字才能 launch expert 计算，每层同步一次，流水线在层与层之间停顿。完全平衡后每个 rank 每层都是 $S\times K$ 份，所有层的计算 shape 静态已知，这个同步就没有了，host 侧的 launch 开销也跟着减轻。

rank 内部：expert 之间仍然偏斜

rank 之间平衡了，rank 内部 $E/R$ 个 expert 各自收到多少 token 仍然是偏的。group GEMM 如果按固定顺序、不看负载地把工作分给 SM，这个偏斜就会变成 SM 之间完成时间的差距：有的 SM 早早做完，干等最后一个大 expert 算完。K3 的 routed expert GEMM 用一个按负载分配的调度器：launch 之前先看这一步每个 expert 分到多少 token，据此选好调度参数，执行中不再改。选参数用的是一个估算耗时的公式，输入是硬件指标，里面的关键系数事先离线调好。shared expert 的 GEMM 则直接放到另一个 stream 上，和主 stream 上的 all-to-all 与 routed GEMM 重叠。

和 QB 的关系

第 5 篇的 Quantile Balancing 也是在做负载均衡，两者不冲突，是分工。QB 通过 router 偏置改变路由本身，让 expert 的负载在统计上平衡，管的是模型训练得好不好；MoonEP 不改路由，接受 router 给出的任何结果，通过冗余副本让每一步的执行形状完全一致，管的是每一步跑得齐不齐。QB 做得越好，MoonEP 需要的冗余越少；但只要 QB 不是逐 micro-batch 精确的（它本来就不是），MoonEP 就有活干。

视觉编码器：塞进气泡里

§5.2.3 讲的是另一类 rank 等 rank：流水线 warmup 和 cooldown 的气泡。这些空隙通常是浪费掉的，K3 往里面塞了视觉编码器的计算。先把时间线放在这里，切到「ViT 填气泡」视图看红色方块：

PP 段数 P = 4micro-batch 数 m = 8每 rank 虚拟段 V = 1

谁在等谁在途激活 ViT 填气泡

前向（1 单位）反向（2 单位）气泡：这个 rank 在等别人ViT 计算

用同一个 micro-batch 的深浅表示它在流水线里的先后； $V > 1$ 时颜色的明暗区分同一 rank 上的不同虚拟段。1F1B 的气泡时间占比理论值是 $(P - 1) / (mV + P - 1)$ ，通常近似写成 $(P - 1) / (mV)$ ；读数里给的是 rank 0 的模拟值。

多模态训练里 MoonViT-V2 的工作量随样本剧烈变化，一张大图或一段长视频会让编码器时间暴涨，并且让持有它的设备成为拖后腿的那个。K2.5 引入过 Decoupled Encoder Process，把 ViT 和文本训练拆成独立阶段，并把视觉的前向和反向在 PP stage 之间摊平。K3 更进一步观察到了 interleaved 1F1B 的一个形状特征：头几个 micro-batch 的文本前向全挤在最开头，末几个 micro-batch 的文本反向全挤在最末尾。于是把 ViT 的计算拆开：头几个 micro-batch 的 ViT 前向同步地先做，其余的 ViT 前向塞进流水线气泡；反向对称处理。上面时间线里的红色方块就是这个安排的示意。结果是大部分 ViT 计算被藏在气泡里，视觉编码器的有效开销基本消失。

塞进气泡的前提是单个样本的编码器计算不能太长，否则塞不下，还会让持有大图的设备拖慢整条流水线。所以配套的是动态上下文并行：一张大图沿 patch 维切到多个设备上，注意力通过跨 CP rank 收集 KV（gather-KV）来算；每个 CP 组再切成若干子组，多张大图按负载均衡分配过去，让通信占比不随规模增长。动态上下文并行把单个样本的编码时间压下来，气泡调度把这些计算放进空闲时间，两者配合，视觉编码器才不再拖慢训练。

小结

这一篇的所有内容可以收进一张表：rank 等 rank、host 等 device、SM 等 SM、计算等通信。MoonEP 用一个定理解决前两类：每个 rank 预留 $E/R$ 个冗余 expert 槽位，任何 router 输出都能补到每个 rank 恰好 $S\times K$ 份 token；平衡带来的静态 shape 又顺手消掉了 host 同步和 $R$ 倍的通信缓冲。第三类靠按负载分配的 GEMM 调度，第四类靠 stream 重叠和把 ViT 塞进流水线气泡。

时间上的等待解决了，下一个问题是空间。流水线让 rank 0 攒着 $P$ 份甚至更多的激活，冗余 expert 让每个 rank 要为一倍的 expert 权重留位置，而 2.8T 参数本身的梯度和优化器状态还没算。第 9 篇从一张每卡显存账本开始。