Megatron Core 代码地图与 RL 训推

v7 的概念齐了，还剩两个缺口。面试官会追问“这在 Megatron 里是怎么写的”，团队现在又在做 RL post-training。这一课把每一版对到 megatron/core/ 的文件上，带读一层 Transformer，再讲 RL 里同一份 weight 怎么在训练和生成两种切法之间来回搬。

面试会这样问How does Megatron Core set up its process groups? · Walk me through a transformer layer's forward pass under tensor parallelism. · In RL post-training, how do the trainer's weights get to the inference engine?

考过NVIDIA 框架岗 2023 电话面（"Column Parallel和Row Parallel Linear Layer的实现和pros & cons"，写少量 pseudocode，知乎）· NVIDIA DL Engineer 慕尼黑 2026（"Can you explain the details of how distributed training works"，Glassdoor RVW104512401）· RL 训推部分还没有面经，依据是团队的招聘帖（见最后一节）

一个仓库，两样东西

parallel_state：给每张卡发一组坐标

Process group 是一组要一起做 collective 的 rank，像微信群，all-reduce 只在群里发。每张卡同时在一个 TP 群、一个 DP 群和一个 PP 群里。

例子：16 张卡（2 台机器 × 8），TP = 2、PP = 2，DP 由代码算出来，16 ÷ (2 × 2) = 4。initialize_model_parallel 的默认 order="tp-cp-ep-dp-pp" 越靠左变化越快，CP、EP 都是 1 时就是按 tp、dp、pp 的顺序编号：

\[ \text{rank} = \text{tp} + 2\cdot\text{dp} + 8\cdot\text{pp} \]

例：rank 5 = 1 + 2 × 2 + 8 × 0，坐标是 (tp 1, dp 2, pp 0)，它在下表的 [4, 5]、[1, 3, 5, 7]、[5, 13] 三个群里。

群个数 × 每群卡数成员在哪
TP8 × 2[0, 1] [2, 3] … [14, 15]相邻 rank，同一台机器（第 8 课的规则）
DP4 × 4[0, 2, 4, 6] [1, 3, 5, 7] [8, 10, 12, 14] [9, 11, 13, 15]这个例子里也在机内
PP8 × 2[0, 8] [1, 9] … [7, 15]跨机器

算成员的是 RankGenerator 类，其他代码用 get_tensor_model_parallel_group() 这类函数拿到群，再在群里调 NCCL。

文件地图

都在 megatron/core/ 下：

对应文件先看这几个名字
process groupparallel_state.pyinitialize_model_parallel、RankGenerator
第 5 课 TPtensor_parallel/layers.pyColumnParallelLinear、RowParallelLinear、VocabParallelEmbedding
第 5 课 f / gtensor_parallel/mappings.pycopy_to_tensor_model_parallel_region（f）、reduce_from_tensor_model_parallel_region（g）
vocab-parallel losstensor_parallel/cross_entropy.pyvocab_parallel_cross_entropy
一层 Transformertransformer/transformer_layer.py、attention.py、mlp.pyTransformerLayer、SelfAttention、MLP
第 6 课 PPpipeline_parallel/schedules.pyget_forward_backward_func，1F1B 和 interleaved 的 forward_backward_pipelining_*
第 4 课 DDPdistributed/distributed_data_parallel.pyDistributedDataParallel（bucket、overlap）
第 4 课 ZeRO-1optimizer/distrib_optimizer.pyDistributedOptimizer
第 7 课 MoEtransformer/moe/router.py、token_dispatcher.pyTopKRouter、MoEAlltoAllTokenDispatcher
checkpointdist_checkpointing/加载时可以换并行配置
本课 RLresharding/swap_model_weights

怎么读：一个文件 15 分钟

  1. 读每个 class 的 docstring，说出它解决第几课的什么问题
  2. 看 forward 调了哪个通信函数，对上前面数过的通信次数
  3. fusion、FP8、CUDA graph 的分支先跳过

先读最核心的四个，约一小时：parallel_state.py → layers.py → mappings.py → attention.py 和 mlp.py。注意默认的 Transformer Engine spec 把 LayerNorm 融进了 linear_qkv（TELayerNormColumnParallelLinear），pre_mlp_layernorm 变成 IdentityOp，第一次读看 models/gpt/gpt_layer_specs.py 里的 get_gpt_layer_local_spec 更清楚。

带读：一层 TransformerLayer 的 forward

用 Qwen3-8B 的维度、TP = 2。Megatron 的 activation 按 [s, b, h] 排，sequence 在前。

  1. input_layernorm：每卡对完整的 [s, b, 4096] 算一遍（没开 SP）
  2. self_attention（SelfAttention）
    • linear_qkv：ColumnParallelLinear，gather_output=False。输出宽度 = Q 4096 + K 1024 + V 1024 = 6144，每卡 3072（16 个 Q head、4 个 KV head）。f 在它内部，backward 时 all-reduce input 的 gradient（参数名 allreduce_dgrad）
    • core_attention：DotProductAttention，用 Transformer Engine 时是 TEDotProductAttention。每卡只算自己的 head，不通信
    • linear_proj：RowParallelLinear，input_is_parallel=True（输入已经切好）。每卡用 [2048, 4096] 的 weight 算出部分和，g 做 all-reduce（forward 第 1 次），开了 SP 换成 reduce-scatter
  3. self_attn_bda：bias-dropout-add，最后加 residual
  4. pre_mlp_layernorm
  5. mlp（MLP）
    • linear_fc1：ColumnParallelLinear。SwiGLU 的 gate 和 up 合成一个矩阵，宽度 2 × 12288 = 24576，每卡 12288
    • activation_func：SwiGLU，在本卡算
    • linear_fc2：RowParallelLinear，每卡 [6144, 4096]，g 做 all-reduce（forward 第 2 次）
  6. mlp_bda：再加一次 residual

