LLM Lab Notes 05: Attention Tensor Parallelism — Splitting QKV by Head, Megatron Style

Code file: llm_lab/03_tensor_parallel/tp_attention_demo.py

Note 04 covered the FFN; this one handles the other half of the Transformer — Self-Attention. Attention has a natural advantage: heads are mutually independent, which is a structural dividend that makes it easier to parallelize than the FFN.

1. Experiment Goal

1. Implement multi-head attention tensor parallelism: split QKV weights by head count;
2. Compute Q/K/V projections and attention locally, then AllGather the full output;
3. Understand Megatron Self-Attention's "three column weights + one gather" structure.

2. Code Walkthrough

Full code (tp_attention_demo.py, 63 lines):

world = dist.get_world_size()

batch = 4
seq = 128
hidden = 4096

heads = 32
local_heads = heads // world      # ① heads per GPU
head_dim = hidden // heads        # dim per head

local_qkv_dim = local_heads * head_dim

wq = torch.randn(hidden, local_qkv_dim, device="cuda")   # ② three column weights
wk = torch.randn(hidden, local_qkv_dim, device="cuda")
wv = torch.randn(hidden, local_qkv_dim, device="cuda")

x = torch.randn(batch, seq, hidden, device="cuda")

q = x @ wq                        # ③ local Q/K/V projections
k = x @ wk
v = x @ wv

q = q.view(batch, seq, local_heads, head_dim)
k = k.view(batch, seq, local_heads, head_dim)
v = v.view(batch, seq, local_heads, head_dim)

scores = torch.matmul(q, k.transpose(-1, -2)) / (head_dim ** 0.5)   # ④ local attention
attn = F.softmax(scores, dim=-1)

out = torch.matmul(attn, v)

out = out.reshape(batch, seq, local_qkv_dim)

gather_list = [torch.zeros_like(out) for _ in range(world)]
dist.all_gather(gather_list, out)         # ⑤ AllGather to reassemble all heads

full_out = torch.cat(gather_list, dim=-1) # full output [batch, seq, hidden]

① local_heads = heads // world

The key decision in attention parallelism: don't cut inside a head's dimension — evenly divide the 32 heads among GPUs. Each GPU owns the complete computation of local_heads heads. Compared with cutting head_dim, cutting heads makes each GPU's computation self-contained (Q/K/V all local): zero communication inside attention.

② Three column weights (Column Parallel style)

All three projection weights (Q/K/V) are [hidden, local_qkv_dim] — split along columns, no AllGather needed to reassemble (because the next layer "concatenate heads" is exactly the sharding state a Row layer wants). This is Megatron's Self-Attention: three Column weights, no intermediate communication.

③④ Local projections and local attention

After projection, reshape to 4D (batch, seq, local_heads, head_dim); then q @ kᵀ / √d, softmax, and attn @ v all happen locally — attention scores, softmax, and weighted sums need no cross-GPU traffic.

⑤ The only communication point: output AllGather

After computing local-head outputs, all_gather collects every GPU's head fragments and concatenates along the last dim into [batch, seq, hidden]. The entire Attention block has just this one collective.

3. How It Works: Attention's "Structural Dividend"

Why is Attention easier to parallelize than the FFN? Three reasons:

1. Heads are naturally independent: 32 heads are mutually independent; split across 4 GPUs and each GPU runs an independent 8-head attention. The parallelism degree equals the head count (32 heads → up to TP=32 here).

2. QKV projections can skip communication: under Column cuts, each GPU only computes its own heads and never needs other heads' Q/K — attention computation is fully localized.

3. Communication only at the end: the single AllGather happens at output concatenation, and when it pairs with a Row-style output projection (as in Note 04's Column→Row handoff), that one can be saved too.

4. Results

Running on 4 GPUs (32 heads → 8 per GPU), expected output:

Attention output shape: torch.Size([4, 128, 4096])

How to tell it succeeded: compare with single-GPU full attention using allclose. Pay special attention to the per-GPU head order — the concatenated head order must match the sharding convention (heads 0–7 on GPU 0, 8–15 on GPU 1, …); otherwise you get "all-wrong values with all-correct shapes," the most insidious kind of bug.

5. Engineering Takeaways

1. Divisibility constraint between head count and TP degree: heads // world must divide evenly. 32 heads support TP=2/4/8/16/32 but not TP=6. When planning cluster parallelism, the model's head count is a hard constraint.

2. Communication analysis: one AllGather (output size = batch × seq × hidden) — smaller than the FFN's one AllReduce (2× output). Attention is the least communication-hungry part of tensor parallelism, which is why Attention is usually not the bottleneck under TP.

3. Write the concatenation convention into comments: the order of cross-GPU heads and the numbering rule for local_heads are the most likely places for TP code to "silently compute wrong results." In production, add a torch.allclose assertion (against single-GPU) right after the split — cheap insurance.

4. Combination with FlashAttention (previewed in Note 07): in production, the local_heads attention computation is replaced with flash_attn_func — TP's sharding logic stays unchanged; only the local compute kernel is swapped. That's the modular power of parallel frameworks.

Next: with both Attention and FFN splittable, assemble them into a complete Megatron-style Transformer Block — Note 06: LLM Lab Notes 06: Building a Megatron-Style Transformer Block from Scratch.