20-Megatron-MCore-TP-PP-EP如何支撑MoE

Megatron/MCore 切开的不是某一个 tensor，而是模型、专家、序列和流水线同时存在的并行坐标。

第 19 篇讲 FSDP/FSDP2：当一次 RL update 进入 worker 内部，训练状态怎样被参数、梯度和 optimizer state 的 shard 承担。本文换到 Megatron/MCore 后端，问题也随之变化：如果模型不只是“大”，还带着 MoE 专家、流水线阶段、长上下文和 tensor parallel，verl 怎样把这一组并行组合接到同一个 BaseEngine合同里？

本文的核心判断是：Megatron/MCore 在 verl 里不是“比 FSDP 更复杂的训练策略”，而是一个并行坐标系统。TrainingWorker仍然只通过 EngineRegistry创建 engine，上层 PPO/GRPO 仍然只看 train_batch()和 infer_batch()；但进入 MegatronEngine以后，TP、PP、VPP、EP、ETP、CP 和 SP 会先被写入 Megatron parallel state，再通过 MBridge/provider 参与模型构造，最后由 Megatron 的 pipeline schedule 执行 micro-batch。这个坐标系统让 MoE 大模型可以训练，也把路由一致性、all-to-all、micro-batch 可整除和权重导出复杂度留在 engine 内部。

先看第 20 篇在第四组里的位置。读图时注意：FSDP 文章关注训练状态怎样 shard；本文关注同一个 worker 边界下，Megatron 如何把模型计算本身放进多个并行维度。

Megatron/MCore 在 RL update 中的位置

这张图对应的源码入口和第 19 篇一致：TrainingWorker通过 EngineRegistry.new(model_type=..., backend=self.engine_config.strategy, ...)创建后端；language model 和 value model 分别注册到 backend="megatron"；BaseEngine.train_batch()仍然定义清梯度、forward/backward、optimizer step 的统一顺序（verl/workers/engine_workers.py:131-139，verl/workers/engine/base.py:112-131，verl/workers/engine/base.py:323-337，verl/workers/engine/megatron/transformer_impl.py:789-790，1004-1005）。

1. 先把并行维度变成配置合同

理解 Megatron/MCore 的第一步，不是先背 TP、PP、EP 的定义，而是看这些维度在 verl 里从哪里进入系统。它们不是 trainer 里的临时变量，而是 McoreEngineConfig的字段；actor/critic 只需要把 strategy设成 megatron，再继承 trainer/config/engine/megatron.yaml。

下面这张图把并行维度拆开。看图时注意：这些维度解决的问题不同，不能互相替代；MoE 主要把 EP/ETP 和 router/all-to-all 带进来，但仍然要和 TP、PP、CP/SP 共存。

Megatron 并行维度地图

McoreEngineConfig直接暴露 tensor_model_parallel_size、expert_model_parallel_size、expert_tensor_parallel_size、pipeline_model_parallel_size、virtual_pipeline_model_parallel_size、context_parallel_size和 sequence_parallel等字段；配置文件也按 TP、EP、PP、CP、SP、distributed optimizer、distributed checkpointing、router replay 分组呈现这些开关（verl/workers/config/engine.py:149-218，verl/trainer/config/engine/megatron.yaml:13-57，108-125）。还有一个细节很重要：当 TP size 为 1 时，__post_init__()会把 sequence_parallel改成 False，因为 SP 在这条路径里依附于 tensor parallel 的切分坐标（verl/workers/config/engine.py:209-216）。

这些字段最终进入 MegatronEngine._init_device_mesh()。verl 调用 mpu.initialize_model_parallel()时，把 TP、PP、VPP、CP、EP、ETP 一次性传进去；如果启用 dynamic context parallel，还会先检查 Megatron 版本是否支持相应参数，并要求配置 max_seqlen_per_dp_cp_rank（verl/workers/engine/megatron/transformer_impl.py:131-158）。所以 Megatron 的第一层工程含义是“建坐标系”：后续模型构造、pipeline schedule、MoE 路由和输出收集都在这个坐标系里解释。

2. TP、PP、EP 切的是三类不同压力

TP、PP、EP 经常被放在一起说，但它们处理的是三种不同压力。

TP 把单层内部的矩阵计算、attention/logits 等 tensor 维度分摊到多个 rank。Megatron forward 路径里，logprob 和 entropy 会走 vocab_parallel_log_probs_from_logits()、vocab_parallel_entropy()这类 vocab-parallel 操作；这意味着输出 tensor 的语义也已经带着 TP 切分（verl/workers/engine/megatron/transformer_impl.py:905-934）。

PP 把层切成 pipeline stage，VPP 再把每个 pipeline rank 上的层切成虚拟 chunk。verl 的输出收集也体现了这一点：is_mp_src_rank_with_outputs()只让 TP rank 0、最后一个 PP rank、CP rank 0 作为可收集输出的 rank（verl/workers/engine/megatron/transformer_impl.py:363-368）。如果不区分 PP stage，就会误以为每个 rank 都能直接返回完整模型输出。

