GPU Systems and Performance Engineering on Synaptic Radio

Megatron Tensor Parallelism: Rebuilding a SwiGLU Feed-Forward Block Across Two Ranks

anshuman264@gmail.com (Anshuman Sahoo) — Thu, 16 Jul 2026 00:00:00 +0000

Megatron training is often introduced through launcher flags: tensor parallelism, pipeline parallelism, sequence parallelism, distributed optimizer, recomputation, and process groups. That is useful when operating the framework, but it is a poor way to learn what the framework is doing.

I wanted to understand one smaller question first:

How can two ranks reproduce the forward pass, backward pass, and parameter update of one dense Transformer feed-forward sublayer without either rank storing the full intermediate width?

So I rebuilt that path with ordinary PyTorch and two CPU processes. The block is not a complete Transformer. It contains the part under test: pre-RMSNorm, a fused gate/up projection, SwiGLU, a down projection, and a residual connection.

The finished two-rank step matched the dense reference to within 1.27e-07 in the forward pass and 1.49e-08 after one SGD update. More importantly, two deliberately broken versions failed in different places. Removing the forward reduction changed the output immediately. Removing the backward reduction left the output and loss correct, but corrupted the gradient entering RMSNorm.

That difference is the useful lesson. Tensor parallelism is not merely a way to divide weights. It is a way to divide one algebraic operation while restoring the missing sums at the exact boundaries where the dense computation requires them.

Start with the dense block

Here is the dense computation the two ranks must reproduce:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41


class DenseSwiGLUBlock(nn.Module):
 r"""
 ┌──────────────────────────────┐
 │ │
 │ residual path │
 │ ▼
 input x ────────────┼───────────────────────────── (+) ── output
 │ ▲
 ▼ │
 RMSNorm │
 │ │
 ┌──────┴──────┐ │
 ▼ ▼ │
 gate projection up projection │
 │ │ │
 ▼ │ │
 SiLU │ │
 │ │ │
 └────── × ────┘ │
 │ │
 ▼ │
 down projection ──────────────────────┘
 """

 def __init__(self, hidden_size: int, ffn_size: int):
 super().__init__()
 self.norm_weight = nn.Parameter(torch.ones(hidden_size))
 # Stack gate and up projections so the tensor-parallel version can shard the shared FFN dimension.
 self.w_gate_up = nn.Parameter(torch.empty(2, ffn_size, hidden_size))
 self.w_down = nn.Parameter(torch.empty(hidden_size, ffn_size))
 self.b_down = nn.Parameter(torch.zeros(hidden_size))

 def forward(self, x):
 residual = x
 x = rms_norm(x, self.norm_weight)

 gate = F.linear(x, self.w_gate_up[0])
 up = F.linear(x, self.w_gate_up[1])
 hidden = F.silu(gate) * up

 return residual + F.linear(hidden, self.w_down, self.b_down)

The executable fixture uses:

sequence length S = 3;
batch size B = 2;
model width H = 8;
feed-forward width F = 12;
tensor-parallel size P = 2.

The dense shape trace is:

1
2
3
4
5
6
7


input [3, 2, 8]
RMSNorm output [3, 2, 8]
gate projection [3, 2, 12]
up projection [3, 2, 12]
SwiGLU activation [3, 2, 12]
down projection [3, 2, 8]
residual output [3, 2, 8]

The equations are:

\[ N = \operatorname{RMSNorm}(X), \]\[ G = N W_g^\top, \qquad U = N W_u^\top, \]\[ A = \operatorname{SiLU}(G) \odot U, \]\[ Y = X + A W_d^\top + b_d. \]

The gate and up matrices each have shape [F, H]. The down matrix has shape [H, F]. That shared feed-forward dimension is the axis we will partition.

Split the expansion

The first change is to divide the gate and up projections by output feature. Each rank receives the full normalized token representation but computes only half of the feed-forward channels.

For two ranks:

\[ W_g = \begin{bmatrix} W_g^{(0)} \\ W_g^{(1)} \end{bmatrix}, \qquad W_u = \begin{bmatrix} W_u^{(0)} \\ W_u^{(1)} \end{bmatrix}. \]

Rank r computes:

