Megatron-LM 的 MoE 并行：以 DeepSeek-V4 为例

本文沿 Megatron-LM 的源码，拆解一个 MoE 层里的参数放置、token 分发、专家计算与反向通信，再把它接回 DP、CP 和 PP。算例采用 DeepSeek-V4-Pro 的模型结构。

阅读 MoE 代码时，可以一直追踪两个对象：当前这行 token 属于哪个样本、哪个位置；当前这块权重属于哪个专家、哪个矩阵分片。通信改变了它们的位置，参数梯度的归约还要求参与者持有同一逻辑参数。把这两件事对上，大部分并行关系就能顺着代码推出来。

本文源码范围为 2026-10-26 的 Megatron-LM，链接固定到同一版本。DeepSeek 官方报告描述的是自有训练基础设施，下面的 Megatron 分组是机制算例，未做完整 V4 训练与性能验证。Megatron-LM、V4 训练框架

模型与入口

先固定后面会用到的模型维度。

这些维度可以对照官方配置。V4-Flash 的层数、hidden size 和专家数不同，本文计算统一使用 Pro。top-6 只包含 routed experts，shared expert 另算。

一个 token 的 MoE 输出写成：

$x_t\in\mathbb R^h$ 是输入，$I_t$ 是选中的专家集合，$g_{t,e}$ 是混合权重，$F_s$ 是 shared expert。先暂时忽略 dropout、padding 和量化，追踪这个式子需要哪些设备一起计算。

384 个专家的权重需要常驻或按需取回，而一个 token 只执行其中 6 个。以 SwiGLU 的三个矩阵计算，单个 routed expert 有：

个参数。这里忽略 bias 与量化元数据。参数的存储量与 token 的计算量按不同规则增长，这也是 MoE 需要单独安排专家并行的原因。

源码的主要入口如下。

先沿普通 AlltoAll 训练路径读。Flex、GTP 和虚拟专家在后面接上，每条支路仍用参数和 token 的归属来解释。

两套并行网格

Megatron 的 MoE Parallel Folding 允许 attention 和 routed experts 使用两套设备网格。论文 §3.2

设总设备数为 $N$，流水线并行度为 $P$。暂时取 GTP、EGTP 为 1，则同一批设备满足：

$D_a$ 是 attention 的数据并行度，$D_e$ 是 expert data parallelism，后面简称 EDP。几种并行分别切不同对象。

SP 使用 TP 组，不额外增加设备维度。EP 已经出现在专家网格中，也不能在 attention 网格的等式右边再乘一次。 当前普通 MoELayer.forward 还会检查训练时 attention TP>1 必须开启 SP；本文 V4 算例取 TP1，不触发这项条件。

为什么能这样折叠？attention 需要知道 token 属于哪条序列、能看哪些上下文；专家 FFN 对每一行独立计算。进入 MoE 后，来自不同序列或不同上下文分片的行，可以按专家重新分组。输出送回时，仍恢复原来的 token 位置。

这里还要区分 dp 和 dp_cp。CP ranks 处理的是同一序列的不同位置，参数共享，但它们各自产生的参数梯度需要相加。dense reducer 通常使用包含 CP 的 dp_cp。expt_dp 已经按专家网格建立，不能再凭习惯给它额外乘一次 CP。

一个分组算例

固定下面一组数值，用于后面的 V4 数据流计算。

N = 1024       PP = 8
TP = 1         CP = 2
ETP = 1        EP = 64
GTP = EGTP = 1

于是：

每个 stage 的 128 张卡，在 attention 看来是 64 个样本副本，每个副本用 2 卡处理上下文；在 expert 看来是 2 份完整专家集合，每份分布在 64 卡上。每 rank 每个 MoE 层保存 $384/64=6$ 个 routed experts。 常规等分实现要求专家数能被 EP 整除，专家矩阵相应维度也要满足 ETP 的整除与算子对齐条件。

这组数只证明设备数能对上。61 个块怎样划给 8 个 stage、是否包含 MTP、显存能否装下，都需要另行安排。V4 的 HybridModel 还使用 layer pattern 描述模型与流水线边界，不能直接把这段数值当成启动命令。模型构建

进程组的成员

initialize_model_parallel 分别建立 dense 和 expert 的 RankGenerator，再检查两套网格的 PP 组一致。普通布局下，dense generator 的 EP 轴取 1，expert generator 的 CP 轴取 1。组初始化

沿用默认顺序 tp-cp-ep-dp-pp，并忽略大小为 1 的轴。stage 0 的 rank 为 0 到 127，可以同时用两种坐标表示：

这里 $c\in[0,1]$，$d_a\in[0,63]$，$e\in[0,63]$，$d_e\in[0,1]$。注意 $e$ 在这个等式中是 EP rank，模型的全局 expert ID 另用下标说明。

