Megatron 面试题（一）：模型抽象与训练数据流

本文是Megatron 专家面试题系列的 Level 1，覆盖 M013-M024。重点是Transformer layer、QKV、GQA、RoPE、SwiGLU、vocab parallel 与精度状态。

上一篇：全局架构 · 系列总索引 · 下一篇：TP 与 SP

---

3. Level 1：模型抽象与训练数据流（M013-M024）

M013. 一个 TransformerLayer 的主数据流是什么？

高分回答要点

• 典型 pre-norm 层是 norm -> self-attention -> residual，再 norm -> MLP -> residual；具体由 spec/config 决定。
• TP/SP/CP 会改变每个边界上的 activation layout，不能只画单卡算子图。
• backward 反向经过同一依赖，并在 parallel linear、attention/context 和 DP buffer 处触发不同通信。

追问：让你标出每个 residual 前后的 shape 和是否在 TP rank 间复制。

前置知识与分析过程：先确定张量布局记号，例如[s,b,h]及SP时[s/t,b,h]，再沿norm、attention、residual、MLP逐边标layout。每经过parallel linear都检查是否发生collective或保持分片。

主问题得分点（10分）：基本layer顺序 2 分；输入输出shape 2 分；TP/SP/CP layout变化 3 分；backward和通信位置 2 分；spec可变性 1 分。

追问分析与参考回答：无SP时residual主干通常在TP ranks复制为[s,b,h]；QKV column-parallel后内部head分片，output row-parallel归约回复制hidden再加residual。开SP时norm/residual可保持sequence shard[s/t,b,h]，在需要column-parallel输入前AG、row-parallel输出后RS。具体边界依版本实现核对trace。

追问得分点（5分）：给shape 1 分；无SP复制语义 1 分；SP shard 1 分；AG/RS位置 1 分；说明实现边界 1 分。

M014. Megatron attention 中 Q/K/V 如何随 TP 分片？

高分回答要点

• 通常按 attention heads 或投影输出维分到 TP ranks，每 rank 计算本地 heads 的 Q/K/V 和 attention。
• output projection 用 row-parallel 汇合各 rank head contribution，forward 需要归约或 SP 下的 ReduceScatter。
• head 数、query groups 与 TP 必须满足 layout 约束；GQA/MQA 可能涉及 KV group 复制或不同分片策略。

追问：给 hidden=8192, heads=64, TP=8，每 rank query heads=8；再把 KV groups 改成 8，讨论每 rank KV ownership。

前置知识与分析过程：先计算head_dim=hidden/heads，再按Q heads与KV groups分别判断可整除性和owner。最后检查QKV projection输出shape与attention kernel支持。

主问题得分点（10分）：QKV/head分片定义 3 分；output projection通信 2 分；GQA/MQA约束 2 分；给具体shape/owner 2 分；性能含义 1 分。

追问分析与参考回答：head_dim=128，MHA时每rank 8个Q/K/V heads。GQA设8个KV groups且TP8时，理想布局为每rank 1个KV group服务本rank 8个Q heads中的对应组；本地Q维8x128，KV维1x128每种。若groups小于TP，就可能需要复制或更复杂映射，不能负数/零head分片。

追问得分点（5分）：head_dim 1 分；Q heads 1 分；KV owner 2 分；groups<TP反例 1 分。

M015. MHA、GQA、MQA 对 Megatron 并行与显存有什么影响？

高分回答要点

• GQA/MQA 减少 KV heads，训练 attention 参数和推理 KV cache 更小，但 Q 与 KV 的 TP 可整除/复制关系更复杂。
• 若 KV groups 少于 TP，某些 rank 可能复制 KV 或使用专门 partition，不能直接假设所有权重都按 TP 等分。
• 性能需看 projection GEMM、attention kernel 支持、通信和负载均衡，不只看参数数目。

追问：为什么 TP degree 可能受 num_query_groups 而不只是 num_attention_heads 限制？

前置知识与分析过程：先把参数/activation减少与layout合法性分开。列Q heads和KV groups两个独立分片轴，再分析复制方案和性能代价。

主问题得分点（10分）：MHA/GQA/MQA定义 3 分；参数/KV收益 2 分；TP分片约束 3 分；kernel/性能边界 2 分。