\[ G^{(r)} = N\left(W_g^{(r)}\right)^\top, \qquad U^{(r)} = N\left(W_u^{(r)}\right)^\top, \]\[ A^{(r)} = \operatorname{SiLU}\left(G^{(r)}\right) \odot U^{(r)}. \]

The nonlinear operation remains local because each gate feature is paired with the corresponding up feature on the same rank.

1
2
3
4
5


normed_parallel = CopyToTensorParallelRegion.apply(normed)

gate_local = F.linear(normed_parallel, w_gate_up_local[0])
up_local = F.linear(normed_parallel, w_gate_up_local[1])
hidden_local = F.silu(gate_local) * up_local

Nothing needs to be summed in the forward pass yet. Rank 0 owns one set of feed-forward features; rank 1 owns the other. Concatenating those slices would recover the dense SwiGLU activation.

The shape correspondence is easier to see in one table:

Tensor	Dense shape	Shape on each rank	Partition
Input `X`	`[3, 2, 8]`	`[3, 2, 8]`	replicated
RMSNorm output `N`	`[3, 2, 8]`	`[3, 2, 8]`	replicated
Gate weight	`[12, 8]`	`[6, 8]`	output features
Up weight	`[12, 8]`	`[6, 8]`	output features
Local SwiGLU output	`[3, 2, 12]`	`[3, 2, 6]`	feed-forward features
Down weight	`[8, 12]`	`[8, 6]`	input features
Partial down output	—	`[3, 2, 8]`	partial sum
Reduced output	`[3, 2, 8]`	`[3, 2, 8]`	replicated

In Megatron terminology, this expansion is the role played by a column-parallel linear layer: each rank owns different output columns of the logical projection.

Split the contraction

Now partition the down projection along its input feature dimension. Each rank consumes the SwiGLU slice it already owns:

\[ W_d = \begin{bmatrix} W_d^{(0)} & W_d^{(1)} \end{bmatrix}. \]

Each rank computes a full-width but incomplete output:

\[ Z^{(r)} = A^{(r)}\left(W_d^{(r)}\right)^\top. \]

The dense result is the sum of those partial products:

\[ Z = \sum_{r=0}^{P-1} Z^{(r)}. \]

That equation tells us exactly where the first collective belongs.

1
2
3


partial = F.linear(hidden_local, w_down_local, None)
mlp_out = ReduceFromTensorParallelRegion.apply(partial)
out = residual + mlp_out + b_down

ReduceFromTensorParallelRegion performs an all-reduce in the forward pass:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class ReduceFromTensorParallelRegion(torch.autograd.Function):
 @staticmethod
 def forward(ctx, x):
 y = x.clone()
 dist.all_reduce(y, op=dist.ReduceOp.SUM)
 return y

 @staticmethod
 def backward(ctx, grad_output):
 return grad_output

The backward pass through this operation is an identity. Every rank receives the same gradient for the replicated output, and each local down-projection shard uses that gradient to compute its own weight and activation gradients.

This is the row-parallel half of the pair: split the input features, compute partial outputs, and sum them.

Put the other collective in backward

The expansion layer has the opposite communication pattern.

Its forward pass needed no reduction because the output features were intentionally partitioned. But during backpropagation, every rank computes only the contribution from its local gate and up features to the gradient of the replicated normalized input.

For the dense input gradient:

$$ \begin{aligned} \frac{\partial L}{\partial N} &= \sum_{r=0}^{P-1} \left(\frac{\partial L}{\partial N}\right)_r . \end{aligned} $$

That sum must happen before the gradient continues through RMSNorm and into the residual input. The custom operation around normed therefore does nothing in forward and reduces gradients in backward:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class CopyToTensorParallelRegion(torch.autograd.Function):
 @staticmethod
 def forward(ctx, x):
 return x

 @staticmethod
 def backward(ctx, grad_output):
 grad = grad_output.clone()
 dist.all_reduce(grad, op=dist.ReduceOp.SUM)
 return grad

The two communication rules are now paired:

Logical layer	Forward	Backward
Column-parallel gate/up projection	keep local feature slices	sum replicated-input gradient
Row-parallel down projection	sum partial outputs	keep local shard gradients