在连续专家分配下，EP rank 0 保存 expert 0–5，EP rank 1 保存 expert 6–11。全局 rank 0 和 64 都保存本 stage 的 expert 0–5；rank 0 和 1 则保存不同专家，不能把它们的权重梯度逐元素平均。

CP 组和 EP 组存在重叠很正常。一个 process 可以参与多种通信组，程序根据当前操作选择 group。逻辑 rank 邻近也不保证物理上位于同一节点，实际映射仍取决于 launcher 与硬件拓扑。

MoELayer 的调用链

MoELayer.forward 的普通训练流程可以缩成下面的伪代码。这里保留调用关系，省略 CUDA Graph、重算和推理分支，不是可直接运行的实现。

shared = shared_experts_compute(x)
weights, routes = route(x, input_ids=ids)
send_x, send_p = preprocess(x, weights, routes)
recv_x, recv_p = dispatch(send_x, send_p)
partial_y, bias = routed_experts_compute(recv_x, recv_p)
returned_y = combine(partial_y)
y = postprocess(returned_y, shared)

shared 分支的实际启动位置可以调整。routed_experts_compute 内部还包含 dispatcher 的后处理和 combine 预处理，不能把它理解成一个单独 GEMM。前向调用

dispatcher 的六个阶段全在前向执行：

combine 需要在前向把专家结果送回，模型才能继续执行下一层。随后 autograd 再为这些操作建立反向。

读代码时，还有一个影响后文所有 shape 的字段：

self.tp_group = pg_collection.expt_tp

这行在 dispatcher 基类里。因此 dispatcher 的 tp_size 指 ETP；shared expert 中的 tp_group 则来自 pg_collection.tp。两处同名变量需要按构造位置解释。dispatcher 基类、shared expert 构造

路由与负载

路由张量

设进入 router 的激活为 $[s_{local},b,h]$，展平后有 $n=s_{local}b$ 行。gating linear 把 $[n,h]$ 映射为 $[n,E]$ logits。

后面的 $n$ 都指实际 tensor 行数。如果输入已经按 CP、SP 切分，就直接数本地行，避免重复除并行度。普通 TP+SP 布局中可写成 $n=bS/(CP\,TP)$；本文 V4 算例 TP1，因此是 $bS/CP$。

V4 使用的 affinity 和选择过程可以表示为：

$\beta_e$ 是负载控制用的 expert bias，$\alpha$ 是路由缩放系数。选择时使用偏置，混合权重取原始 affinity；expert bias 也与 gating linear 的普通 bias 不同。官方 Gate、Megatron 路由工具

普通 AlltoAll 路径输出 probs[n,E] 和 bool routing_map[n,E]。前者在选中位置存权重，其余为零；后者表示 assignment。部分 Flex 路径改用紧凑的 indices[n,k]，阅读后端时要核对具体格式，不能把所有 routing_map 都当 bool 矩阵。

Hash 路由

V4 前三层由 token ID 查表取得专家 ID，后面的专家通信仍然存在。相同 token ID 在不同上下文里会得到不同 hidden states，专家输出和混合权重也可能不同。

Megatron 的 _hash_routing 保留了 gating 与权重计算，只替换专家选择。还有一个容易漏掉的排列关系：hidden 从 [S,B,H] 展平，而 input_ids 以 [B,S] 传入，因此源码先 input_ids.T.reshape(-1)，使 token ID 与 hidden 的每行对齐。Hash 实现

当前初始化中的 round-robin tid2eid 是带 warning 的占位表。使用已有 V4 checkpoint 时需要正确加载表；从头训练则需要明确初始化策略。占位公式本身不能当成 DeepSeek 公布的 Hash 构造算法。表初始化

均衡统计

router 的通信不只有发送 token。它还可能汇总专家计数、计算 aux loss，或更新 expert bias。

local aux loss 的统计默认涉及 TP×CP，global aux loss 将范围扩到 TP×DP×CP，并维护累积步的统计。expert bias 也汇总 global batch 的负载，按更新规则调整偏置，走的是统计更新，不靠 optimizer 对 bias 求梯度。统计组、bias 更新

MoEAuxLossAutoScaler 在 forward 原样传递主激活，在 backward 给 aux loss 接上缩放后的梯度。所以只在主模型里搜索 loss += aux_loss，可能找不到完整的训练目标。辅助损失接图

group-limited routing 先选专家组，再在选中组里选 top-k。组只有与专家的物理放置对应时，才可解释成限制目的节点。V4 已取消路由目标节点数限制，不应把 V3 的 node-limited 设置放进 V4 的必要条件。

最后区分三件事。Dropless 保留有效路由，允许各专家 token 数变化；capacity dropping 会移除超额 assignment；量化 padding 则补齐矩阵行数。训练样本数、有效 assignment 数、GEMM 实际执行行数需要分别统计。

AlltoAll 的前向