数一下：forward 2 次 all-reduce，backward 2 次（linear_qkv、linear_fc1 的 input gradient），和第 5 课一致。每层约 193M 个 parameter（attention 42M + MLP 151M），TP = 2 时每卡约 96M。

RL post-training：一个训练引擎，一个 inference engine

GRPO、PPO 的每一轮先 rollout，用当前 policy 生成一批回答，主要是 decode，memory-bound（第 0 课）。再 training，算 reward 和 advantage 后做 gradient update，就是第 1–8 课的内容。两件事各用最合适的引擎：

为什么每轮都要重切 weight：update 之后 inference engine 手里还是旧 weight，不同步的话 rollout 用的就是旧 policy。两边切法又不同。例如 8 张卡，训练用 TP4 × PP2（每卡一半的层、每层 1/4 的矩阵），vLLM 用 TP8（每卡全部的层、每层 1/8 的矩阵），所以每个 parameter 都要先拼起来再重切。存储格式也不同，Megatron 的 linear_qkv 按 query group 交错存（q1 q2 k1 v1 | q3 q4 k2 v2 | …），HF 和 vLLM 的 Q、K、V 是分开的，还要重排、改名。NeMo RL 把这一步叫 refit，Megatron Core 有专门的 resharding/。

要搬多少：8B 的 bf16 weight 是 16 GB，8 张卡每卡 2 GB，走 NVLink 约 4.4 ms，走 IB 约 40 ms。70B 就是 140 GB，而且每轮都要搬。

colocated（同一批卡）separate（分开的卡）
怎么跑训练和 rollout 轮流用同一批卡，训练时 inference 端释放 KV cache一批卡只训练，另一批只生成
weight 怎么传不出机器，NeMo RL 用 CUDA IPC 直接共享 GPU memory跨网络，NeMo RL 默认用 NCCL broadcast
取舍没有卡闲着，传得快，但同一时刻只有一边在干活两边能同时跑、各选并行配置，但要用落后几步的 policy（off-policy），否则一边空转
例子verl 的 hybrid engine（hybrid_engine: true，默认）；NeMo RL colocated.enabled: trueNeMo RL colocated.enabled: false

去查一下OpenManus-RL 基于 verl。打开你的启动脚本或 config：actor_rollout_ref.actor.strategy 是 fsdp 还是 megatron？actor_rollout_ref.rollout.name 是 vllm 还是 sglang，rollout 的 tensor_model_parallel_size 是几？hybrid_engine 是不是 true？查清楚就能讲出你那套训推每轮怎么同步 weight，这正是团队在做的事。NVIDIA 2025-10-31 的 “Deep Learning Framework Interns” LinkedIn 招聘帖把 Megatron-Core、Megatron-Bridge、NeMo-RL 列在一起，主题之一是用 NeMo-RL 做 RL post-training。2026 年夏天一位 NeMo 和 Megatron 组实习生的简历写的是给 NeMo-RL 加 delta-compressed weight 和 GPUDirect RDMA 的 refit。

30 秒英文回答

Megatron Core is NVIDIA's library of composable parallel building blocks; Megatron-LM, Megatron Bridge and NeMo RL build on it. parallel_state gives every rank TP, DP and PP coordinates. TP groups are adjacent ranks on NVLink, and PP groups cross nodes. In a layer, linear_qkv and linear_fc1 are column-parallel, linear_proj and linear_fc2 are row-parallel: two all-reduces forward, two backward. In RL, a trainer like Megatron or FSDP alternates with an inference engine like vLLM. Their layouts differ, say TP4 × PP2 versus TP8, so after every update the weights are resharded and synced.

自测

parallel_state.py 做什么？

按各维的大小和 order 给每个 rank 算坐标，建好各个 process group。默认 tp 变化最快，所以 TP 群是相邻 rank，落在机内。

16 卡，TP = 2、PP = 2、DP = 4，rank 13 在哪几个群？

13 = 1 + 2 × 2 + 8 × 1，坐标 (tp 1, dp 2, pp 1)。TP 群 [12, 13]，DP 群 [9, 11, 13, 15]，PP 群 [5, 13]。

ColumnParallelLinear 在哪个文件？一层里哪些 linear 用它？

tensor_parallel/layers.py。linear_qkv 和 linear_fc1 用它，linear_proj 和 linear_fc2 用 RowParallelLinear。f、g 在 mappings.py。

RL 为什么要 reshard weight？

update 后 inference 端的 weight 过期，而两端的切法和格式都不同（TP4 × PP2 vs TP8，QKV 交错存 vs 分开存），要拼起来重切再发过去。

Megatron Core 给了什么，是只用 PyTorch FSDP 得不到的？

FSDP 只在 DP 方向切 model state（约等于 ZeRO-3），activation 不切。Megatron Core 另外提供 TP + SP、PP schedule、CP、EP、distributed optimizer、FP8 kernel，以及能换并行配置加载的 distributed checkpoint，70B 以上、长序列、MoE 时缺的正是这些。

想看细一点 6.1 TP Linear · 6.2 TP communication · 6.5 Pipeline Parallel · 6.6 DP、ZeRO、FSDP · 6.8 MoE 与 Expert Parallel · 6.9 怎么选 parallelism 配置 · 3.15 RLHF、DPO、GRPO