This maps to the core Megatron pattern:

Reconstruction	Megatron-style concept	Why it exists
Fused local gate/up weights	`ColumnParallelLinear` role	partition expansion outputs
Local down-projection weights	`RowParallelLinear` role	consume partitioned hidden features
Identity forward, reduce backward	copy-to-tensor-parallel region	combine contributions to replicated input gradients
Reduce forward, identity backward	reduce-from-tensor-parallel region	reconstruct the dense contraction output

The fixture implements these semantics directly with torch.distributed; it does not import Megatron-Core classes.

Reproduce one dense training step

First inspect the actual environment, block, and tensor shapes:

1
2


python -m torch.distributed.run --standalone --nproc_per_node=2 \
 code/megatron_tp_mlp.py --mode inspect

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


{
 "python": "3.13.5",
 "pytorch": "2.10.0+cpu",
 "backend": "gloo",
 "world_size": 2,
 "block_path": "pre-RMSNorm -> fused gate/up -> SwiGLU -> down projection -> residual add",
 "dense_shapes": {
 "x": [3, 2, 8],
 "norm_weight": [8],
 "w_gate_up": [2, 12, 8],
 "w_down": [8, 12],
 "output": [3, 2, 8]
 },
 "rank_local_shapes": {
 "gate_up_weight": [2, 6, 8],
 "gate": [3, 2, 6],
 "up": [3, 2, 6],
 "swiglu": [3, 2, 6],
 "down_weight": [8, 6],
 "partial_output": [3, 2, 8]
 }
}

The launcher’s Gloo connection lines are omitted from this embedded excerpt; the complete retained output is in data/terminal-01-inspect.txt. The environment record is in data/environment.json.

Now compare one complete sharded training step with the dense reference initialized from the same tensors:

1
2


python -m torch.distributed.run --standalone --nproc_per_node=2 \
 code/megatron_tp_mlp.py --mode equivalence

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


{
 "mode": "equivalence",
 "world_size": 2,
 "forward_max_abs_diff": 1.2665987014770508e-07,
 "loss_abs_diff": 0.0,
 "input_grad_max_abs_diff": 2.9802322387695312e-08,
 "norm_grad_max_abs_diff": 1.4901161193847656e-08,
 "gate_up_grad_max_abs_diff": 4.470348358154297e-08,
 "down_grad_max_abs_diff": 2.2351741790771484e-08,
 "down_bias_grad_max_abs_diff": 2.2351741790771484e-08,
 "post_step_weight_max_abs_diff": 1.4901161193847656e-08,
 "passed": true
}

The complete transcript is data/terminal-02-equivalence.txt.

Check	Maximum absolute difference
Forward output	`1.27e-07`
Loss	`0.00`
Input gradient	`2.98e-08`
RMSNorm-weight gradient	`1.49e-08`
Fused gate/up gradient	`4.47e-08`
Down-projection gradient	`2.24e-08`
Down-projection bias gradient	`2.24e-08`
Parameters after one SGD step	`1.49e-08`

The differences are consistent with floating-point operation ordering. Within this CPU fixture, the two-rank path reproduces the dense forward pass, backward pass, and parameter update.

Remove each collective

An equivalence result is useful only if the test can reject plausible wrong implementations. The two broken versions below are designed to fail at different points.

Remove the forward reduction

Prediction: each rank will treat a partial down-projection result as the complete MLP output, so the first mismatch should appear in the forward output.

1
2


python -m torch.distributed.run --standalone --nproc_per_node=2 \
 code/megatron_tp_mlp.py --mode missing_forward_reduce

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


{
 "mode": "missing_forward_reduce",
 "world_size": 2,
 "forward_max_abs_diff": 1.1021764278411865,
 "loss_abs_diff": 0.19488954544067383,
 "input_grad_max_abs_diff": 0.03851410746574402,
 "norm_grad_max_abs_diff": 0.04135049879550934,
 "gate_up_grad_max_abs_diff": 0.05925934761762619,
 "down_grad_max_abs_diff": 0.026660695672035217,
 "post_step_weight_max_abs_diff": 0.0036689937114715576,
 "passed": false
}