接下来读 MoEAlltoAllTokenDispatcher。先取 ETP1、dropless、没有额外 padding，hidden 为 $[n,h]$，每个 token 恰好有 $k$ 个 assignment。

计数与重排

preprocess 沿 routing_map 的 token 维求和，得到每个专家的 assignment 数。假设每 rank 放 $L=E/EP$ 个连续专家，本地发往目标 EP rank $j$ 的行数为：

计数在 ETP×EP 域内交换，再推导当前 rank 从各来源接收多少行，形成 output_splits；ETP>1 时还需要后续 AllGather 的变长信息。计数与 split

第一次 permute 把同一专家的行排到一起，并复制多路选择的 token：

probs 按相同顺序重排。由于同一个 rank 的专家 ID 连续，按 expert 排序后，也就形成按目标 rank 分块的发送缓冲。

先用四个 token 算一次 split。令 EP2、ETP1、专家数4，rank 0 保存 expert 0/1，rank 1 保存 expert 2/3，每个 token 选两个专家。

rank 0 发送给两个目的 rank 的行数为 [2,2]，rank 1 为 [1,3]。因此 rank 0 的 output_splits 为 [2,1]，rank 1 为 [2,3]，实际各接收3行和5行。两边原本各有2个 token、各产生4个 assignment，接收量仍然可以不同。返回后，a 的原位置累加 expert 0 与2的加权结果，其他 token 同理。

这里的“permutation”包含复制。$k>1$ 时，输入输出行数不同，后面求反向时不能把它当成方阵的可逆置换。

EP 交换

token_dispatch 用 self.ep_group 执行可变长 AlltoAll，hidden 和概率各自交换。所有参与者遵循同一组的 rank 顺序与 split 约定，目标设备只接收分配给本地专家的 assignment。dispatch

以两 rank、每 rank 两个专家为例，接收缓冲可能排列为：

来源 rank 0: expert A 的块 | expert B 的块
来源 rank 1: expert A 的块 | expert B 的块

GEMM 需要的却是：

expert A: 来源 0 的块 | 来源 1 的块
expert B: 来源 0 的块 | 来源 1 的块

这就是 dispatch_postprocess 再做一次 chunk 排序的原因。第一次重排服务发送目的地，第二次重排服务本地专家计算，二者解决不同的布局问题。接收后处理

普通可变长 collective 还需要 host 可读的 split。代码通过专门的 DtoH stream 和延后同步覆盖部分开销。即使 dropless 下总发送量 $nk$ 已知，各目标 rank 的动态接收量仍可能需要同步；总行数静态不代表整个路径没有 GPU→CPU 等待。

ETP 里的矩阵

EP 决定本地有哪些专家，ETP 决定每个专家的矩阵如何切开。

采用数学上的 $XW$ 记号，一个 SwiGLU expert 为：

为突出矩阵分片，这个式子省略 V4 的激活 clamp；clamp 改变逐元素计算，不改变后面的 shape。设 ETP=$q$，expert $e$ 收到 $m_e$ 行。在不加 bias 的情况下：

源码实际存储可采用转置的权重布局，表格按计算含义表示。每个 ETP rank 需要同一个专家的完整输入行，只计算自己那份 intermediate。FC2 的结果还要沿 ETP 相加，才能得到完整输出。专家实现、TE 的显式专家通信

普通 AlltoAll dispatcher 在 EP 交换后执行 ETP AllGather，在 expert 计算后执行 ETP ReduceScatter。因此整条前向为：

permute
→ EP AlltoAll
→ ETP AllGather
→ sort by local expert
→ expert FC1 / activation / FC2
→ undo local sort
→ ETP ReduceScatter
→ EP AlltoAll
→ unpermute

ETP 的 AllGather 收集 token 行，ReduceScatter 同时求和 FC2 部分结果、恢复 token 分片。hidden 的 $h$ 维在这两步中不被切成 $h/q$。 ETP1 时这两次 collective 退化；EP1 时 EP 交换退化。EP=ETP=1 仍然需要本地路由、重排、专家计算与加权累加。

再用4张卡看 ETP2×EP2 的组。默认顺序下，ETP 组为 [0,1]、[2,3]，EP 组为 [0,2]、[1,3]。rank 0/1 保存同一批专家的不同矩阵片：先在各自 EP 组收取 token，再在 [0,1] 收齐这批专家的输入，计算后沿该 ETP 组归约。这个小例子只解释专家侧分组，不作为 V4 attention TP>1 的运行配置。

TEGroupedMLP 将本 rank 的多个专家组织为 Grouped GEMM。各专家的权重独立，$m_e$ 也可不同。Grouped GEMM 是本地计算组织方式，不额外建立一个 GPU 并行维度。