追问分析与参考回答：Q可按attention heads分，但K/V只有num_query_groups组；若TP大于或不能整除groups，简单“每rank若干KV组”失效。框架需复制KV组、改变partition或拒绝配置，都会影响参数节省和通信，因此config validation要同时检查heads与query groups。

追问得分点（5分）：指出独立约束 2 分；给不可整除场景 1 分；列复制/拒绝方案 1 分；性能后果 1 分。

M016. RoPE、learned absolute position 和长上下文 scaling 在架构上如何接入？

高分回答要点

• learned absolute position 通常在 embedding 相加；RoPE 在 attention 内对 Q/K 施加位置旋转，不增加 position embedding table。
• PP 首 stage 提供 position ids，但实际 RoPE 计算在持有 attention layer 的各 stage；CP 还需正确处理全局 position offset。
• scaling 改变模型语义和 checkpoint 配置，不是纯 runtime 优化。

追问：为什么 CP rank 只持局部 token 仍必须知道全局 position？

前置知识与分析过程：先区分位置编码注入点，再分析PP传position、CP切sequence后的global offset。最后说明checkpoint/config语义。

主问题得分点（10分）：三类位置编码区别 3 分；RoPE接入attention 2 分；PP/CP position流 3 分；scaling语义边界 2 分。

追问分析与参考回答：RoPE相位由token在完整序列中的绝对/规则位置决定；CP rank虽只计算局部chunk，但若把本地索引从0重新编号，跨chunkQ/K相位关系错误。应由全局sequence offset生成position或传递正确position ids，packing时还需按样本边界重置/映射。

追问得分点（5分）：RoPE相位语义 2 分；global offset 1 分；错误后果 1 分；packing边界 1 分。

M017. SwiGLU MLP 的参数与 TP 切分有什么特殊点？

高分回答要点

• 第一投影产生 gate 和 value 两路，局部激活后相乘，再经第二投影回 hidden；中间维与参数量口径需考虑两路。
• gate/up 常用 column parallel 在输出维切，down 用 row parallel 在输入维切，激活可保持本地。
• fused SwiGLU/GroupedGEMM 的 shape、对齐和 checkpoint layout 都与普通 GELU MLP 不同。

追问：为什么从 GELU 4h 切到 SwiGLU 常调整 FFN hidden，而不是仍机械使用 4h？

前置知识与分析过程：先写GELU与SwiGLU两种FFN参数公式，再按TP轴推local shape。比较参数/FLOPs对齐后才能讨论中间维选择。

主问题得分点（10分）：SwiGLU数据流 3 分；两路参数口径 2 分；TP切分 2 分；fusion/shape 2 分；中间维调整原因 1 分。

追问分析与参考回答：GELU FFN约有h*4h + 4h*h=8h^2权重；SwiGLU第一投影有gate/up两路，若每路仍4h，总量约12h^2，显著增参增FLOPs。为匹配dense FFN预算，常把每路中间维降至约8h/3并做硬件对齐，具体值依模型recipe。

追问得分点（5分）：写两种参数公式 2 分；得出需降中间维 1 分；给约8h/3及非固定边界 1 分；对齐/fusion 1 分。

M018. Vocab-parallel embedding 如何工作？

高分回答要点

• vocabulary rows 在 TP ranks 分片；每 rank 对不属于本 shard 的 token mask，查本地表后在 TP group 归并结果。
• 参数显存按 TP 降低，但 forward 有通信；词表 padding/整除会影响本地 shard 大小。
• tied output weight 必须保持输入 embedding 与输出 projection 的 shard/同步语义一致。

追问：输入 token 只属于一个 shard，为什么最终每个 rank 的 hidden activation 都能得到正确 embedding？

前置知识与分析过程：先写词表切分区间和token mask，再对单token追踪各rank局部lookup值。最后判断使用reduce还是gather及输出layout。

主问题得分点（10分）：vocab row分片 2 分；mask/local lookup 2 分；跨rank归并 2 分；padding/整除 2 分；tied权重 2 分。