EP 则来自 MoE：token 会被 router 分配给不同专家，专家可以分布在 expert parallel group 上；ETP 进一步切专家内部 tensor。verl 在初始化 Megatron parallel state 时传入 EP/ETP，MBridge provider 路径也把 expert_model_parallel_size和 expert_tensor_parallel_size写进去，并把 MoE token dispatcher 设成 alltoall（verl/workers/engine/megatron/transformer_impl.py:148-155，207-222）。因此 EP 不是“更多 data parallel”，而是 token 到专家之间的通信和路由问题。

3. MBridge 把 HF 模型变成 MCore 分布式模型

有了并行坐标，还要把模型按这个坐标建出来。verl 这里的关键桥梁是 MBridge/provider：它把 HF 配置和本地模型路径转成 Megatron/MCore 可以构造的 distributed model，同时把 TP、PP、EP、CP、SP 等并行字段注入 provider。

下面这张图按源码顺序展开。看图时注意：并行配置不是训练开始后才生效，而是在 provider 和 module 构造时已经决定了模型的形态。

MBridge provider 如何构造 Megatron 模型

在非 vanilla MBridge 路径里，MegatronEngine._build_tf_config()调 AutoBridge.from_hf_pretrained()，再通过 bridge.to_megatron_provider(load_weights=False)取得 provider；随后设置 dtype、TP、PP、EP、ETP、VPP、CP、SP，以及 attention_backend、variable_seq_lengths、MoE dispatcher 等字段，最后 provider.finalize()（verl/workers/engine/megatron/transformer_impl.py:190-238）。如果启用 router replay，engine 会先打 patch，并把 enable_routing_replay写入 provider 或 transformer config（verl/workers/engine/megatron/transformer_impl.py:117-123，234-250）。

模型本体在 _build_megatron_module()里创建。verl 先构造 McoreModuleWrapperConfig，决定是否 value model、是否 DDP、是否 distributed optimizer、是否 Megatron-FSDP；然后调用 make_megatron_module()，加载 HF 权重或 distributed checkpoint（verl/workers/engine/megatron/transformer_impl.py:263-317）。在 provider 分支里，make_megatron_module()会先注册 PEFT、value head、freeze router 等 pre-wrap hook，再创建 DDP config，最后调用 provider.provide_distributed_model(...)（verl/utils/megatron_utils.py:215-331）。这解释了为什么 Megatron/MCore 后端不只是“换一个 optimizer”：模型结构、包装方式、权重加载和并行 group 在初始化阶段就绑定了。

4. 训练执行交给 Megatron pipeline schedule

从 BaseEngine看，训练仍然是 optimizer_zero_grad()、forward_backward_batch()、optimizer_step()。进入 Megatron 后，真正的执行单位会变成 micro-batch，并交给 Megatron 的 pipeline schedule；这和 FSDP 里逐个 micro-batch 手动 loss.backward()的直觉不同。

下面这张图展示 micro-batch 如何进入 PP/VPP 调度。看图时注意三点：第一，batch 会先按 dynamic batch 或固定 micro-batch 切分；第二，VPP 会要求 micro-batch 数满足分组整除；第三，只有 pipeline last stage 负责把结果后处理回上层。

Megatron pipeline micro-batch 执行路径

源码里，forward_backward_batch()先把 loss_mask的 token 数在 data parallel group 上 all-reduce，用于全局 loss 归一化；然后把 sp_size设成 context_parallel_size，调用 prepare_micro_batches()切 batch（verl/workers/engine/megatron/transformer_impl.py:599-622）。prepare_micro_batches()在 dynamic batch 路径下按 max_token_len_per_gpu * sp_size重排样本；固定路径下按 micro_batch_size_per_gpu切 TensorDict（verl/workers/engine/utils.py:58-96）。

如果启用 VPP，Megatron engine 会读取 microbatch_group_size_per_vp_stage，并断言 micro-batch 数可整除；随后每个 micro-batch 都写入 num_micro_batch，再通过 get_forward_backward_func()取得 Megatron schedule，使用 make_batch_generator(micro_batches, vpp_size=len(self.module))为虚拟 pipeline chunk 准备迭代器（verl/workers/engine/megatron/transformer_impl.py:610-674，verl/utils/megatron/pipeline_parallel.py:49-70）。forward step 里还会处理 dynamic CP split、remove padding 数据格式、vocab-parallel logprob/entropy 和 router replay；postprocess 时，loss 会乘回 num_micro_batch，因为 Megatron pipeline schedule 内部会按 micro-batch 做缩放（verl/workers/engine/megatron/transformer_impl.py:808-1001）。

这部分的工程含义是：Megatron/MCore 后端把“怎么跑完一次 backward”交给 pipeline schedule，而不是让上层 trainer 感知 PP stage、VPP chunk 或 P2P 通信形状。代价是 micro-batch 数、动态重排、输出收集 rank 和 loss 缩放都必须在 engine 内处理正确。

5. MoE 还需要路由一致性和权重导出边界

MoE 的难点不只是专家多。一次 RL update 里，forward、logprob、loss 和 backward 需要围绕同一批 token 的路由结果保持一致；同时，训练完的 Megatron/MCore 权重还要导出给 rollout 或 checkpoint engine。verl 因此把 router replay 和 per-tensor 权重导出都放在 Megatron engine 的职责里。