TE 的 explicit expert communication 路径还会调整本地矩阵尺寸并关闭对应 linear 的自动 TP 通信，因为 dispatcher 已经完成 ETP 的通信。阅读时不能再为这些 linear 凭空补一轮 AllReduce。

概率在哪里乘

数学式写的是 $g_eF_e(X)$，代码不一定等到 combine 才乘 $g_e$。无 bias 时，FC2 是线性映射：

TEGroupedMLP 可以在 SwiGLU 激活后融合乘 permuted_probs，再执行 FC2。沿这条路径，combine 只负责把已经加权的结果汇回；再乘一次就会得到错误的平方权重。加权激活

这个等价变形依赖线性位置。不能把概率随意移到 SiLU 之前，非线性通常不满足同样关系。源码对 moe_apply_probs_on_input 也有单独限制。

结果怎样送回

combine_preprocess 先撤销第二次 chunk 排序，再执行 ETP ReduceScatter。之后 token_combine 沿 EP 返回，前后两次 AlltoAll 使用交换后的 splits：

dispatch: output_splits, input_splits
combine:  input_splits,  output_splits

combine_postprocess 撤销第一次重排，将同一 token 的 $k$ 份结果累加到原位置，再恢复 $[s_{local},b,h]$。返回路径

这里建议直接看执行语句。当前 AlltoAll 类的 token_combine docstring 仍提到 DeepEP，局部 shape 注释也残留了 $H\to H/TP$ 的描述。该方法实际调用的是 all_to_all(self.ep_group, ...)，保持 trailing hidden 维度，ETP 归约已在前面完成。函数名和注释可以帮助定位，shape 需要沿 tensor 操作核对。

一轮 V4 的数据量

回到前面的 1024 卡算例，取一个 microbatch 的 packed 序列长度为 8192，batch size 为 1，CP2。V4 当前 CP 路径使用 packed THD，激活可表示为 [T_local,1,H]；这里每个 rank 有 4096 个有效 token 行，便于按固定长度计算。V4 CP 输入

一个 EP64 组覆盖的唯一 token 数为：

每 token 产生 6 个 routed assignment，全组为 $1,572,864$ 个。假设路由均匀，单个 expert 的期望接收量为：

每 rank 有 6 个 expert，期望接收 $6\times4096=24,576$ 行。第一次 permute 的发送缓冲固定有 24,576 行，接收缓冲只有在均衡假设下才同样大。

这些都是形状推导，没有把期望值当成实际负载。一个 rank 保存多个专家，既要看单专家热点，也要看该 rank 所有专家接收量之和。

BF16 hidden 每元素 2 byte，单方向、包含留在本 rank 的 assignment，其 payload 为：

普通 A2A、均匀路由下，约 $63/64$ 的 assignment 需要离开当前 rank，对应 330.75 MiB 的远程发送量。前向 dispatch 与 combine 合计约 661.5 MiB；若反向采用相同元素字节数，四次 EP hidden 交换合计约 1323 MiB。

这个口径统计发送 payload，不把接收再加一遍。它未含概率、计数、padding、协议开销、重算和 ETP 通信，也没套用融合后端的去重。计算图中有 6 个 assignment，不能直接推出某种分层网络实现必定跨节点发送 6 份完整 token。

权重也可以对照计算。每 rank 每层的 6 个 routed experts，BF16 权重为：

如果某 stage 放置 8 个这样的 MoE 层，routed 权重约 5.906 GiB。还需要加 shared expert、attention、梯度、优化器状态与通信缓冲。这部分账只说明专家怎样放置，不证明完整训练能装下。

另外两类 dispatcher

AllGather

MoEAllGatherTokenDispatcher 沿 ETP×EP 组收集全部 hidden、概率和路由图。每个 rank 从完整 token 集中选择自己负责的专家，计算后先恢复到全域 token 位置，再用同组 ReduceScatter 求和并归还原来的 token 分片。AllGather 实现

它同样会跨 GPU 搬 token。与 AlltoAll 的区别在于，hidden 收集阶段没有先按目的 expert 做稀疏过滤。每个 rank 的临时输入会随组内总 token 数增加，之后才筛出本地工作。

ETP1、EP=$p$、每 rank 有 $n$ 行时，AllGather 后每 rank 的 hidden 是 $[pn,h]$；AlltoAll 的均衡接收期望则为 $[kn,h]$。这是布局与存储量比较，不能单凭 $p$ 和 $k$ 就给出后端耗时结论，collective 实现、消息大小和拓扑还会影响结果。

Flex

MoEFlexTokenDispatcher 把 ETP×EP 当成一个通信域。_initialize_metadata 将逻辑专家展开为对应的物理矩阵分片目的地，后端统一执行 dispatch/combine。Flex 初始化

例如逻辑 top-6、ETP2，需要把每个被选 expert 的输入提供给它的两个矩阵分片。manager 中可能出现 topk=12、num_experts=768 这样的数值，它们编码的是物理目的地；模型仍然是 384 个专家、每 token 选 6 个。