The retained output is data/terminal-03-missing-forward.txt.

The prediction holds. The first changed signal is the model output, with a maximum error of 1.10. Once the loss is computed from the wrong output, every downstream gradient is also wrong.

Remove the backward reduction

Prediction: the forward output and loss will still match because the row-parallel forward sum remains intact. The first mismatch should appear only when gradients from the two local expansion shards need to be combined before RMSNorm.

1
2


python -m torch.distributed.run --standalone --nproc_per_node=2 \
 code/megatron_tp_mlp.py --mode missing_backward_reduce

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


{
 "mode": "missing_backward_reduce",
 "world_size": 2,
 "forward_max_abs_diff": 1.2665987014770508e-07,
 "loss_abs_diff": 0.0,
 "input_grad_max_abs_diff": 0.13864503800868988,
 "norm_grad_max_abs_diff": 0.20240136981010437,
 "gate_up_grad_max_abs_diff": 4.470348358154297e-08,
 "down_grad_max_abs_diff": 2.2351741790771484e-08,
 "post_step_weight_max_abs_diff": 0.010120153427124023,
 "passed": false
}

The retained output is data/terminal-04-missing-backward.txt.

This is the more revealing failure. The forward output still agrees to 1.27e-07, and the loss is identical. The local gate/up and down-projection gradients also match. Yet the input gradient differs by 0.139, the RMSNorm-weight gradient differs by 0.202, and one optimizer step moves the parameters apart by 0.0101.

A forward-only test would approve this broken training graph.

The distributed test suite checks all four cases:

1

pytest -q code/test_megatron_tp_mlp.py

1
2


.... [100%]
4 passed in 26.94s

The retained test output is data/terminal-05-tests.txt. The complete implementation is code/megatron_tp_mlp.py, and the direct tests are code/test_megatron_tp_mlp.py.

What this reconstruction proves—and what it does not

The small implementation establishes one precise result: for this pre-RMSNorm SwiGLU feed-forward block, partitioning the expansion by output features and the contraction by input features reproduces dense training when the partial contraction outputs are summed in forward and the partial replicated-input gradients are summed in backward.

It also shows why the two reductions need separate tests. A missing forward collective is visible immediately. A missing backward collective can survive output and loss checks while silently corrupting upstream learning.

The fixture does not execute Megatron-Core. It does not test:

self-attention or a complete Transformer block;
sequence parallelism;
CUDA or NCCL collective behavior;
BF16 or FP8 arithmetic;
fused RMSNorm, SwiGLU, or linear kernels;
asynchronous communication overlap;
gradient accumulation;
pipeline parallelism;
distributed optimizer state;
more than two tensor-parallel ranks.

Those are not cosmetic differences. For example, sequence parallelism changes which activations are replicated, fused kernels change numerical ordering, and larger process groups change communication cost and failure behavior.

Still, the reconstruction changes how I read Megatron code. I no longer start with “which tensors are split?” I start with the dense equation and ask two questions:

Which sum has been replaced by independent local products?
At what point must those products be combined so forward and backward remain equivalent?

For the feed-forward block, those questions lead directly to the matched column-parallel and row-parallel projections—and to one collective in each direction.

References

NVIDIA, Megatron Core User Guide.
NVIDIA, Megatron Core parallelism strategies.
NVIDIA, Tensor-parallel layers API.

Tensor-Parallel GQA Attention, One Head Group at a Time

anshuman264@gmail.com (Anshuman Sahoo) — Thu, 16 Jul 2026 00:00:00 +0000

Grouped-query attention creates a partitioning problem that ordinary multi-head attention can hide. There are more query heads than key/value heads, so several query heads reuse the same KV head. If a tensor-parallel split cuts projection matrices at an arbitrary width, it can separate a query head from the KV head it is supposed to share.

I wanted to reconstruct the smallest attention block that still preserves that constraint. The reference is a pre-RMSNorm, causal GQA layer with rotary position embeddings, four query heads, two KV heads, and a residual connection. I then turn it into two rank-local attention paths in visible steps and compare the resulting training update against the dense block.

The result is narrow but useful: the natural partition is a complete query/KV group, not an arbitrary matrix slice. The forward pass also needs an output reduction, while the replicated input path needs a separate backward reduction. A forward-only test catches the first bug and misses the second.

