数据截至 2026-08-27(上游 commit e79cb4c1bae1)

06 · GPT 模型组装与端到端主线

这一章讲什么： 收官章，把前五章串起来。GPTModel 怎么把 embedding、TransformerBlock、并行输出层组装成「一个 PP 段」，pretrain_gpt.py 怎么把一个 microbatch 从 token 算到 loss，以及一个完整训练 step 的端到端地图。读完你应该能从头口述一次前向+反向的完整旅程。

1. 它要解决的小问题

前五章讲的都是零件：TP 层、PP 调度、DP 优化器、MoE。但真正跑起来还需要回答：谁把这些零件拼成一个模型？在流水并行的世界里，「模型」从来不是一个完整对象——每张卡只持有模型的一段。

Megatron 的答案很朴素：同一个 GPTModel 类，用两个布尔开关 pre_process / post_process 决定这张卡是首段（带 embedding）、末段（带输出层和 loss）、还是中间段（只有 TransformerBlock 的一段）。

2. 组装：三段式结构

GPTModel.__init__（megatron/core/models/gpt/gpt_model.py:98）按三个开关位组装：

部件	只在哪段出现	本体	行的位置
`self.embedding`	首段（`pre_process`）	`LanguageModelEmbedding`（词表并行 embedding）	`gpt_model.py:176`
`self.decoder`	每段都有	`TransformerBlock`（本段的那几层）	`gpt_model.py:235`
`self.output_layer`	末段（`post_process`）	列切 `ColumnParallelLinear`	`gpt_model.py:279`

三个细节值得停下来看：

输出层是列切的（02 章的刀法用在词表维上），且默认 gather_output=not self.parallel_output（gpt_model.py:290）——parallel_output=True 时 logits 保持按词表切开的状态，不做 all-gather，直接喂给词表并行交叉熵。这是省显存的关键一步：vocab=256k 的完整 logits 是张巨物。
权重共享：share_embeddings_and_output_weights 时输出层复用 embedding 权重（skip_weight_param_allocation 开在 gpt_model.py:291-292），由 setup_embeddings_and_output_layer（megatron/core/models/common/language_module/language_module.py:208）把首段的 embedding 权重接到末段输出层——PP 下这两段在不同卡上，背后有专门的同步组（embedding group，01 章建组时顺带建的）。
弃用警告：构造时即打印「GPTModel IS DEPRECATED」（gpt_model.py:126-130），官方迁移方向是 HybridModel。它是存量主线，不是未来方向。

3. 前向：三段式调用

GPTModel.forward（gpt_model.py:566）的结构干净得可以背下来：

input_ids, position_ids, attention_mask, labels
        │
        ▼
① _preprocess      首段: embedding(input_ids) → decoder_input
                   中间/末段: decoder_input 来自上一段的 P2P 接收
        │
        ▼
② self.decoder(...)   TransformerBlock 跑本段的 N 层
        │              层内部: TP 配对层 + 可选 MoE 四拍
        ▼
③ _postprocess     末段: output_layer → 词表并行 CE loss
                   其他段: 原样返回 hidden, 由调度器 P2P 发给下一段

真实代码里 ② 就是一次 self.decoder(hidden_states=decoder_input, ...) 调用（gpt_model.py:631）。

末段的 loss 在 compute_language_model_loss（language_module.py:161）：labels 转置成 [s, b] 后，默认调 tensor_parallel.vocab_parallel_cross_entropy（language_module.py:200）——logits 全程保持 TP 切片状态，归一化项跨 TP 组 all-reduce 算（机制见 02 章 §4.5）。也有 TE/native 融合实现可选（cross_entropy_loss_fusion）。

4. 原理演示：一个 PP 段的最小形态

把「同构类 + 开关」的心智模型演出来：

# 示意，非源码
class GPTModelStage(nn.Module):
    def __init__(self, config, pre_process, post_process):
        if pre_process:                       # 我是首段
            self.embedding = VocabParallelEmbedding(...)
        self.decoder = TransformerBlock(my_layers(config))  # 每段只拿自己的层
        if post_process:                      # 我是末段
            self.output_layer = ColumnParallelLinear(hidden, vocab, gather_output=False)

    def forward(self, input_ids=None, hidden_from_prev=None, labels=None):
        x = self.embedding(input_ids) if self.pre_process else hidden_from_prev
        x = self.decoder(x)
        if self.post_process:
            logits = self.output_layer(x)     # 词表切片, 不 gather
            return vocab_parallel_ce(logits, labels)
        return x                              # 交给调度器发给下一段

重点看：非首段的输入不是 token 而是上一段发来的激活。调度器侧靠 TransformerBlock.set_input_tensor（megatron/core/transformer/transformer_block.py:463）把 P2P 收到的张量塞进 block——这就是 03 章的 recv_forward 与模型代码之间的接缝。

5. 端到端主线：一个 step 走一遍