当前源码包含四种 manager。

Flex 不要求 Python 层逐项执行“EP A2A→ETP AG→ETP RS”。后端可以把它们融合成另一种物理实现，比如跨节点先去重、节点内再展开，combine 时先局部求和再跨节点返回。HybridEP 工程说明

NCCL EP 的 zero-copy 路径还允许专家 FC2 输出、FC1 输入梯度直接写入后续通信缓冲，省掉中间复制。它改变缓冲的所有权和生命周期，逻辑上仍需完成专家结果汇总。通信缓冲接口

这些训练路径需要与推理 decode 的低延迟 kernel 区分。DeepEP 提供多个场景的接口，选一个库名不足以说明采用了哪种执行策略。

反向的通信

现在给 MoE 输出一个上游梯度 $\bar y_t=\partial L/\partial y_t$。对选中的 expert，有：

$x_t$ 还影响路由权重，所以输入梯度包含 router 分支，以及 shared expert 分支。Top-k 的离散 ID 不求连续导数，选中分支的连续 score 与混合权重仍然可导，Hash 层也保留这部分学习。

从通信上看，反向按前向的依赖倒序执行。

MoE 输出梯度
→ unpermute 的反向：展开到 assignment
→ combine A2A 的反向：送到专家侧
→ ETP RS 的反向：AG
→ expert backward：dgrad 与 wgrad
→ ETP AG 的反向：RS
→ dispatch A2A 的反向：返回 token owner
→ permute 的反向：累加多路 dgrad
→ router 与原输入

$dgrad$ 表示对输入激活的梯度，$wgrad$ 表示对参数的梯度。多个专家对同一 token 的输入梯度要相加；这是第一次 permute 包含复制的直接结果。

源码中 _AllToAll.backward 再调用一次 AlltoAll，并交换 input/output split。默认 token 维 AllGather 的 backward 是 ReduceScatter，ReduceScatter 的 backward 是 AllGather。autograd 封装

不过 gather_from_sequence_parallel_region 也有 tensor_parallel_output_grad=False 的 split-only 分支。它服务不同的梯度语义，所以阅读 collective 的反向要连同 wrapper 参数一起看。

到这里，梯度已经沿计算图回到应该在的位置。每份参数在不同数据副本上产生的贡献，还要接着同步。

参数梯度同步

dense 与 expert 的归约组

routed expert 参数上经常能看到：

param.allreduce = False

这个标记用于选择 expert 拓扑。选择 expert 拓扑的条件也不只有 EP>1，ETP 与 TP、EGTP 与 GTP 的差异同样会影响它。DDP 会把这些参数放进独立 buffer，沿 expt_dp 同步；普通参数放进 dense buffer，沿 dp_cp 同步。参数标记、buffer 分组

回到 rank 示例，rank 0 和 64 的 expert 0–5 相同，它们产生的对应梯度需要合并。rank 0 和 1 分别保存 expert 0–5、6–11，参数身份不同，不能直接做同一个权重 AllReduce。

router 的全专家 gating 矩阵与 shared expert 通常归入 dense 路径。TP+SP 时，router 和 LayerNorm 等复制参数还需要收齐 TP ranks 上不同 token 的梯度，finalize_model_grads 对这类标记参数执行额外 TP 归约。不同 PP stage 通常持有不同层，普通 router 权重并不会统一跨 PP 平均；tied embedding 等共享参数有专用同步。梯度收尾

为什么不能直接除 EDP

专家副本通过 dispatch 接收了多个 dense 数据位置的 token。它的本地 $wgrad$ 已包含这些输入的贡献，归一化不能只看参与最后一次 collective 的 EDP 大小。

固定 GTP/EGTP1、关闭 per-token loss，DDP 的 SUM 路径对 dense 和 expert 都预乘：

若 collective 本身使用 AVG，expert 则先乘：

再由 EDP AVG 除以 $|EDP|$，得到相同的归一化口径。源码中的两种分支可以直接对照。缩放实现

本文例子 $|DP\times CP|=128$、$|EDP|=2$。SUM 分支的 factor 为 $1/128$；AVG 分支的 expert factor 为 $2/128$，之后在两副本间平均。这里也不能再手动对 expert 梯度除一次 EDP。

per-token loss 使用另一条路径：先保留 SUM，再按有效 token 总数缩放。流水线的 loss scale、梯度累积和 GTP 修正也各有位置。比较两个并行配置的梯度时，需要把这些缩放连起来，不能只看最后的 reducer。

优化器的通信

普通 DDP 在相应副本组内 AllReduce 梯度；distributed optimizer 则通常先 ReduceScatter，只把负责的梯度片交给本 rank 的 optimizer state。更新后再 AllGather 参数，让下一次计算获得所需权重。梯度 bucket、优化器构建