Start with the dense block

The dense reference is short enough to inspect:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38


def apply_rope(x):
 """
 x[..., :half] ─┐
 ├─ rotate by position/frequency phase ── RoPE(x)
 x[..., half:] ─┘
 """
 # x: [S,B,H,D], D even
 s = x.shape[0]
 half = x.shape[-1] // 2
 pos = torch.arange(s, dtype=x.dtype, device=x.device).view(s,1,1,1)
 freq = torch.arange(half, dtype=x.dtype, device=x.device).view(1,1,1,half)
 theta = pos / (10000.0 ** (freq / max(1, half)))
 c, si = theta.cos(), theta.sin()
 a, b = x[..., :half], x[..., half:]
 return torch.cat([a*c - b*si, a*si + b*c], dim=-1)


def repeat_kv(x, groups):
 """
 KV heads [S,B,Hkv,D] ── repeat each head `groups` times ── [S,B,Hq,D]
 """
 return x.repeat_interleave(groups, dim=2)


def dense_forward(x, p, cfg):
 """
 x ── RMSNorm ── Q,K,V ── RoPE ── repeat KV ── attention ── output proj ─┐
 └──────────────────────────────────────── residual ─────────────────── (+)
 """
 residual = x
 n = rms_norm(x, p['norm'], cfg.eps)
 q = F.linear(n, p['wq']).view(cfg.seq,cfg.batch,cfg.query_heads,cfg.head_dim)
 k = F.linear(n, p['wk']).view(cfg.seq,cfg.batch,cfg.kv_heads,cfg.head_dim)
 v = F.linear(n, p['wv']).view(cfg.seq,cfg.batch,cfg.kv_heads,cfg.head_dim)
 q, k = apply_rope(q), apply_rope(k)
 groups = cfg.query_heads // cfg.kv_heads
 ctx = attention(q, repeat_kv(k, groups), repeat_kv(v, groups))
 return residual + F.linear(ctx.reshape(cfg.seq,cfg.batch,-1), p['wo'])

The model dimensions are intentionally tiny, but the semantic axes are real:

Tensor	Shape	Meaning
Input `x`	`[5, 2, 16]`	sequence, batch, hidden
Query `q`	`[5, 2, 4, 4]`	four query heads
Key/value `k`, `v`	`[5, 2, 2, 4]`	two KV heads
Repeated KV	`[5, 2, 4, 4]`	one KV head per two query heads
Context	`[5, 2, 4, 4]`	one context vector per query head
Output	`[5, 2, 16]`	residual-width result

The first terminal command records the actual fixture and versions:

1

PYTHONPATH=code python code/tp_gqa_attention.py --mode inspect

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


{
 "python": "3.13.5",
 "pytorch": "2.10.0+cpu",
 "config": {
 "seq": 5,
 "batch": 2,
 "hidden": 16,
 "query_heads": 4,
 "kv_heads": 2,
 "head_dim": 4,
 "world_size": 2,
 "eps": 1e-06
 },
 "dense_shapes": {
 "x": [5, 2, 16],
 "q": [5, 2, 4, 4],
 "k_v": [5, 2, 2, 4]
 },
 "rank_local": {
 "query_heads": 2,
 "kv_heads": 1
 },
 "rule": "keep complete query/KV groups on each rank; sum output-projection partials"
}

The full transcript is retained in data/terminal-01-inspect.txt, and the executable reference is in code/tp_gqa_attention.py.

First change: split complete head groups

Each KV head is shared by exactly two query heads. That gives two indivisible groups:

1
2


group 0: query heads 0, 1 + KV head 0
group 1: query heads 2, 3 + KV head 1

Rank 0 receives group 0. Rank 1 receives group 1. In matrix terms, the Q projection is split by blocks of 2 × head_dim = 8 output rows, while K and V are split by blocks of 1 × head_dim = 4 output rows.

The local code is the dense code with smaller head counts:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


q_local = F.linear(n, wq_local).view(S, B, 2, 4)
k_local = F.linear(n, wk_local).view(S, B, 1, 4)
v_local = F.linear(n, wv_local).view(S, B, 1, 4)

