TP：Column / Row Parallel Linear 怎么切？
MLP 和 attention 怎么切？

英文问法What is tensor parallelism? · Explain column-parallel vs row-parallel linear layers, pros and cons. · Which MLP matrix is split by columns, which by rows?

考过NVIDIA 框架岗 2023 电话面（"Column Parallel和Row Parallel Linear Layer的实现和pros & cons"，写少量 pseudocode，知乎）· NVIDIA DL Algorithm Engineer（Glassdoor RVW102989764，年份未核实）· NVIDIA High Performance LLM Training Engineer 电面（Glassdoor RVW105730158）· 华为机考（Megatron MLP 的正确切法）· 牛客多篇（"mlp中第一个矩阵和第二个矩阵分别是行切还是列切"）· 腾讯混元 AI Infra 一面 · OpenAI MLE onsite（PracHub，未核实）。手写练习见 ml-coding。

先弄清楚：每个概念是什么

Tensor parallelism (TP)
把每层的 weight matrix 切成 \(t\) 份放到 \(t\) 张卡上，每卡只存、只算自己那份。例：\(h = 8192\)，MLP 第一层 weight 是 [8192, 32768]，\(t = 8\) 时每卡只放 [8192, 4096]。
Column parallel
weight 按列切。每卡拿完整输入 \(X\)，算出输出的几列，拼起来就是完整输出。例：\(A\) 是 [8, 32]，\(t = 2\) → 每卡 \(A_i\) 是 [8, 16]，\(Y_i = XA_i\) 是 [n, 16]。
Row parallel
weight 按行切，输入跟着按列切。每卡算出完整形状的部分和，加起来才对。例：\(XB = X_1B_1 + X_2B_2\)，两项都是 [n, 8]。
All-reduce
每卡一个同形状 tensor，求和后每卡都拿到总和，详见 all-reduce。
All-reduce 与 concatenate 的区别
这里的 all-reduce 做逐元素求和，不是等待所有 GPU 后把结果拼长。例如 column backward 的两份 dX contribution 是 [1, 3] 和 [2, 4]，结果是 [3, 7]，不是 [1, 3, 2, 4]。all-gather 才是收集各卡的块，让每卡拿到完整集合（可沿切分维 concatenate）。基本 TP MLP 中 column forward 的输出保持分块，直接交给 row，无须真的拼起来。这里假设不开 sequence parallelism，row 输入已分块、上游 dY 已在各卡可用；row backward 各卡直接计算自己的 dX_i = dY·A_iᵀ，无须 TP communication。
f / g operator
Megatron 插在并行区两头的一对通信算子。f 只在 backward all-reduce，g 只在 forward all-reduce，另一个方向原样传。

看图：column / row 怎么切，MLP 怎么接

GPU0 的那份GPU1 的那份两卡各存一整份Column parallel：A 按列切X×A1A2=XA1XA2[bs, h][h, 2h][bs, 2h][h, 2h][bs, 2h]拼起来= XA [bs, 4h]forward 不通信；backward 时 dX 是部分和，要 all-reduce（f）Row parallel：B 按行切，输入 Y 也按列切Y1Y2[bs, 2h][bs, 2h]×B1B2[2h, h][2h, h]=两个部分和，各 [bs, h]Y1B1Y2B2all-reduce（g）：Z = Y1B1 + Y2B2Z：[bs, h]，两卡都拿到一整份forward 的输出是部分和，要 all-reduce（g）；backward 不通信
图 1：\(t = 2\)，每卡分到 \(4h/t = 2h\)。Column parallel 的每卡输出拼起来就是 \(XA\)。Row parallel 的每卡输出是 [\(bs\), \(h\)] 的部分和，all-reduce 相加后才对。
GPU0GPU1X [bs, h]：两卡各一整份f：forward 原样复制；backward 在这里 all-reduce dXX × A1 → [bs, 2h]X × A2 → [bs, 2h]GeLU（本卡算）→ Y1 [bs, 2h]GeLU（本卡算）→ Y2 [bs, 2h]Y1 × B1 → 部分和 [bs, h]Y2 × B2 → 部分和 [bs, h]g：all-reduce，Z = Y1B1 + Y2B2中间零通信
图 2：MLP 在两张卡上的数据流。括号里的三步都在本卡算，forward 只在最后的 g 做 1 次 all-reduce，虚线框 f 只在 backward 时 all-reduce dX。

30 秒回答

Tensor parallelism splits each weight matrix across t GPUs. Column-parallel splits the output columns: every GPU sees the full input and computes a slice of the output, so forward needs no communication, but the input gradient is a partial sum, all-reduced in backward. Row-parallel splits the input rows: each GPU computes a full-size partial sum, so forward needs an all-reduce. Megatron's MLP is column then row, since GeLU is elementwise and runs locally: one all-reduce forward, one backward. Attention splits by heads: QKV column-parallel, output projection row-parallel. Memory drops by t, but the all-reduces are on the critical path, so TP stays within a node.