因此同一段 profile 中的 AllGather 可能搬 token，也可能搬更新后的参数。看 kernel 名称不足以判定它属于哪种并行。

多个 distributed optimizer instances 还会把副本域分层：instance 内分片，instance 间同步相同 shard。它细分现有副本域，不额外增加一个独立 world-size 因子。

Megatron-FSDP 进一步允许参数本身在副本域分片，计算前取回、使用后释放。dense 与 routed expert 分别使用 DP×CP 和 EDP mesh；HSDP、HFSDP 再区分内层分片与外层副本。它们需要接在已有 TP/EP 分片上理解。FSDP mesh

shared expert

shared expert 对每个有效输入 token 都执行，不参加 routed top-k。它按 dense TP 组织，因此 TP 和 ETP 不同时，shared 与 routed 专家的矩阵布局也不同。

一个 shared expert 是一组逻辑参数，可以沿数据副本复制、沿 TP 分片。这个“1”不能解释成所有设备只在某一张卡保存共享分支。

开启 shared-expert overlap 后，SharedExpertMLP 被拆成前置通信、FC1 与激活、FC2、后置通信。dispatcher 在 EP dispatch/combine 周围插入这些操作，通过 stream 与 event 保证最终相加前结果就绪。shared overlap

普通 AlltoAll 路径还专门安排了发射顺序：

token A2A → shared FC1 → probability A2A

代码这样处理，是为了避免某些连接数设置下，概率通信先排入队列而阻塞 shared FC1 的发射。combine 一侧再安排 shared FC2 与后置通信。它需要控制具体依赖和排队顺序，只有 async_op=True 并不能完成这些调度。

CP 与 PP

CP 的两条数据流

CP ranks 在 attention 里交换上下文；进入 MoE 后，FFN 可以对本地 token 独立路由。expert generator 中 CP=1，只表示专家网格不再单独保留这个轴，原来的局部 token 仍由 dispatcher 接管。

V4 的 CSA/HCA 还需要专门处理压缩边界。官方报告描述的 CP 先向相邻 rank 发送尾部未压缩 KV，使跨边界的压缩块能够计算；之后 AllGather 压缩 KV，再整理可见范围与 padding。V4 的压缩注意力 CP

这条 attention 数据流与 MoE 的 EP dispatch 作用不同，可以在同一个 rank 上先后出现。前者建立上下文，后者按 expert 重组 token。

当前 Megatron V4 attention 要求 TP1；CP>1 时要求 packed THD，CP 布局使用 contiguous，且只支持 static CP。通用 MoE 支持 CP，不代表任意 attention 的 CP backend 都能直接套在 V4 上。attention 校验、packed CP 校验

PP 的边界

PP 把不同模型块分配到 stage，前向传激活，反向传输入梯度。普通 MoE 的 EP 通信发生在本 stage 的设备中，结果恢复之后才进入下一段模型。

V4 采用 mHC，残差流的扩张会增加某些 PP 边界的激活与通信量。expert 内部的 $[n,h]$ 是经过相应映射后的输入，不能把它当成全部 PP 边界的 tensor shape。官方训练框架也为 mHC 调整了 DualPipe 的重叠方案。mHC 与流水线

VPP 让一个物理 stage 持有多个模型 chunk，增加可交错的任务。它改变执行时序和激活生命周期，专家的全局编号与参数身份仍然需要保持一致。

在本文普通布局、固定 packing 口径下，设全局 batch 为 $B$，microbatch size 为 $b$，每次更新累积 $m$ 个 microbatch：

CP 切同一条序列，EP/ETP 切专家计算，它们都不额外增加独立样本。取 $B=4096$、$b=1$，上面的 $D_a=64$ 对应 $m=64$。这个数会影响流水线能否连续取得工作，不能用 EDP2 代入得到 2048 个 microbatch。

重叠与重算

同一份源码中有几种不同的 overlap。

以专家 linear 的 backward 为例，输入梯度 $dX=\bar YW^T$ 决定上一层何时能继续；参数梯度 $dW=X^T\bar Y$ 可以在满足 optimizer 依赖的前提下延后。backward_dw、CUDA event 和延后 wgrad stream 为此提供接口。延后专家 wgrad

Megatron 的 combined_1f1b 将 attention、route、expert 等任务拆开组织。调度实现 DeepSeek 的 DualPipe 有自己的模型与调度接口，不能仅凭都能重叠通信，就认定两者是同一个实现。

当前版本也确实存在组合限制：Hash MoE 禁止 overlap_moe_expert_parallel_comm；mHC 与 PP>1 的组合也拒绝该开关。shared overlap 与 batch-level EP overlap 还存在互斥条件。因此本文 V4 配置不能直接抄一组通用 EP overlap flags。配置检查