q_local = apply_rope(q_local)
k_local = apply_rope(k_local)

k_local = repeat_kv(k_local, groups=2)
v_local = repeat_kv(v_local, groups=2)
context_local = causal_attention(q_local, k_local, v_local)

No communication is needed inside this local attention calculation. RoPE operates within each head, the causal mask is shared, and the rank already owns the KV head needed by both of its query heads. This matches the original Megatron observation that attention heads can be computed locally after Q, K, and V are column-partitioned by head ownership.

The important restriction is divisibility. With two TP ranks, two KV heads divide cleanly. A configuration with two KV heads and four TP ranks cannot assign one complete KV head to every rank without replication or a different layout. The toy fixture does not solve that more general case.

Second change: split the output projection

Local attention produces only half of the dense context width:

Tensor	Dense shape	Rank-local shape	Ownership
Normalized input	`[5, 2, 16]`	`[5, 2, 16]`	replicated
Query	`[5, 2, 4, 4]`	`[5, 2, 2, 4]`	partitioned by query heads
Key/value	`[5, 2, 2, 4]`	`[5, 2, 1, 4]`	partitioned by KV heads
Local context	—	`[5, 2, 2, 4]`	partitioned
Output partial	—	`[5, 2, 16]`	partial sum
Final output	`[5, 2, 16]`	`[5, 2, 16]`	replicated after reduction

The dense output projection has shape [hidden, query_heads × head_dim] = [16, 16]. Each rank owns the columns corresponding to its local context features. If C^(r) is rank r’s local context and W_o^(r) is the matching column slice, then

\[ Y = X + \sum_{r=0}^{P-1} C^{(r)}\left(W_o^{(r)}\right)^\top. \]

Each rank can compute a full-width partial:

1
2
3
4


partial = F.linear(
 context_local.reshape(S, B, -1),
 wo_local,
)

But no partial is the dense answer. The rank-local outputs must be summed before the residual addition.

This is the row-parallel half of the reconstruction. It also explains why gathering the local contexts first would be wasteful: the output projection can consume the partitioned context directly, and only its full-width partial outputs need to be reduced.

Rebuild the complete training step

The compact tensor-parallel path now has two transformations:

partition Q, K, and V by complete GQA head groups;
partition the output projection by the corresponding context features and sum the partial outputs.

The implementation uses one process to make the algebra easy to inspect, while treating the two branches as rank-local computations with shared parameters sliced exactly as two TP ranks would own them. It is therefore a numerical reconstruction of the TP layer, not a benchmark of distributed communication.

Run the direct comparison:

1

PYTHONPATH=code python code/tp_gqa_attention.py --mode equivalence --seed 11

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


{
 "seed": 11,
 "mode": "equivalence",
 "forward_max_abs_diff": 1.1920928955078125e-07,
 "loss_abs_diff": 0.0,
 "input_grad_max_abs_diff": 3.725290298461914e-09,
 "parameter_grad_max_abs_diff": 2.2351741790771484e-08,
 "post_step_weight_max_abs_diff": 1.4901161193847656e-08,
 "passed": true
}

The comparison uses identical inputs, parameters, target tensor, MSE loss, and SGD learning rate. It checks more than the visible output:

Check	Seed 11 maximum absolute difference	Five-seed maximum
Forward output	`1.19e-07`	`2.38e-07`
Loss	`0.00`	`0.00`
Input gradient	`3.73e-09`	`1.12e-08`
Parameter gradients	`2.24e-08`	`4.47e-08`
Parameters after SGD	`1.49e-08`	`2.98e-08`

Those residuals are consistent with different floating-point accumulation order, not a different computation. The retained five-seed results are in data/results.csv and data/results.json.

Break forward and backward separately

A reconstruction is more convincing when a plausible wrong version fails for the reason predicted by the algebra.

Remove the output reduction

Prediction: the first wrong signal should be the forward output because one rank-local partial is being mistaken for the complete output.

1
2


PYTHONPATH=code python code/tp_gqa_attention.py \
 --mode missing_output_reduce --seed 11

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


