LLM Lab 实验手记 06：从零拼一个 Megatron 风格 Transformer Block
代码文件：
llm_lab/04_transformer/（mini_transformer_block.py、mini_megatron_layer.py、mini_transformer_train.py、mini_megatron_profile.py）
前 5 篇把分布式训练的"零件"都过了一遍：进程组、DDP、AllReduce 自动微分、FFN 与 Attention 的张量并行。本篇把它们焊成一个完整的 Transformer Block——这也是 Megatron-LM 中一个 transformer 层的最小复刻，且支持反向传播与逐段性能剖析。
1. 实验目标
- 组装 Megatron 风格 Transformer Block：LayerNorm → Attention（TP）→ 残差 → LayerNorm → FFN（TP）→ 残差；
- 让整个块可反向传播（用第 3 篇的
AllReduceautograd 算子）； - 对 Attention / FFN / Backward 做逐段计时，建立"哪段耗时最多"的第一手感觉。
2. 代码逐段拆解（mini_megatron_layer.py，最完整版）
class AllReduce(torch.autograd.Function): # ① 复用第 3 篇的通信算子
@staticmethod
def forward(ctx, x):
x = x.contiguous()
dist.all_reduce(x)
return x
@staticmethod
def backward(ctx, grad_output):
grad_output = grad_output.contiguous()
dist.all_reduce(grad_output)
return grad_output
x = torch.randn(batch, seq, hidden, device="cuda", requires_grad=True)
# ---- LayerNorm ----
x_norm = F.layer_norm(x, (hidden,)) # ② 先归一化
# ---- Attention（TP：按 head 切）----
wq = torch.randn(hidden, local_heads * head_dim, device="cuda", requires_grad=True)
wk = ...
wv = ...
q = (x_norm @ wq).view(batch, seq, local_heads, head_dim)
k = (x_norm @ wk).view(batch, seq, local_heads, head_dim)
v = (x_norm @ wv).view(batch, seq, local_heads, head_dim)
scores = torch.matmul(q, k.transpose(-1, -2)) / (head_dim ** 0.5)
attn = F.softmax(scores, dim=-1)
attn_out = torch.matmul(attn, v).reshape(batch, seq, local_heads * head_dim)
gather = [torch.zeros_like(attn_out) for _ in range(world)]
dist.all_gather(gather, attn_out) # ③ Attention 唯一通信点
attn_full = torch.cat(gather, dim=-1)
# ---- Residual ----
x = x + attn_full # ④ 残差连接
# ---- FFN（TP：Column→GELU→Row）----
x_norm = F.layer_norm(x, (hidden,)) # ⑤ 第二层 Norm（Post-LN）
w1 = torch.randn(hidden, ffn_chunk, device="cuda", requires_grad=True)
w2 = torch.randn(ffn_chunk, hidden, device="cuda", requires_grad=True)
mid = F.gelu(x_norm @ w1)
local_out = mid @ w2
local_out = AllReduce.apply(local_out) # ⑥ FFN 唯一通信点
# ---- Residual ----
out = x + local_out # ⑦ 残差
# ---- Loss ----
loss = out.mean()
loss.backward() # ⑧ 端到端反向
if rank == 0:
print("Loss:", loss.item())
print("Megatron layer backward OK")① 通信算子的复用
第 3 篇手写的 AllReduce 在这里直接复用——这正是自定义算子的价值：写一次，全模型复用。
②⑤ Post-LN 结构
本实验是 Post-LayerNorm 结构：先残差相加，再做 LayerNorm（与原始 Transformer 一致）。Megatron 与当前主流 LLM 多用 Pre-LN（先 Norm 再算），但结构骨架一致——只需挪动 Norm 与残差的位置。
③⑥ 两个通信点，对应第 4/5 篇
- Attention 输出端
all_gather（拼 head）； - FFN 输出端
AllReduce（汇总部分和）。 - 中间没有通信——因为 Column→Row 的分片状态自然衔接（第 4 篇的结论）。
④⑦ 残差：免通信
残差连接是逐元素加法，完全本地。这让"通信点只存在于 Attention/FFN 收尾"成为可能。
⑧ 端到端反向
loss.backward() 会沿残差、Norm、Attention、FFN 一路回传，并在 AllReduce.apply 的反向里再次触发通信（第 3 篇验证过的机制）。打印"backward OK" = 梯度流完整贯通。
3. 原理讲解：一个 Block 的通信账本
把整个 Block 的集合通信列成一张表：
| 位置 | 集合操作 | 通信量 | 对应实验 |
|---|---|---|---|
| Attention 输出 | AllGather | batch × seq × hidden | 第 5 篇 |
| FFN 输出 | AllReduce | 2 × batch × seq × hidden | 第 4 篇 |
| 残差 / GELU / LayerNorm | 无 | 0 | — |
| 反向（两个通信点各一次） | AllReduce×2 | ≈ 前向反向合计 | 第 3 篇 |
结论：一个 TP 的 Transformer Block，前向只有 2 次集合通信（1 次 AllGather + 1 次 AllReduce），反向同样 2 次——整个大模型的通信开销可以精确预估，这就是规划并行方案时的理论依据。
4. 实验结果
可训练版（mini_transformer_train.py）：4 卡运行，预期 rank 0 打印：
Loss: 0.2305
Backward done性能剖析版（mini_megatron_profile.py，seq=512，4 卡）：每步输出分段耗时：
Attention: 18.42 ms
FFN : 21.76 ms
Backward : 40.35 ms（数值因硬件而异；核心观察点是三个结论：① Backward ≈ 前向之和（反向要通信+重算梯度）；② FFN 与 Attention 量级相当（FFN 参数量是 Attention 的 4 倍左右，但 Attention 的 softmax 与 matmul 也吃时间）；③ 改变 seq（128→256→512）时 Attention 时间按平方增长、FFN 线性增长——这是定位长上下文瓶颈的第一手证据。）
5. 工程启示
"从零拼装"是最好的调试方法：生产框架（Megatron/DeepSpeed）几千行代码里藏着并行细节，真出问题时很难定位。先用这种 100 行级别的 Block 验证"通信点正确 + 梯度贯通"，再对照生产实现，排障效率完全不同。
Post-LN vs Pre-LN 不只是学术问题：Pre-LN 训练更稳定（缓解梯度消失），是 GPT 系标配；Post-LN 是原版 Transformer。改动只是 Norm 位置，但会影响可训练性与最终效果——迁移架构时别忽略。
性能剖析要分段打点：
torch.cuda.synchronize()+time.time()是最朴素也最可靠的剖析手段。生产上可升级为nsys/nsight，但先做粗粒度分段、确定瓶颈在哪一段，再上细粒度工具，方向才不会错。本系列的通路：这个 Block 加上 Embedding、位置编码、多层堆叠、分布式 checkpoint 与数据加载，就是 Megatron 训练一个 7B/70B 模型的骨架——你看到的每个生产组件，都能在这个最小实现里找到对应物。
下一步：本地计算内核还能再升级——第 7 篇对比 FlashAttention 与标准注意力：LLM Lab 实验手记 07：FlashAttention 为什么快。