下面这张图把 MoE 路由和导出边界放在一起。看图时注意：EP/all-to-all 解决 token 到专家的分发，router replay 解决路由结果在不同阶段复用，get_per_tensor_param()则把 Megatron 内部权重重新变成外部可同步的 tensor 流。

MoE 路由与 Megatron 权重导出边界

MCore 配置转换里，MoE 模型会设置专家数量、top-k、router score function、router load balancing、grouped GEMM 等字段；基础配置也默认把 moe_token_dispatcher_type设成 alltoall（verl/models/mcore/config_converter.py:60-98，187-218，221-245）。在 Megatron engine 里，router replay 配置不为 disabled时会应用 patch；patched router 可以记录 top-k indices，也可以在 replay forward/backward 中复用给定 indices（verl/utils/megatron/router_replay_patch.py:38-130，141-199，337-443）。forward_backward_batch()会根据 enable_routing_replay设置全局 replay/record action，forward_step()则在对应 VP stage 上写入 replay 数据、切换到 backward replay，并在需要时收集 routed experts（verl/workers/engine/megatron/transformer_impl.py:652-713，849-959）。

最后，MegatronEngine.get_per_tensor_param()会把 module 加载回 GPU，再通过 bridge 导出 HF 权重、adapter 权重或 vanilla bridge 权重；如果启用 QAT，还会经过 QAT weight exporter（verl/workers/engine/megatron/transformer_impl.py:715-744）。这条路径和第 19 篇 FSDP 的导出边界相呼应：训练 engine 可以在内部使用不同分布式形态，但对 rollout/weight sync 层来说，最终仍然需要一条逐 tensor 的同步接口。

小结：Megatron/MCore 是 worker 内部的并行坐标合同

本文从 TrainingWorker进入 MegatronEngine，看到 Megatron/MCore 后端在 verl 里的职责链：

actor/critic strategy = megatron
-> EngineRegistry selects MegatronEngine
-> McoreEngineConfig carries TP / PP / VPP / EP / ETP / CP / SP
-> mpu.initialize_model_parallel()
-> MBridge provider builds distributed MCore module
-> Megatron pipeline schedule runs micro-batches
-> MoE router replay and all-to-all expert dispatch
-> get_per_tensor_param() exports weights for rollout sync

放回系列地图，第 20 篇补上的是 workers/resources -> training engine/distributed core里的第二块：当 FSDP 的状态切分不再足够表达 MoE 大模型训练时，Megatron/MCore 如何把模型计算、专家路由和 pipeline 调度放进同一个 engine 合同。下一篇会回到抽象层，专门解释 EngineRegistry为什么能让上层 trainer 不关心 FSDP、Megatron 或其他后端的具体实现，却仍然要求每个后端完整实现训练、推理、rank、checkpoint 和权重导出边界。

本文源码索引

• verl/workers/engine_workers.py:131-139：TrainingWorker如何用 EngineRegistry.new()根据 strategy创建后端。
• verl/workers/engine/base.py:112-131、323-337：BaseEngine.train_batch()的统一训练顺序，以及 registry 如何实例化 engine。
• verl/trainer/config/actor/megatron_actor.yaml:1-18、verl/trainer/config/critic/megatron_critic.yaml:1-28：actor/critic 如何继承 Megatron optimizer 和 engine 配置。
• verl/workers/config/engine.py:149-218、verl/trainer/config/engine/megatron.yaml:13-57、108-125：Megatron/MCore 并行配置、distributed optimizer/checkpointing 和 router replay 入口。
• verl/workers/engine/megatron/transformer_impl.py:75-158：MegatronEngine初始化、offload 标记、router replay 开关和 Megatron model parallel 初始化。
• verl/workers/engine/megatron/transformer_impl.py:190-238：MBridge/provider 如何注入 dtype、TP、PP、EP、ETP、VPP、CP、SP 和 MoE dispatcher。
• verl/workers/engine/megatron/transformer_impl.py:263-345、verl/utils/megatron_utils.py:215-331：MCore module、DDP/distributed optimizer、权重加载和 provider distributed model 构造。
• verl/workers/engine/megatron/transformer_impl.py:599-713、verl/workers/engine/utils.py:58-96、verl/utils/megatron/pipeline_parallel.py:49-70：Megatron micro-batch 切分、VPP batch generator 和 pipeline schedule 调用。
• verl/workers/engine/megatron/transformer_impl.py:808-1001：LM forward step、vocab-parallel logprob/entropy、router replay 和 loss postprocess。
• verl/models/mcore/config_converter.py:60-98、187-245：MCore transformer 配置如何写入并行字段与 MoE 字段。
• verl/utils/megatron/router_replay_patch.py:38-130、141-199、337-443，verl/utils/megatron/router_replay_utils.py:276-335、486-555：router replay 的记录、回放、patch 和本地 layer/router 定位。
• verl/workers/engine/megatron/transformer_impl.py:363-368、715-744：哪些 model-parallel rank 负责返回输出，以及 Megatron 权重如何逐 tensor 导出。