重算同样需要沿调用边界分析。重算整个 MoE 可能再次执行路由、重排与通信；只重算局部激活则有不同代价。router 计数和状态更新还要避免被重算重复执行，源码为部分统计检查 checkpoint 状态。

CUDA Graph 要处理动态接收行数与 host 同步。capacity padding 可以让形状固定，但 dropping 会改变路由，padding 会增加计算；sync-free kernel、固定上限缓冲与分段捕获则采用其他代价。开启 graph 后仍应逐段看 dispatcher，而不能默认动态通信已经消失。融合算子说明

GTP 与 EGTP

前面固定了 GTP/EGTP1。当前源码还有一条权重重物化路径，需要与常驻矩阵分片的 TP/ETP 分开理解。GTP 设计与实现

设 dense 的 GTP rematerialization 轴大小为 $R$，expert 对应轴为 $R_e$。先不启用额外序列折叠，权重视图为：

$C$ 是配置的 CP，$D_w$ 是不含 GTP 轴的 dense 权重副本维度。这里换用 $D_w$，避免与前面普通布局的样本并行度混用。

ETP=$q$、EGTP=$R_e$ 时，同一逻辑 expert 的权重永久分成 $qR_e$ 份。执行前，在 expt_gtp_remat 组中 AllGather，恢复当前 ETP rank 的矩阵片；GEMM 仍按 ETP 的 intermediate 分片执行。反向产生的完整临时 wgrad 再沿 EGTP ReduceScatter 回永久权重片。

永久 expert weight shard
→ EGTP AllGather
→ 当前 ETP rank 所需的 weight slice
→ expert GEMM
→ 临时 wgrad
→ EGTP ReduceScatter
→ 永久 grad shard
→ 真正副本域的 EDP 归约

这里同时存在三种动作：EP 把不同专家放到不同设备；ETP 让多个设备计算一个专家的不同矩阵片；EGTP 在执行时临时取回更完整的权重。EGTP 的 AllGather 搬权重，dispatcher 的 ETP AllGather 搬 token，二者不能合并记账。

实现可以从 GTPShardedParam.materialize_group_for_forward/backward 和 finalize_group_grads 往下读，分别接到 all_gather_and_prefetch、wgrad_reduce_scatter。重物化实现

router、norm 等未做 GTP 权重切片的参数，还需要在相应 GTP 轴上归约；已经通过重物化 wgrad ReduceScatter 完成该轴的参数会被区别处理。DDP 中也有 $R_e/R$ 的 expert gradient correction，不能把普通布局的缩放公式无条件沿用。GTP 梯度收尾

权重与 token 双视图

GTP 还允许在 $R$ 轴内部继续分序列。令 $q_s$ 为 gtp_remat_num_sequence_shards，要求 $q_s\mid R$，则：

$R$ 的一部分处理同一样本的不同序列片，剩余部分处理不同样本。batch 关系变成：

权重网格依然是 $TP\times R\times C\times D_w\times P$。如果把已扩大的 $CP_{runtime}$ 再与完整 $R$ 相乘，就重复计算了 $q_s$。布局对象

例如官方文档的 128 卡、TP1、PP1、$R=64$、$C=2$、$q_s=64$，得到 runtime CP128、$D_w=1$、$D_{batch}=1$。这些设备处理同一 microbatch 的不同上下文分片，权重则在两份副本内各分 64 片。

这是通用 GTP 布局例子。当前 GTP sequence sharding 要求 Megatron DDP 与 torch_dist，dynamic CP、FSDP 和部分 packing scheduler 有明确限制；不能据此声称已验证完整 V4 的任意 GTP 组合。

虚拟专家

静态 EP 中，专家 ID 对应固定 owner。动态路由可能让某个 owner 收到很多 token，即使其他 rank 空闲，也不能直接拿这些 token 去算别的专家。

当前 HybridEP 接入的虚拟专家负载均衡允许给热点逻辑 expert 临时分配额外执行槽。planner 选择要复制的专家权重，搬到 peer 的临时槽，再把部分 assignment 发过去。反向生成的临时副本 wgrad 回收到原 owner，之后继续正常的参数同步。负载均衡实现

逻辑专家数和模型 top-k 保持原意。临时执行副本不注册为另一套独立可训练参数，否则 optimizer 会把同一个专家的两份计算更新成不同模型。

它同时移动计算和被选权重，因此要一起核算权重传输、临时缓冲与梯度回流，不能只看 token 负载变均匀。当前路径有 ETP1、固定本地 token 数、禁用 expert bias 和 FP4 等配置要求，FSDP 及部分 overlap 也有限制。本文把它作为通用实现支路介绍，未将其当成 V4 原始配置可直接开启的选项。当前约束

容量与低维通信

capacity 的作用范围

前面的 shape 都假设 dropless。普通 AlltoAll 的 drop-and-pad 路径会按当前来源 rank 的 assignment 预算计算每专家容量：