{
 "seed": 11,
 "mode": "missing_output_reduce",
 "forward_max_abs_diff": 0.700843870639801,
 "loss_abs_diff": 0.05957293510437012,
 "input_grad_max_abs_diff": 0.013707960024476051,
 "parameter_grad_max_abs_diff": 0.12953613698482513,
 "post_step_weight_max_abs_diff": 0.0064768195152282715,
 "passed": false
}

The prediction holds. Across five seeds, the forward error ranges from 0.324 to 0.937. Everything downstream is already operating on the wrong activation.

Remove the replicated-input gradient reduction

The less obvious bug is in backward. Both column-partitioned projection branches read the same normalized input. In a real TP implementation, each rank computes only its local contribution to the gradient with respect to that replicated input. Those contributions must be summed.

The fixture simulates a missing reduction by allowing both branches to contribute in forward while dropping rank 1’s contribution to the upstream input gradient.

Prediction: the forward output and loss should still match, but the first wrong signal should be dX.

1
2


PYTHONPATH=code python code/tp_gqa_attention.py \
 --mode missing_input_grad_reduce --seed 11

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


{
 "seed": 11,
 "mode": "missing_input_grad_reduce",
 "forward_max_abs_diff": 1.1920928955078125e-07,
 "loss_abs_diff": 0.0,
 "input_grad_max_abs_diff": 0.01103300042450428,
 "parameter_grad_max_abs_diff": 0.0307764932513237,
 "post_step_weight_max_abs_diff": 0.0015388727188110352,
 "passed": false
}

The forward check passes within the same tolerance as the correct implementation. Training is still wrong. Across five seeds, the input-gradient error ranges from 0.00998 to 0.0182.

The test hierarchy now follows directly from the mechanism:

compare forward outputs to catch a missing row-parallel output reduction;
compare input and parameter gradients to catch missing backward communication;
compare updated parameters to confirm that the entire state transition matches.

Run the suite:

1

PYTHONPATH=code pytest -q code/test_tp_gqa_attention.py

1
2


..... [100%]
5 passed in 9.92s

The full test output is retained in data/terminal-05-tests.txt.

How this maps to Megatron Core

The reconstruction is deliberately written with ordinary PyTorch functions, but the objects map cleanly to Megatron’s tensor-parallel design:

Reconstruction object	Megatron-style object	Role	Important difference
Sliced Q/K/V weights	column-parallel projections	partition output features by head ownership	production code uses TP process groups and optimized kernels
Local RoPE and causal attention	rank-local attention heads	compute attention without immediate TP communication	production may use FlashAttention and packed sequences
Sliced output weight columns	row-parallel output projection	consume partitioned context	production performs an actual collective
Sum of output partials	reduce-from-TP region	reconstruct replicated hidden state	this fixture sums Python tensors in one process
Sum of input-gradient contributions	copy-to-TP backward reduction	recover gradient for replicated input	this fixture simulates omission by detaching one branch

The original Megatron-LM design partitions Q, K, and V so complete attention heads stay local, then row-partitions the output projection. NVIDIA’s current Megatron Core documentation still describes tensor parallelism as splitting individual layers, while the production implementation adds sequence parallelism, process-group management, communication overlap, Transformer Engine kernels, and hardware-specific collectives.

This post does not establish performance, communication cost, or correctness for Megatron Core itself. It does not run CUDA, NCCL, FlashAttention, sequence parallelism, context parallelism, mixed precision, fused QKV layouts, KV replication for awkward GQA ratios, dropout RNG tracking, or distributed optimizer state. Those are the next layers of the reconstruction.

What the reconstruction made visible

I started with the intuition that tensor parallelism partitions projection matrices. That is mechanically true but not precise enough for GQA. The reusable rule is semantic: keep each query/KV reuse group intact, let attention remain local, and place reductions at the boundaries where independently computed contributions must become one tensor.

The forward reduction and backward reduction are separate obligations. The second one is easy to miss because the model can produce the correct output while training with the wrong gradient. After this reconstruction, I would not approve a tensor-parallel attention change from forward parity alone.

References

Joshua Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, 2023.
Mohammad Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism”, 2019.
NVIDIA, Megatron Core parallelism strategies.
NVIDIA, Megatron-LM source repository.