Parallelism Strategies #1 - Tensor Parallelism
make big models fit on tiny devices ....
Quick note: While the term tensor parallelism is quite overloaded in general this post refers to the implementation advocated in the Megatron-LM paper where Shoeybi et al proposed a clever partitioning scheme for the weights such that overall communication cost is reduced.
Imagine you have access to 2 connected GPU devices with 8GB VRAM memory in each but you want to run inference on a 7B model which roughly occupies 7*4bytes = 28 GB memory in fp32 precision. Let’s say you put the model in bf16 precision cutting down memory to 14 GB you still can’t fit the entire model on one of your devices. The natural next step you try is split the 2D weight matrices on either the row or column dimension but if you naively do this for all matrices we’d have to do all-gather before forward pass + all-gather before backward pass and reduce_scatter gradients, that’s a LOT of communication (3 collectives per matrix) but that’s basically how FSDP does it and it works out within a node where NVLink has fast transfer speeds.
What if there’s a smarter way to do this eliminating all-gathers entirely and just relying on 2 all-reduce collectives ? that would save a lot of communication costs and that’s basically what Megatron style sharding promises.
Consider a simple operation to simulate a transformer’s FFN block as Y = O @ W_2.t() where O = X @ W_1.t(), now let’s work on the underpinning intuitions listed below:
If we split W1 row-wise and W2 column-wise we can eliminate all-gathers entirely to compute final correct output.
Forward pass output requires all-reduce to get correct result on each device.
Backward pass needs all-reduce to pass on X’s gradients to previous layer.
Now let’ see why (1) and (2) hold with a sketched out example from the matmul backprop post where we shard the weight matrix row-wise resulting in an output that is sharded column-wise, note that PyTorch uses weight transposed so we do the same:
Given Y = (X @ W_1.t()) @ W_2.t(), let’s explore the case where W2 is sharded column-wise , what does such partitioning even achieve ? If we look at the correct final output it’s composed of terms that can be computed using what we have locally on each device … this naturally makes it so that we can compute partial outputs on each device then all-reduce to get correct final output on each device as seen below.
To see this detail even more clearly look at below decomposition :
Now for the backward pass let’s first work through all_reduce gradient which is as gradient of a simple sum. If c = a + b then dL/da = dL/db = dL/dc , so that would allow us to compute correct gradient of Y_partial as dL/dY. From matmul backprop post we know that compute
dL/dO_shard = dL/dY_partial @ W_2_shard
dL/dW_2 = dL/dY_partial.t() @ O_shard
Now to backprop further we have O_shard = X @ W_1_shard.t(), where O_shard is column-sharded so we have
dL/dW_1_shard = (dL/dO_shard).t() @ X
dL/dX_partial = dL/dO_shard
dL/dX = all_reduce(dL/dX_partial)
We need the last all_reduce as we only have access to parts of the full gradient given we shards of O on each device and not the entire contents of O. This brings us to the operators introduced in Megatron paper to make this whole concept simpler :
Here f acts as no-op in forward pass but acts as all-reduce in backward pass, g is exactly the opposite … acting as all-reduce in forward pass and no-op in backward pass. Now let’s implement the forward pass in code and verify that our sharded forward pass output matches unsharded final output :
"""
╔══════════════════════════════════════════════════════════════════════════════╗
║ MEGATRON-LM: TENSOR PARALLELISM FOR FFN BLOCK ║
║ (Column + Row Parallel Linear Layers) ║
╚══════════════════════════════════════════════════════════════════════════════╝
The FFN block computes: Y = ReLU(X @ W1ᵀ) @ W2ᵀ
(simplified, no ReLU here for clean gradient verification)
UNSHARDED (single GPU):
┌───────────┐ W1 (2×3) ┌───────────┐ W2 (3×2) ┌───────────┐
│ X (3×3) │ ─────────────▶ │ O (3×2) │ ─────────────▶ │ Y (3×3) │
└───────────┘ └───────────┘ └───────────┘
SHARDED ACROSS 2 GPUs (Megatron style):
┌──────────────────────────────────────────────────────────┐
│ KEY IDEA: Shard W1 by ROWS, Shard W2 by COLUMNS │
│ → Each GPU handles half the "hidden" dimension │
│ → Final outputs are summed via All-Reduce │
└──────────────────────────────────────────────────────────┘
GPU 0: X @ W1_shard0ᵀ → O_shard0 → O_shard0 @ W2_shard0ᵀ → Y_partial0
GPU 1: X @ W1_shard1ᵀ → O_shard1 → O_shard1 @ W2_shard1ᵀ → Y_partial1
↓
Y = Y_partial0 + Y_partial1
(All-Reduce / sum)
"""
import torch
torch.manual_seed(42)
# ═══════════════════════════════════════════════════════════════
# SETUP: Input and Weight Matrices
# ═══════════════════════════════════════════════════════════════
#
# X : (3×3) — sequence of 3 tokens, each with 3 features
# W1 : (2×3) — projects 3 → 2 (hidden dim = 2)
# W2 : (3×2) — projects 2 → 3 (output dim = 3)
#
# ┌─────────────────┐
# │ X (3 × 3) │ ← 3 tokens × 3 input features
# └─────────────────┘
# │
# │ @ W1ᵀ (3×2) W1 (2×3):
# │ ┌──────────────┐
# ▼ │ row 0 (1×3) │ ← W1_shard0 → GPU 0
# ┌─────────────────┐ │ row 1 (1×3) │ ← W1_shard1 → GPU 1
# │ O (3 × 2) │ └──────────────┘
# └─────────────────┘ ↑ Row-parallel sharding
# │
# │ @ W2ᵀ (2×3) W2 (3×2):
# │ ┌──────────────┐
# ▼ │ col 0 (3×1) │ ← W2_shard0 → GPU 0
# ┌─────────────────┐ │ col 1 (3×1) │ ← W2_shard1 → GPU 1
# │ Y (3 × 3) │ └──────────────┘
# └─────────────────┘ ↑ Column-parallel sharding
X = torch.rand(3, 3)
W1 = torch.rand(2, 3)
W2 = torch.rand(3, 2)
X.requires_grad = True
W1.requires_grad = True
W2.requires_grad = True
# ═══════════════════════════════════════════════════════════════
# UNSHARDED FORWARD PASS (ground truth on a single GPU)
# ═══════════════════════════════════════════════════════════════
O = X @ W1.t() # 3x2
Y = O @ W2.t() # 3x3
W1.retain_grad()
W2.retain_grad()
# ═══════════════════════════════════════════════════════════════
# SHARDED FORWARD PASS
# ═══════════════════════════════════════════════════════════════
#
# STEP 1: Shard W1 along ROWS (row-parallel, dim=0)
# ────────────────────────────────────────────────────
#
# W1 (2×3): W1_shard0 (1×3): W1_shard1 (1×3):
# ┌───────────┐ ┌───────────┐ ┌───────────┐
# │ row 0 │ → │ row 0 │ GPU 0 │ row 1 │ GPU 1
# │ row 1 │ └───────────┘ └───────────┘
# └───────────┘
#
# Each GPU independently computes its slice of the hidden dim:
# GPU 0: O_shard0 = X @ W1_shard0ᵀ → (3×1)
# GPU 1: O_shard1 = X @ W1_shard1ᵀ → (3×1)
#
# Together they reconstruct O = [O_shard0 | O_shard1] → (3×2)
# (No communication needed here! Each GPU uses the full X)
W1_shard0, W1_shard1 = W1.chunk(chunks=2, dim=0) # 1x3 , RowParallel
O_shard0 = X @ W1_shard0.t() # 3x1 = 3x3 @ 3x1
O_shard1 = X @ W1_shard1.t() # 3x1 = 3x3 @ 3x1
W2_shard0, W2_shard1 = W2.chunk(chunks=2, dim=1) # 3x1 , ColParallel
Y_partial0 = O_shard0 @ W2_shard0.t() # 3x3 = 3x1 @ 1x3
Y_partial1 = O_shard1 @ W2_shard1.t()
Y_all_reduced = Y_partial0 + Y_partial1
# Verify if final all_reduced Y matches unsharded Y.
torch.testing.assert_close(Y_all_reduced, Y)
# Verify if intermediate output O shards all-gathered matches unsharded O.
O_sharded = torch.concat([O_shard0, O_shard1], dim=1)
torch.testing.assert_close(O_sharded, O)Now let’s implement the sharded backward pass and verify that it matches torch’s autograd computations for unsharded input.
# ═══════════════════════════════════════════════════════════════
# UNSHARDED BACKWARD PASS (ground truth)
# ═══════════════════════════════════════════════════════════════
#
# Loss = mean(Y) → dL/dY = ones(3,3) / 9
#
# dL/dO = dL/dY @ W2 (chain rule through W2)
# dL/dX = dL/dO @ W1 (chain rule through W1)
# dL/dW1 = dL/dO.T @ X
# dL/dW2 = dL/dY.T @ O
# Y_all_reduced.retain_grad()
loss = Y.mean()
loss.backward()
# ═══════════════════════════════════════════════════════════════
# SHARDED BACKWARD PASS
# ═══════════════════════════════════════════════════════════════
#
# Forward recap (for gradient derivation):
# ┌──────────────────────────────────────────────────────┐
# │ GPU g: Y_partial_g = O_shard_g @ W2_shard_gᵀ │
# │ Y = Σ_g Y_partial_g (all-reduce) │
# │ loss = mean(Y) │
# └──────────────────────────────────────────────────────┘
#
# Since Y = Y_partial0 + Y_partial1 and loss = mean(Y):
# dL/dY_partial_g = dL/dY = ones(3,3)/9 ← same on both GPUs
# (broadcast, not reduce)
dL_by_dY_all_reduced = (torch.ones(3, 3)) / 9
dL_by_dY_partial0 = dL_by_dY_all_reduced # 3x3
dL_by_dY_partial1 = dL_by_dY_all_reduced # 3x3
dL_by_dO_shard0 = dL_by_dY_partial0 @ W2_shard0 # 3x1 = 3x3 @ 3x1
dL_by_dO_shard1 = dL_by_dY_partial1 @ W2_shard1 # 3x1 = 3x3 @ 3x1
# ─────────────────────────────────────────────────────────────
# Gradient w.r.t. X (needs All-Reduce!)
# ─────────────────────────────────────────────────────────────
#
# dL/dX_partial_g = dL/dO_shard_g @ W1_shard_g
#
# GPU 0: (3×1) @ (1×3) → (3×3)
# GPU 1: (3×1) @ (1×3) → (3×3)
#
# ┌──────────────────┐ ┌──────────────────┐
# │ dL/dX_partial0 │ + │ dL/dX_partial1 │
# │ (3×3) │ │ (3×3) │
# └──────────────────┘ └──────────────────┘
# ↑ All-Reduce (sum) → dL/dX (3×3)
#
# WHY? Because X contributed to O_shard0 AND O_shard1 via
# different weight shards, so its total gradient is the sum
# of contributions from all GPUs.
dL_by_dX_partial0 = dL_by_dO_shard0 @ W1_shard0 # 3x1 @ 1x3
dL_by_dX_partial1 = dL_by_dO_shard1 @ W1_shard1 # 3x1 @ 1x3
# all-reduce to get correct gradient
dL_by_dX = dL_by_dX_partial0 + dL_by_dX_partial1
torch.testing.assert_close(dL_by_dX, X.grad)
# gradients for weight shards required for optimizer update
# ─────────────────────────────────────────────────────────────
# Gradient w.r.t. weight shards (local, no communication!)
# ─────────────────────────────────────────────────────────────
#
# Each GPU only needs its local activations to update its weight shard.
# This is the beauty of Megatron sharding: weight gradients are local!
#
# dL/dW2_shard_g = dL/dY_partial_gᵀ @ O_shard_g
# GPU 0: (3×3)ᵀ @ (3×1) → (3×1) ← gradient for W2_shard0 on GPU 0
# GPU 1: (3×3)ᵀ @ (3×1) → (3×1) ← gradient for W2_shard1 on GPU 1
#
# dL/dW1_shard_g = dL/dO_shard_gᵀ @ X
# GPU 0: (3×1)ᵀ @ (3×3) → (1×3) ← gradient for W1_shard0 on GPU 0
# GPU 1: (3×1)ᵀ @ (3×3) → (1×3) ← gradient for W1_shard1 on GPU 1
dL_by_dW2_shard0 = dL_by_dY_partial0.t() @ O_shard0 # 3x1 = 3x3 @ 3x1
dL_by_dW2_shard1 = dL_by_dY_partial1.t() @ O_shard1 # 3x1 = 3x3 @ 3x1
dL_by_dW1_shard0 = dL_by_dO_shard0.t() @ X # 1x3 @ 3x3
dL_by_dW1_shard1 = dL_by_dO_shard1.t() @ X # 1x3 @ 3x3
torch.testing.assert_close(W1.grad, torch.vstack([dL_by_dW1_shard0, dL_by_dW1_shard1])) # row sharded W1
torch.testing.assert_close(W2.grad, torch.hstack([dL_by_dW2_shard0, dL_by_dW2_shard1])) # row sharded W1