追问分析与参考回答：只有owner rank查到真实向量，其他rank将该token局部结果置零；对TP group做sum-reduce后，每rank都得到owner贡献，即完整embedding hidden。若后续SP需要sequence shard，归并形式可能与layout组合优化，但逻辑oracle仍是每token唯一owner贡献。

追问得分点（5分）：owner/mask 2 分；sum语义 2 分；layout优化边界 1 分。

M019. Vocab-parallel cross entropy 如何避免 gather 全量 logits？

高分回答要点

• 每 rank 只持 vocab shard logits；先跨 rank 求 global max 保数值稳定，再求 global exp sum。
• target logit 只在所属 shard 提取，跨 rank归并后计算 loss；backward 直接生成本地 logits gradient。
• 通信规模主要与 token 数相关，而不是完整 tokens x vocab AllGather，显著降低峰值。

追问：手推 log-sum-exp 的 global max/sum 两阶段，以及 target 不在本 shard 时的 mask。

前置知识与分析过程：从单token交叉熵-z_y+log sum exp(z)出发，把vocab维拆为shards。为数值稳定先求全局max，再求全局exp和与target logit。

主问题得分点（10分）：避免full logits gather 2 分；global max 2 分；global exp sum 2 分；target mask/reduce 2 分；backward本地gradient 2 分。

追问分析与参考回答：各rank算local max，AllReduce(MAX)得m；算sum_j exp(z_j-m)的local sum并AllReduce(SUM)得S。target owner取z_y，非owner置0，再SUM得到全局target logit；loss=log(S)+m-z_y。backward每rank对本地softmax减本地target one-hot。

追问得分点（5分）：max阶段 1 分；sum阶段 1 分；target owner 1 分；loss公式 1 分；backward 1 分。

M020. tied embedding 在 PP 首尾 stage 之间如何保持一致？

高分回答要点

• 首 stage 输入 embedding 与末 stage output weight 逻辑共享，但物理上可能位于不同 PP ranks。
• 初始化/加载后需要 embedding group 同步参数，backward 后还需归并两端 gradient。
• checkpoint metadata 必须表达共享关系；只让两个 stage 各自 optimizer step 会发生漂移。

追问：这也是 finalize_model_grads() 不能只做普通 DP gradient sync 的原因之一。

前置知识与分析过程：先区分逻辑共享与物理副本，再列参数初始化、forward读取、backward gradient、optimizer四个同步点。判断普通DP group是否包含PP首尾。

主问题得分点（10分）：PP首尾共享问题 3 分；参数/grad同步 3 分；checkpoint表达 2 分；optimizer漂移风险 2 分。

追问分析与参考回答：普通DP只连接相同model-parallel坐标的replicas，不会连接PP first/last。tied embedding在首尾物理各有一份，需要embedding group初始化/broadcast并在optimizer前AllReduce两端grad；finalize_model_grads集中处理此类跨普通DP语义。否则两份权重分别更新后失去tie。

追问得分点（5分）：指出DP不覆盖 1 分；embedding group 1 分；grad归并 2 分；漂移后果 1 分。

M021. Megatron 参数初始化为什么要考虑 TP？

高分回答要点

• 分片参数应与未分片 reference 的逻辑初始化分布一致，且不同 TP ranks 持有正确、不重复的 shard。
• CPU initialization、GPU initialization 和 checkpoint load 的内存峰值/速度不同；大模型常用分布式或 meta 初始化。
• seed、partition axis 和 master weight 构造错误会在第一步前造成跨并行配置不一致。

追问：如何验证 TP1 与 TP2 初始化等价？gather 逻辑权重后逐 tensor 比较，并跑同输入 forward/backward。

前置知识与分析过程：先定义“逻辑全局权重”与“rank local shard”，再区分随机分布等价和逐元素等价。测试需固定seed、partition axis和padding。

主问题得分点（10分）：初始化为何依赖TP 3 分；CPU/GPU路径 2 分；seed/shard一致性 2 分；完整oracle 3 分。

追问分析与参考回答：用相同master seed构造TP1权重；TP2按真实partition axis生成或加载shards，gather后与TP1逐元素比较。再关闭dropout，用同batch比较logits、loss和gathered gradients；一步optimizer后再比。只比较均值/方差不能发现shard重复或轴切错。