推导 / 关键公式

\(X\) 是 [\(bs\), \(h\)]，MLP 是 \(Z = \mathrm{GeLU}(XA)\,B\)，\(A\) 是 [\(h\), \(4h\)]，\(B\) 是 [\(4h\), \(h\)]：

Column parallel（\(A\)）Row parallel（\(B\)）
每卡 weight\(A_i\)：[\(h\), \(4h/t\)]\(B_i\)：[\(4h/t\), \(h\)]
输入 → 输出完整 \(X\) → 一块列 \(XA_i\)一块列 \(Y_i\) → 部分和 \(Y_iB_i\)
forward 通信无all-reduce（g）
backward 通信dX 是部分和 → all-reduce（f）无

为什么通信位置正好相反：forward \(Y_j = \sum_k X_k A_{kj}\) 对输入维 \(k\) 求和，backward \(dX_k = \sum_j dY_j A_{kj}\) 对输出维 \(j\) 求和。哪一步的求和维度被切开了，哪一步就要 all-reduce。column 切 \(j\)，backward 的 dX 是 partial sum。row 切 \(k\)，forward 的输出是 partial sum。例：\(X = [1, 2]\)、\(A = [[1, 2], [3, 4]]\)、\(dY = [1, 1]\)，column 下两卡算出 [1, 3] 和 [2, 4]，相加才是 dX = [3, 7]。

MLP 为什么先 column 后 row

\[ \mathrm{GeLU}([XA_1, XA_2]) = [\mathrm{GeLU}(XA_1), \mathrm{GeLU}(XA_2)] \]

GeLU 逐元素，每卡对自己那几列直接算，结果正好是 row 要的输入块。反过来先按行切：

\[ \mathrm{GeLU}(X_1A_1 + X_2A_2) \ne \mathrm{GeLU}(X_1A_1) + \mathrm{GeLU}(X_2A_2) \]

GeLU 之前就得 all-reduce，数据从 \(h\) 宽变成 \(4h\) 宽。第二层 column 的输出接 residual 前还要一次 all-gather。通信量约 4.5 倍，省的 memory 也少。

Attention：按 head 切

\(W_Q, W_K, W_V\) 做 column parallel，每卡拿 \(a/t\) 个完整的 head，attention 在本卡算完；\(W_O\) 做 row parallel，最后 1 次 all-reduce。所以 \(a\) 要能被 \(t\) 整除。

代码（NumPy 模拟 t = 2）

rng = np.random.default_rng(0)
n, h, t = 4, 8, 2                     # n = b·s 个 token
X = rng.standard_normal((n, h))
A = rng.standard_normal((h, 4 * h))   # h → 4h
B = rng.standard_normal((4 * h, h))   # 4h → h
gelu = lambda x: x / (1 + np.exp(-1.702 * x))   # GeLU 近似
Z_ref = gelu(X @ A) @ B               # 单卡

A_cols = np.split(A, t, axis=1)       # column: 每卡 [h, 4h/t]
B_rows = np.split(B, t, axis=0)       # row:    每卡 [4h/t, h]
partial = [gelu(X @ Ai) @ Bi          # 每卡本地算，中间不通信
           for Ai, Bi in zip(A_cols, B_rows)]
print(np.allclose(sum(partial), Z_ref))     # g: all-reduce → True

# backward：column 的 dX 是 partial sum → f 做 all-reduce
dY = rng.standard_normal((n, 4 * h))
dX = sum(dYi @ Ai.T for dYi, Ai in zip(np.split(dY, t, 1), A_cols))
print(np.allclose(dX, dY @ A.T))            # True

# 反例：第一层按行切，GeLU 前必须先 all-reduce
X_cols, A_rows = np.split(X, t, 1), np.split(A, t, 0)
bad = sum(gelu(Xi @ Ai) for Xi, Ai in zip(X_cols, A_rows))
print(np.allclose(bad, gelu(X @ A)))        # False

必答点

追问

bias 怎么处理？

Column 的 bias 跟着列切。Row 的 bias 不切，all-reduce 之后加一次，否则会被加 \(t\) 次。

dropout 的随机数怎么处理？

TP 区域内的 dropout 每卡用不同 seed，否则各卡的 head 拿到同一个 mask。区域外用同一个 seed，各卡的副本才一样。

怎么验证你的实现是对的？

固定 seed，和单卡 nn.Linear 对比 output 和 gradient，\(t\) = 1、2、4、有无 bias 都测。没有多卡就用 gloo backend 起多进程。

来源Shoeybi et al. 2019, Megatron-LM §3（pp.4–5，eq. 1–3，Fig. 3–4）、Appendix B.2（pp.11–12，dropout RNG）· Narayanan et al. 2021 §2.3（p.4）· Megatron Core layers.py（ColumnParallelLinear、RowParallelLinear：row 的 bias 不切、reduce 之后再加）