LLM Lab Notes 06: Building a Megatron-Style Transformer Block from Scratch

Code files: llm_lab/04_transformer/ (mini_transformer_block.py, mini_megatron_layer.py, mini_transformer_train.py, mini_megatron_profile.py)

The first five notes covered all the "parts" of distributed training: process groups, DDP, AllReduce autograd, FFN and Attention tensor parallelism. This note welds them into a complete Transformer Block — a minimal reproduction of one Megatron-LM transformer layer that supports backpropagation and per-segment performance profiling.

1. Experiment Goal

1. Assemble a Megatron-style Transformer Block: LayerNorm → Attention (TP) → residual → LayerNorm → FFN (TP) → residual;
2. Make the whole block backprop-able (using Note 03's AllReduce autograd operator);
3. Time Attention / FFN / Backward separately to build first-hand intuition about "which segment costs the most."

2. Code Walkthrough (mini_megatron_layer.py, most complete version)

class AllReduce(torch.autograd.Function):        # ① reuse Note 03's comm operator
    @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,))              # ② normalize first

# ---- Attention (TP: split by 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's only comm point
attn_full = torch.cat(gather, dim=-1)

# ---- Residual ----
x = x + attn_full                                # ④ residual connection

# ---- FFN (TP: Column→GELU→Row) ----
x_norm = F.layer_norm(x, (hidden,))              # ⑤ second 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's only comm point

# ---- Residual ----
out = x + local_out                              # ⑦ residual

# ---- Loss ----
loss = out.mean()
loss.backward()                                  # ⑧ end-to-end backward

if rank == 0:
    print("Loss:", loss.item())
    print("Megatron layer backward OK")

① Reusing the communication operator

The AllReduce hand-written in Note 03 is reused directly — the value of custom operators: write once, reuse across the whole model.

②⑤ Post-LN structure

This experiment uses the Post-LayerNorm structure: residual first, then LayerNorm (consistent with the original Transformer). Megatron and today's mainstream LLMs mostly use Pre-LN (Norm first, then compute), but the skeleton is identical — just move the Norm and residual positions.

③⑥ Two communication points, from Notes 04/05

• all_gather at the Attention output (reassemble heads);
• AllReduce at the FFN output (aggregate partial sums).
• No communication in between — because Column→Row sharding states connect naturally (Note 04's conclusion).

④⑦ Residual: communication-free

Residual addition is element-wise, fully local. This is what makes "communication only at Attention/FFN tails" possible.

⑧ End-to-end backward

loss.backward() flows back through residual, Norm, Attention, FFN — and triggers communication again inside AllReduce.apply's backward (the mechanism verified in Note 03). Printing "backward OK" = gradient flow is fully connected.

3. How It Works: The Communication Ledger of One Block

List every collective in the block as a table:

Location	Collective	Traffic	Corresponding note
Attention output	AllGather	batch × seq × hidden	Note 05
FFN output	AllReduce	2 × batch × seq × hidden	Note 04
Residual / GELU / LayerNorm	none	0	—
Backward (one at each comm point)	AllReduce×2	≈ forward+backward total	Note 03

Conclusion: a TP Transformer Block has only 2 collectives in forward (1 AllGather + 1 AllReduce) and 2 in backward — the whole model's communication overhead can be estimated precisely, which is the theoretical basis for planning a parallel scheme.

4. Results

Trainable version (mini_transformer_train.py): on 4 GPUs, expected rank 0 output:

Loss: 0.2305
Backward done

Profiled version (mini_megatron_profile.py, seq=512, 4 GPUs): per-step segment timing:

Attention: 18.42 ms
FFN      : 21.76 ms
Backward : 40.35 ms

(Values vary by hardware; the core observations are three: ① Backward ≈ sum of forward (backward communicates and recomputes gradients); ② FFN and Attention are comparable in cost (FFN has ~4× the parameters of Attention, but Attention's softmax and matmuls also take time); ③ as seq grows (128→256→512), Attention time grows quadratically while FFN grows linearly — first-hand evidence for locating long-context bottlenecks.)

5. Engineering Takeaways

1. "Assembling from scratch" is the best debugging method: production frameworks (Megatron/DeepSpeed) hide parallel details in thousands of lines; when something breaks, locating it is hard. First validate "communication points correct + gradients connected" with a 100-line block like this, then compare against production implementations — much more efficient troubleshooting.

2. Post-LN vs Pre-LN is not just academic: Pre-LN trains more stably (mitigates vanishing gradients) and is standard in GPT-family models; Post-LN is the original Transformer. The change is just the Norm position, but it affects trainability and final quality — don't overlook it when porting architectures.

3. Profile in segments: torch.cuda.synchronize() + time.time() is the simplest and most reliable profiling approach. In production, upgrade to nsys/nsight, but do coarse segmentation first, determine which segment is the bottleneck, then go fine-grained — right direction beats right tools.

4. The path of this series: this Block plus embeddings, positional encodings, layer stacking, distributed checkpointing and data loading is the skeleton of training a 7B/70B model in Megatron — every production component has a counterpart in this minimal implementation.

Next: the local compute kernel can be upgraded too — Note 07 compares FlashAttention vs standard attention: LLM Lab Notes 07: Why FlashAttention Is Fast.