Tensor-Parallel MLP (Megatron)

The problem

When a layer is too large for one GPU (graphics processing unit), its matrices can be cut into pieces, one per GPU. Megatron-LM does this for the MLP (multilayer perceptron) of a transformer,

with a single communication at the end. Here $X$ holds the token vectors, $A$ of shape $(D, F)$ widens them to $F$ numbers, GeLU (Gaussian error linear unit) is the nonlinearity, applied entry by entry, and $B$ of shape $(F, D)$ narrows them back. You will simulate the split on $N$ GPUs and check that it gives the same result.

The split

• $A$ is split by columns into $A_1, \dots, A_N$. GPU $i$ computes $H_i = \mathrm{GeLU}(XA_i)$ on its own: each column of $XA$ needs only one column of $A$, and GeLU works entry by entry, so no other GPU is needed.
• $B$ is split by rows into the matching $B_1, \dots, B_N$: the rows of $B$ that multiply the columns of $H$ this GPU computed. GPU $i$ computes the partial output $Y_i = H_i B_i$, already of the full output shape.
• One all-reduce adds up the partial outputs: $Y = \sum_i Y_i$.

Example. One token $X = (1, 2)$, $D = 2$, $F = 4$, $N = 2$ GPUs, with $A$ = [[1, 0, -1, 2], [0, 1, 1, -1]] and $B$ = [[1, 0], [0, 1], [1, 1], [2, 0]].

• Unsplit: $XA = (1, 2, 1, 0)$, $\mathrm{GeLU}(XA) \approx (0.841, 1.955, 0.841, 0)$, and $Y \approx (1.682, 2.796)$.
• GPU 0 has columns 0–1 of $A$ and rows 0–1 of $B$: $H_0 \approx (0.841, 1.955)$, $Y_0 \approx (0.841, 1.955)$.
• GPU 1 has columns 2–3 of $A$ and rows 2–3 of $B$: $H_1 \approx (0.841, 0)$, $Y_1 \approx (0.841, 0.841)$.
• All-reduce: $Y_0 + Y_1 \approx (1.682, 2.796)$, the same as unsplit.

Why not split $A$ by rows? Then each GPU would hold part of a sum that must be completed before the GeLU, and GeLU of a sum is not the sum of the GeLUs: $\mathrm{GeLU}(1 + 1) \approx 1.955$, but $\mathrm{GeLU}(1) + \mathrm{GeLU}(1) \approx 1.682$. That would need a second communication.

GeLU

GPT-2 (Generative Pre-trained Transformer 2) uses this approximation of GeLU:

where $x$ is one entry and $u$ is a helper value. For example $\mathrm{GeLU}(1) \approx 0.841$ and $\mathrm{GeLU}(0) = 0$.

What to implement

With $F$ divisible by $N$:

• gelu(x): the approximation above, elementwise.
• shard_mlp(A, B, n): return the list of $n$ pairs (A_i, B_i).
• tp_mlp(X, shards): compute each GPU's partial output from its own (A_i, B_i) only, then sum them (the simulated all-reduce).

Requirements

• GPU $i$ gets columns $iF/n$ to $(i+1)F/n - 1$ of $A$ and the same range of rows of $B$.
• The result must equal the unsplit $\mathrm{GeLU}(XA)\,B$.