多卡张量：拆开计算，再合回同一结果

用两张逻辑卡计算一个两层 MLP。每卡保留哪些矩阵、产生什么局部结果、何时求和，前向与反向都能和未分片基线核对。

01 / COLUMN THEN ROW

先切中间维度，就能保留局部激活

\begin{aligned}W_1&=[W_{1a}\;W_{1b}],\quad W_2=\begin{bmatrix}W_{2a}\\W_{2b}\end{bmatrix},\\H_a&=\phi(XW_{1a}),\quad H_b=\phi(XW_{1b}),\\Y&=H_aW_{2a}+H_bW_{2b}.\end{aligned}

真实小矩阵计算 · 无通信时延模拟

查看阶段

X=[[1,2],[3,−1]]；固定上游 G=[[0.2,−0.4],[0.6,0.1]]。反向目标可写为 L=ΣY⊙G。

Rank	本地 W₁	H 或 GZ	局部 Y 或 GX

未分片参考：

局部结果求和：

最大绝对误差：

02 / BACKWARD

\begin{aligned}G_{H_r}&=G_YW_{2r}^\top,\quad G_{W_{2r}}=H_r^\top G_Y,\\G_{Z_r}&=G_{H_r}\odot(1-H_r^2),\quad G_{W_{1r}}=X^\top G_{Z_r},\\G_X&=\sum_r G_{Z_r}W_{1r}^\top.\end{aligned}

逻辑输出 Y 在两卡可复制保存。反向从这一个逻辑输出的 GY 出发，每片使用同一个上游梯度，不能先把复制的 GY 再加一次。数据并行平均和张量并行合并属于不同轴，系数需要分别推导。

本例所有权重和输入的梯度都可通过有限差分验证。实际训练还需在相同全局样本、随机状态和精度下比较 loss、梯度及一次 optimizer step；浮点规约顺序允许按约定容差检查。

03 / COLLECTIVES

集合操作

R0=[1,2,3,4]，R1=[10,20,30,40]。AllToAll 将前半片发 R0、后半片发 R1；输出按来源 rank 排列。

Rank	输入	输出

ReduceScatter(sum) 后接 AllGather，在匹配分片与 rank 顺序时能得到 AllReduce(sum) 的结果。集合操作要求各 rank 对调用次序、形状、dtype 和通信组达成一致；某卡提前退出或顺序不同，可能等待甚至死锁。

逻辑算子定义不决定实际链路字节数。Ring、Tree、节点拓扑和分层规约会改变消息数量与路径；量化性能时需注明算法、参与 rank、发送/接收口径和通信是否被计算重叠。

04 / VOCABULARY SHARDS

\begin{aligned}m_r&=\max_{j\in V_r}z_j,\quad m=\max_r m_r,\\s_r&=\sum_{j\in V_r}e^{z_j-m},\quad s=\sum_rs_r,\\t_r&=\mathbf1[y\in V_r](z_y-m),\quad L=\log s-\sum_rt_r,\\\frac{\partial L}{\partial z_j}&=e^{z_j-m}/s-\mathbf1[j=y].\end{aligned}

单个监督位置 · 无 label smoothing

词表分片数目标 token ID全部 logits 同时平移

原始 logits=[1.2,−0.5,2,0.1,1.4,−1]。平移不改变数学上的 loss 和概率。

全局最大值

全局指数和

Loss

目标所在 rank：；与未分片参考的 loss 绝对差：

Rank	词表范围	Logits	局部最大	局部指数和	目标平移值

各词表位置的概率与梯度

Token	所属 rank	概率	梯度

输出头按词表列分片后，局部 logits 与梯度留在对应 rank；共享 hidden 的梯度还需要跨这些输出分片求和。词表维度的合并与 DP 样本维度的平均各有作用，不能重复除以并行数。

本例展示简单 CE。真实实现还需处理词表 padding、ignore_index、label smoothing、融合 kernel 和不同精度。所有边界应与完整词表参考的 loss、梯度和归一化规则核对。

两 rank MLP 与六类词表 CE 使用固定数值，输出和梯度与未分片参考核对。集合通信显示逻辑数据排列；不运行通信库或估算网络耗时。

公式与机制的原始依据 · 4 项

NVIDIA · Megatron-LM (2019)2026-10-29 核对作者说明的 Figure 2a：列/行分片间逐元素激活留在本地，MLP 前向归约输出、反向归约输入梯度；不是每个模块前后均同步
NVIDIA NCCL · Collective Operations集合通信的 rank、数据排列、求和与分片语义
PyTorch · Autograd mechanics计算图、保存中间值与自动微分规则
Megatron Core · VocabParallelCrossEntropy局部 logits、全局最大值和指数和、目标 token 与梯度