第一次 permute 输出固定为 $Ec$ 行。接收端一个本地专家会汇集多个来源的容量块，不能把单个来源的 $c$ 当成全 EP 域最终的专家总容量。容量分支

dropping 只删除某些 token→expert assignment，不一定删除这条 token 的全部分支，shared expert 也仍有自己的路径。改变并行分组后，局部容量预算与竞争集合可能改变，dropping 的数值等价不能只靠检查总 token 数证明。

Latent MoE

MoELayer.preprocess/postprocess 还有可选的 latent projection：先把 routed 分支从 $h$ 映射到更小维度，dispatch、expert、combine 在该维度工作，返回后再投影回 $h$。投影分支

router 和 shared expert 仍需按各自的输入路径处理。通信 payload 因 routed hidden 宽度改变而变化，但多了投影参数与计算；它属于模型结构选项。本文 V4 算例没有启用这一选项，不能因为同在 MoELayer 内就替 V4 自动加上。

Muon 的另一轮 A2A

V4 官方训练还提供了一个很适合对照阅读的例子：同样叫 AlltoAll，搬的却是参数梯度。V4 Muon 实现说明

Muon 更新需要逻辑完整的矩阵。官方为 MoE 参数安排保持独立专家矩阵完整的 bucket，再按相同形状批量执行 Newton–Schulz。梯度以随机舍入方式压到 BF16 传输，通过 A2A 交换后在本地 FP32 累加，避免低精度反复归约的误差。

这轮 A2A 合并的是同一参数的梯度贡献。前面 dispatcher 的 A2A 则按专家分配激活与输入梯度，二者的张量、分组和调用时机都需要分别确认。本文不把官方 Muon 同步算法写成 Megatron 普通 distributed optimizer 的默认行为。

精度也要按阶段说明。V4 报告的 MXFP4 QAT 位于后训练部分，与 Megatron 的 NVFP4 原生训练支持属于不同实现。checkpoint 中的 FP4 权重类型不能证明预训练从头到尾都采用 FP4，也不能直接据此把本文的 BF16 payload 除以四。V4 QAT

保存与恢复

并行度改变后，本地 expert 下标不一定保持原意。EP64 的 local expert 0，恢复到 EP32 后，需要依据全局 expert ID 定位，再处理 ETP 与状态分片。专家 checkpoint

Megatron 的 distributed checkpoint 记录逻辑 shape、分片偏移和 replica 信息，加载时据此重分片。完整恢复还涉及 router 权重、Hash 表、expert bias、optimizer state，以及所需 RNG 状态。只按旧 rank 文件名复制数据，无法表达这些对应关系。checkpoint 包

MTP 会增加预测模块或重复使用的计算路径，也可能再次调用 MoE。它不增加一个独立设备维度，但会改变每个 step 的工作量、参数归属与调度。阅读 layout 和 checkpoint 时，应把主模型与 MTP 明确对应到各自模块，不能只按主体 61 个块猜测全部参数位置。

串起一次更新

最后把前面的路径按一次 optimizer update 接起来。仍取普通 folding、ETP1、dropless 的 V4 算例。

64 个 microbatch，逐个进入当前 pipeline 副本
  │
  ├─ attention：按 CP 取得上下文
  ├─ router：本地 token → expert IDs + continuous weights
  ├─ dispatch：重排 → EP 交换 → 本地专家排序
  ├─ experts：FC1 → weighted SwiGLU → FC2
  ├─ combine：撤销排序 → EP 返回 → 恢复原 token
  ├─ shared：对原输入计算，与 routed 结果相加
  └─ backward：沿上述依赖返回 dgrad，累积 wgrad
  │
  ├─ dense wgrad：DP×CP 归约
  ├─ routed wgrad：EDP 归约
  ├─ 处理 TP 复制参数、共享参数与路由统计
  └─ optimizer：更新状态，按所选分片策略同步参数

实际流水线会交错这些任务，bucket 规约也可能提前启动。图里表达的是依赖与归属，不能当成所有 GPU 的串行时间线。

继续阅读某个优化分支时，可以在调用处记下四项：操作前后的 shape、使用的 process group、数据是否已经加权、backward 要把哪些贡献相加。再沿参数标记找到 reducer 与 optimizer，才能确认这一轮计算最后更新了哪一份权重。

阅读资料

机制与源码部分主要对照 MoE Parallel Folding（2025，v3） 的双布局、Scalable Training of MoE Models with Megatron Core（2026，v2） 的系统分解，并以固定源码核实具体调用。V4 的模型和训练边界来自原始技术报告及官方配置。

工程讲解另读了 Ultra-Scale Playbook 的 EP、SP/CP 与组合并行章节，以及 HybridEP、融合专家算子。这些文章帮助建立机制图景；具体版本限制和张量变化仍按上面的源码路径检查。