追问得分点（5分）：gather逐元素 2 分；同输入forward/backward 1 分；一步更新 1 分；拒绝只比统计量 1 分。

M022. FP16/BF16 训练中 model params、main params、main grads 各是什么？

高分回答要点

• 低精度 model params 用于 forward/backward；optimizer 通常维护 FP32 main params 与 moment states。
• gradients 可在低精度或 FP32 contiguous main-grad buffer 累积，影响每参数 bytes、数值稳定和通信 dtype。
• optimizer step 后要把更新后的 main params 同步/拷回 model params，distributed optimizer 还要 AllGather shards。

追问：为什么只说“FP16 参数 2 bytes”会严重低估 Adam 训练显存？

前置知识与分析过程：对每参数列model param、model/main grad、FP32 master、Adam m/v，再标dtype与是否分片。最后加入bucket/padding和临时buffer边界。

主问题得分点（10分）：三类param/grad定义 3 分；optimizer状态 2 分；通信/更新数据流 3 分；显存公式边界 2 分。

追问分析与参考回答：典型FP16 Adam至少含2B model param、可能2/4B grad、4B FP32 master、4B m和4B v，另有buffer/副本，常接近16-20B/param而非2B。distributed optimizer按DP分片其中大部分，但低精度forward参数和临时AG仍存在。

追问得分点（5分）：列五类状态 2 分；给数量级 1 分；分片影响 1 分；临时/allocator边界 1 分。

M023. forward_step_func 与 loss_func 的接口为何对 PP 很关键？

高分回答要点

• schedule 控制何时运行哪个 microbatch，但不知道业务 batch/loss；callback 让应用提供 stage-local forward。
• 只有 last stage 通常计算 loss，其他 stage 返回 activation；loss dict 用于日志和跨 stage 汇总。
• callback 必须处理 first/last stage 输入差异、packed sequence/attention mask 和 gradient scaling。

追问：为什么每个 stage 都直接读取完整 batch 既浪费又可能破坏数据迭代语义？

前置知识与分析过程：先写schedule对callback的调用次数与stage差异，再明确数据iterator应代表同一逻辑microbatch。分析first/last需要哪些字段。

主问题得分点（10分）：callback控制反转 3 分；first/last输入差异 2 分；loss closure 2 分；数据/scale风险 3 分。

追问分析与参考回答：若每stage独立推进dataset iterator，可能取得不同样本；中间stage只需上游activation和mask/meta，不需token/label完整IO。通常同model replica由指定rank准备/broadcast必要数据，schedule保证microbatch id一致，last stage只用对应label算loss。

追问得分点（5分）：不同样本风险 2 分；字段最小化 1 分；一致microbatch id 1 分；合理broadcast/iterator方案 1 分。

M024. 如何建立 TP/PP 模型 correctness oracle？

高分回答要点

• 小模型上构造 TP1/PP1 reference，gather sharded weights，固定输入、dropout/RNG，比较 logits/loss/grad/一步更新。
• 分开检查 parameter layout、forward、backward、optimizer 与 checkpoint；loss finite 只能排除最粗错误。
• 允许浮点容差但报告 max/relative error，并用多 step 检测小偏差累积。

追问：同 shape collective 顺序错配可能得到 finite loss，因此还要 tensor checksum 和 collective fingerprint。

前置知识与分析过程：把correctness拆为weight layout、forward、backward、optimizer和checkpoint五层；每层选择独立reference与误差口径，不能把最终loss当唯一oracle。

主问题得分点（10分）：TP1 reference 2 分；五层验证 3 分；RNG/容差 2 分；多step与静默通信错误 3 分。

追问分析与参考回答：对每个logical collective记录group、sequence、op、shape/dtype和输入checksum；业务tensor再记录logical id。若A/B同shape跨rank顺序反转，NCCL可完成且输出finite，但fingerprint的logical id/sequence在首个分歧点暴露。小模型同时比较gathered tensor reference。

追问得分点（5分）：指出finite不足 1 分；fingerprint字段 2 分；logical id作用 1 分；reference交叉验证 1 分。

---

---

上一篇：全局架构 · 系列总索引 · 下一篇：TP 与 SP