LLM Lab Notes 04: Tensor Parallelism 101 — Column/Row Parallel and the Megatron FFN

Code files: llm_lab/03_tensor_parallel/tensor_parallel_demo.py, row_parallel_demo.py, megatron_ffn_demo.py

Tensor Parallelism (TP) solves "a single weight matrix doesn't fit": split a large matrix across GPUs, each holding a slice, and use collectives to reassemble the result. This note uses three minimal experiments to explain both cutting directions (Column / Row) and Megatron's classic FFN split.

1. Experiment Goal

1. Implement Column Parallel (split weights along columns): each GPU computes part of the output, AllGather reassembles the full output;
2. Implement Row Parallel (split weights along rows): each GPU computes a partial sum, AllReduce aggregates;
3. Reproduce Megatron's two-layer FFN (Column → GELU → Row) and explain why the two directions pair up.

2. Code Walkthrough

Experiment A: Column Parallel (tensor_parallel_demo.py)

world = dist.get_world_size()

hidden = 4096
chunk = hidden // world              # per-GPU column count

weight = torch.randn(hidden, chunk, device="cuda")   # ① column cut: hidden x chunk

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

local_out = x @ weight               # ② each GPU computes part of the output

gather_list = [torch.zeros_like(local_out) for _ in range(world)]
dist.all_gather(gather_list, local_out)              # ③ AllGather to concatenate columns

out = torch.cat(gather_list, dim=1)  # ④ full output [32, hidden*world]

• ① The weight is split into world column slices; each GPU holds [hidden, chunk];
• ② Each GPU computes only its output columns, shape [32, chunk];
• ③④ all_gather collects all GPUs' output columns and concatenates them along the feature dim into [32, hidden] — element-wise identical to the full-matrix result.

In one sentence: Column Parallel = concatenate outputs (AllGather).

Experiment B: Row Parallel (row_parallel_demo.py)

chunk = hidden // world

weight = torch.randn(chunk, hidden, device="cuda")   # ① row cut: chunk x hidden

x = torch.randn(32, chunk, device="cuda")            # ② input split by rows too

local_out = x @ weight               # ③ each GPU computes a partial sum

dist.all_reduce(local_out)           # ④ AllReduce aggregation

• ① The weight is split into world row slices; each GPU holds [chunk, hidden];
• ② The input x must be split the same way (each GPU holds different row fragments);
• ③④ Each GPU's local_out is a partial sum of the full result; all_reduce sums them into the complete output — and every GPU receives the same one.

In one sentence: Row Parallel = sum outputs (AllReduce).

Experiment C: Megatron FFN (megatron_ffn_demo.py)

chunk = ffn_hidden // world

w1 = torch.randn(hidden, chunk, device="cuda")   # Column: hidden x chunk
w2 = torch.randn(chunk, hidden, device="cuda")   # Row: chunk x hidden

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

local_mid = x @ w1                 # first layer: Column Parallel
local_mid = F.gelu(local_mid)      # activation, element-wise, no communication

local_out = local_mid @ w2         # second layer: Row Parallel
dist.all_reduce(local_out)         # aggregation

Key observation: with Column followed by Row, the intermediate AllGather is eliminated.

3. How It Works: Why Column → Row Needs No Intermediate Communication

This is one of Megatron's most elegant designs. Start with the naive decomposition:

x ──(Column w1)──> mid (needs AllGather to reassemble) ──(act)──> mid ──(Row w2)──> out (needs AllReduce)

If the Column layer's output is "logically reassembled," you need an AllGather. But if the next layer is exactly Row Parallel, the Row layer expects its input split by rows — and the Column layer's sharded output is exactly the Row layer's sharded input! The intermediate AllGather disappears entirely:

x ──(Column w1)──> mid_local (naturally sharded) ──(GELU element-wise)──> mid_local ──(Row w2)──> part ──(AllReduce)──> out

The two layers share a single communication: the whole FFN block needs only one AllReduce (at the Row layer's output) instead of two. Extended to the full Transformer: each Block's FFN communicates once, and Attention's QKV (Column) pairs with the output projection (Row) to also skip the intermediate communication — communication points across the whole model drop dramatically. This is key to why Megatron tensor parallelism actually runs fast.

(GELU being element-wise and communication-free is also part of why the intermediate sync can be dropped.)

4. Results

Running the three experiments on 4 GPUs, expected output (printed by rank 0):

# Experiment A: Column Parallel
Output shape: torch.Size([32, 4096])

# Experiment B: Row Parallel
Output shape: torch.Size([32, 4096])

# Experiment C: Megatron FFN
FFN output shape: torch.Size([32, 4096])

How to tell it succeeded: compare against the single-GPU full-matrix output — the distributed output must match (add a torch.allclose assertion). Correct shape is only the first step; numerical equivalence is the proof that communication and partition are correct.

5. Engineering Takeaways

1. The two cuts have different communication costs: Column uses AllGather (traffic = output size); Row uses AllReduce (traffic ≈ 2× output). So for communication-sensitive spots, prefer Column cuts (and backward needs no communication there); Row is used only where pairing demands it (FFN second layer, Attention output projection).

2. "Partition direction × neighbor layers" determines the number of communication points: when designing a parallel scheme, draw a "sharding state propagation" graph layer by layer — what sharding state the input is in, what state the next layer needs, and whether a collective is required in between. Megatron's essence is letting sharding states connect naturally, minimizing communication points.

3. TP's applicability boundary: TP is the most communication-dense strategy and communicates every step; it is typically used within a node (NVLink). Cross-node TP explodes inter-node traffic — which is the hardware reason TP degrees are usually 4–8 in 3D parallelism.

4. world_size must divide the dimension: this experiment's chunk = hidden // world assumes divisibility. Real code must check explicitly; otherwise misaligned slices cause silent wrong results — far harder to debug than an error.

Next: with the FFN split down, how about Attention? Note 05 splits QKV by head: LLM Lab Notes 05: Attention Tensor Parallelism.