Vocab parallelism：embedding 和 LM head 怎么切？
cross-entropy 怎么算才不用 gather logits？

英文问法How do you shard a huge embedding table and the LM head? · Why split along the vocabulary, and where does communication happen? · How do you compute cross-entropy without gathering the full logits?

考过Microsoft MAI infra onsite 2026（整轮围绕大 embedding 和 LM head 的 distributed 切分，一直追问"为什么这么切、通信在哪里发生"，Pass，1p3a 1189843）· 快手（"这个列划分并行是什么？这里没用到all-gather啊？"）

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

Vocab-parallel embedding
[\(V\), \(h\)] 的 embedding 表按行（vocab）切，每卡管一段连续的 id。token 发给所有卡，只有管那段的卡输出那一行，其他卡输出 0，all-reduce 求和（图 1）。
LM head
[\(h\), \(V\)] 的矩阵，把 hidden 投成 \(V\) 个 logits。按 vocab 切就是 column parallel，每卡算 [\(b, s, V/t\)] 的 logits。
Vocab-parallel cross-entropy
loss 只需要每个 token 的三个标量：logits 的 max、\(\sum e^{z}\)、目标 logit。每卡在自己的 \(V/t\) 列上算局部值再 all-reduce，不用把 \(V\) 个 logits 拼回来。

看图：embedding 怎么查，cross-entropy 怎么算

token id = 5：两张卡都收到GPU0：E 的第 0–3 行0E[0]1E[1]2E[2]3E[3][V/t, h] = [4, h]5 不在本卡：查 0 号行后置 0GPU1：E 的第 4–7 行4E[4]5E[5]6E[6]7E[7][V/t, h] = [4, h]5 在本卡：本地第 5 − 4 = 1 行输出：0 向量（h 维）输出：E[5]（h 维）all-reduce 求和：0 + E[5] = E[5]，两张卡都拿到
图 1：\(V = 8\)、\(t = 2\)，每卡存 [\(V/t\), \(h\)] 的一段行。token 5 发给所有卡，只有管 4–7 的 GPU1 查到 \(E[5]\)，GPU0 输出 0，all-reduce 求和后正好是 \(E[5]\)。
rank 0rank 1logits：id 0–3 的列 [b, s, V/t]logits：id 4–7 的列 [b, s, V/t]局部 max：m0局部 max：m1① all-reduce MAX：m = max(m0, m1)S0 = Σ exp(z − m)，本卡的列S1 = Σ exp(z − m)，本卡的列② all-reduce SUM：S = S0 + S1目标 y 不在本卡：填 0目标 y 在本卡：取 zy③ all-reduce SUM：zy = 0 + zyloss = log S − (zy − m)：每张卡都能算
图 2：目标 \(y\) 落在 rank 1。三次 all-reduce 依次求全局 max \(m\)、\(\sum e^{z-m}\) 和目标 logit，每次每个 token 只传一个数，共 [\(b, s\)] 个。

30 秒回答

Megatron splits both along the vocabulary. The embedding is split by rows: each rank owns a range of token ids, looks up those it owns, outputs zeros for the rest, and an all-reduce sums the results, since exactly one rank contributes each token. The LM head is column-parallel over the same split, so each rank gets logits for its V/t slice. Cross-entropy runs without gathering the b·s·V logits: all-reduce MAX of the local max, all-reduce SUM of the local sum of exps, and all-reduce SUM of the target logit, which only one rank has. That's three b·s tensors instead of b·s·V, and logits memory drops by t.

推导 / 关键公式

Embedding

非本卡的 id 先映射到 0 号行去查，查完把输出置 0，再 all-reduce（开 SP 时换成 reduce-scatter）。

Cross-entropy：减 max 的 log-sum-exp 拆到各卡

\[ \ell = \log\sum_j e^{z_j - m} - (z_y - m) \] \[ m = \max_r m_r \quad \text{all-reduce MAX} \] \[ \textstyle\sum_j e^{z_j - m} = \sum_r S_r \quad \text{all-reduce SUM} \] \[ z_y = \sum_r z_y^{(r)} \quad \text{all-reduce SUM，非本卡填 0} \]

backward 全在本地：\(\partial\ell/\partial z_j = \mathrm{softmax}_j - \mathbb{1}[j = y]\)，只有拥有 \(y\) 的那张卡减 1。LM head 的 dX 再按 column parallel 的规矩 all-reduce。

代码（NumPy 模拟 t = 4）

rng = np.random.default_rng(0)
n, V, t = 6, 16, 4                    # n = b·s 个 token
z = rng.standard_normal((n, V)) * 5   # logits
y = rng.integers(0, V, n)             # target id
i = np.arange(n)
zs = z - z.max(1, keepdims=True)      # 单卡 reference
ref = np.log(np.exp(zs).sum(1)) - zs[i, y]

shards = np.split(z, t, axis=1)       # rank r 管 V/t 个词
m = np.max([zr.max(1) for zr in shards], 0)     # all-reduce MAX
S = sum(np.exp(zr - m[:, None]).sum(1)          # all-reduce SUM
        for zr in shards)
zy = 0
for r, zr in enumerate(shards):
    lo = r * V // t
    own = (y >= lo) & (y < lo + V // t)       # target 在本卡吗
    loc = np.where(own, y - lo, 0)            # 不在：index 置 0
    zy = zy + np.where(own, zr[i, loc], 0)    # all-reduce SUM
print(np.allclose(np.log(S) - (zy - m), ref))  # True

手算例子：省多少

\(b = 1\)、\(s = 4096\)、\(h = 8192\)、\(V = 131072\)（128K）、\(t = 8\)：

每卡不切按 vocab 切
embedding weight + Adam state（16 B/param，见 optimizer）16 GiB2 GiB
fp32 logits [\(b, s, V\)]2 GiB0.25 GiB

gather logits 要搬 \(bsV\) ≈ 5.4 亿个元素，vocab-parallel CE 只 all-reduce \(3 \times bs\) 个 fp32，共 48 KiB。

必答点

追问

为什么不按 hidden 切？

Embedding 按 hidden 切，每卡只有部分维度，后面的层要完整 \(h\)，得 all-gather。LM head 按 hidden 切就是 row parallel，要 all-reduce \(bsV\) 个元素的部分和，太大。

还有别的切法吗？

按 token id 把查询路由到拥有那段的卡，查完再 all-to-all 发回。只搬用到的行，但高频 token 集中的卡会 load imbalance。

weight tying 碰上 pipeline parallel 怎么办？

embedding 在第一个 stage、LM head 在最后一个 stage，各存一份。初始化时 all-reduce 让两份相同，每步再 all-reduce 两份的 gradient。

为什么 vocab 要 padding？

让每卡的 \(V/t\) 是 128 的倍数，logits GEMM 更高效。Megatron-LM 把 GPT-2 的 50257 补到 51200（\(128 \times 8\) 的倍数）。

来源Shoeybi et al. 2019 §3（pp.4–5：embedding 按 vocab 切 + all-reduce；output GEMM 和 CE 融合，通信从 \(b \times s \times v\) 降到 \(b \times s\)）、§5.1（p.6，padding 到 51200）· Korthikanti et al. 2022 §4.3（p.8，logits 用 fp32，\(4sbv/t\)）· Megatron Core cross_entropy.py（MAX / SUM 三次 all-reduce）、layers.py（VocabParallelEmbedding）、language_module.py（PP 下 tied embedding）