把所有章节串成一条线。入口脚本 pretrain_gpt.py 只提供两个钩子，其余交给框架：

model_provider（实际定义在仓库根的 model_provider.py:20，pretrain_gpt.py:80 导入）：按 pre_process/post_process/vp_stage 建出本卡的 GPTModel 段；
forward_step（pretrain_gpt.py:295）：get_batch（:110）取 token/labels → 调 model(tokens, position_ids, attention_mask, labels=..., loss_mask=...)（:361-369）→ 返回 (output_tensor, partial(loss_func, loss_mask))；loss_func（pretrain_gpt.py:222）做 mask 加权求和并打包上报指标。

然后 pretrain(...)（pretrain_gpt.py:540）接管。一个 step 的完整旅程：

pretrain_gpt.py                megatron/training            megatron/core
─────────────                  ────────────────              ─────────────
                                  │ pretrain() (training.py:1500)
                                  │   └ initialize_model_parallel   ← ① 建网格 [01 章]
                                  ▼
model_provider ──► GPTModel 段    │
                                  │ train() 循环:
                                  ▼
                                  │ train_step (training.py:3010)
                                  │   ② zero_grad_buffer       ← grad buffer [04 章]
                                  ▼
forward_step ────► ③ forward_backward_func (schedules.py)
                     PP 调度器逐个 microbatch:                ← [03 章]
                       段内: GPTModel.forward
                         embedding→decoder→output_layer       ← [02 章 TP / 05 章 MoE]
                       段间: P2P 收发激活/梯度
                     反向: 梯度落 grad buffer, 桶满即
                       异步 reduce-scatter                     ← [04 章]
                                  ▼
                                  │ ④ optimizer.step()
                                  │   分片 Adam → 写回 buffer
                                  │   → all-gather 参数        ← [04 章]
                                  ▼
                                  │ 下一个 step

怎么读这张图： 左列是用户要写的（两个钩子），右两列是框架做的。①②④ 各对应前面一章，③ 是调度器与模型代码的交汇点。

6. 关键细节与坑

forward_step 返回的是「输出 + loss 闭包」，不是 loss。 调度器只在末段、且反向需要时才调用这个闭包（partial(loss_func, loss_mask)）——中间段的 output_tensor 只是待发送的激活。写自定义 forward_step 时搞错这个契约是头号 bug 源。
首末段不对称是特性。 embedding 只在首段、loss 只在末段——所以 PP 配置里首末段的显存/算力天然和其他段不同，调 pipeline_model_parallel_layout 做非均匀分段时 first/last stage 通常是调节对象。
parallel_output=False 会很贵。 它强制末段对 logits 做 all-gather（gather_output=not parallel_output），vocab 大时通信和显存都爆——除非你确实需要完整 logits（如某些蒸馏），别关。
词表并行 CE 有三条实现路径。 默认 tensor_parallel.vocab_parallel_cross_entropy（language_module.py:200），或 cross_entropy_loss_fusion 下的 TE / native 融合版（:172-199）——融合版省一次 logits 物化，但各有各的版本约束（代码里对 TE 版本有硬断言）。
packed sequence 走 PackedSeqParams。 变长打包序列（thd 格式）在 forward_step 里组装（pretrain_gpt.py:335-347），CP 组信息也在这里挂上（cp_group=hybrid_cp_group）——长序列+CP 的正确性依赖这条路径，而不是 attention_mask。
NaN/Inf 检查在 DP 规约之前。 loss_func 先在本卡校验 NaN/Inf 再上报（pretrain_gpt.py:258-279），配合 rerun state machine 可以自动重跑出错批次——大规模训练里「坏 batch 毒化全局 loss」的真实防线。

7. 代码地图

主题	文件路径	符号名
GPT 模型组装	`megatron/core/models/gpt/gpt_model.py`	`GPTModel`（`__init__`/`forward`/`_postprocess`）
Transformer 块	`megatron/core/transformer/transformer_block.py`	`TransformerBlock`（`set_input_tensor`/`forward`）
词表并行 loss	`megatron/core/models/common/language_module/language_module.py`	`compute_language_model_loss`
embedding/输出层权重共享	`megatron/core/models/common/language_module/language_module.py`	`setup_embeddings_and_output_layer`
GPT 预训练脚本	`pretrain_gpt.py`	`forward_step`、`loss_func`、`get_batch`
模型建造钩子	`model_provider.py`	`model_provider`
训练主循环	`megatron/training/training.py`	`pretrain`、`train`、`train_step`
库用法最小示例	`examples/run_simple_mcore_train_loop.py`	`initialize_distributed`、`model_provider`

回到全局：index.md；横向对比兄弟实现：DeepSpeed、verl。

1. 它要解决的小问题​

2. 组装：三段式结构​

3. 前向：三段式调用​

4. 原理演示：一个 PP 段的最小形态​

5. 端到端主线：一个 step 走一遍​

6. 关键细节与坑​

7. 代码地图​