行并行张量并行原理

可视化演示

展示SQL查询从解析到执行的计划树构建过程，对比不同访问路径和连接策略的代价估算，动态演示优化器选择最优执行计划的全流程。

行并行概念

图书馆管理员发现，与其把整本书交给一个抄书员，不如按章节把书拆开。每个抄书员拿到属于自己的那几章，抄完后把内容重新拼成一本书。这个“按章节拆书”的动作，就是行并行的核心直觉——矩阵乘法中的权重矩阵是按“行”组织的，每一行都独立参与计算，所以可以像撕书一样把矩阵从中间撕开。

在深度学习中，一个线性层做的事情就是 Y = XW。其中 W 是权重矩阵，X 是输入。行并行的做法是：把 W 从中间水平切一刀，上半部分放到 GPU0，下半部分放到 GPU1。于是 GPU0 只需要计算 X 与 W 上半部分的乘积，GPU1 计算 X 与 W 下半部分的乘积。两边的计算结果合在一起，就是完整的输出。

这个方案的美妙之处在于“互不干扰”——因为矩阵乘法中，输出 Y 的第 i 行只依赖于 W 的第 i 行（严格说是 Y[i,:] = X × W[i,:] 的转置）。切分 W 的行，就像把一条生产线上的工序按产品型号分开，每个工厂生产自己负责的那几种型号，互不依赖。

数学原理

为什么行并行在数学上是成立的？因为我们利用了矩阵乘法的分块性质。具体来说，如果把权重矩阵 W 按行分成上下两块：

这意味着，输出矩阵 Y 的上半部分是 X 与 W 上半行的乘积，Y 的下半部分是 X 与 W 下半行的乘积。两块计算完全独立，可以分给不同的 GPU 同时执行。

回到图书馆的比喻：把书按章节撕开后，每个抄书员手里的内容合起来就是完整的书，章节的顺序没有变，内容没有增删，只是写的人多了。行并行就是这样——数学保证了“拼回去”和“没撕开”的结果一模一样。

前向传播流程

前向传播是行并行最直观的舞台。假设我们有 2 张 GPU（叫 TP=2），一个权重矩阵 W 被切成上下两半，分别放在 GPU0 和 GPU1 上。现在的任务是计算 Y = XW。

第一步，把输入 X 广播给所有 GPU。这个过程就像图书馆管理员把原始书籍的复印本发给每位抄书员——所有人都拿到完整的 X。第二步，GPU0 和 GPU1 各自用自己手里的那半块 W，与完整的 X 做矩阵乘法，算出各自的“部分输出”。第三步，把部分输出沿行方向拼接起来（这个过程对应集合通信中的 All-gather），得到完整的 Y。

这个过程就像四个抄书员各自抄完自己的章节后，把稿纸交回给主编，主编按章节号把稿纸钉成一本完整的书。前向传播的结束标志，就是这本完整的书（输出 Y）被送到了下一层。

反向传播流程

反向传播是训练中“往回传误差”的阶段。这时候，图书馆的抄书员们需要根据书评（梯度）来修改自己抄写的内容（权重）。行并行的反向传播稍微绕一点，因为梯度需要“沿着来路返回”。

关键直觉还是那句话：行切分让前向输出是“拼接”出来的，所以反向传播时梯度要先“拆开”。具体来说，输出 Y 的梯度 dY 被拆成上下两块，各自回到对应的 GPU。每张卡用自己收到的梯度块 dY 和自己的权重切片 W，计算出对输入 X 的梯度 dX；同时用 dY 和缓存的前向输入 X 计算自己那部分权重 W 的梯度 dW。最后，所有卡对 dX 做一次 All-reduce（把大家算出的 dX 梯度求和），因为完整的 dX 是所有卡贡献之和。

回到抄书的比喻：主编发现书的第二章有错（梯度），就通知负责第二章的抄书员修改。每个抄书员只修改自己负责的章节，不需要管别人的部分。但有一点特殊——如果某个观点（输入特征）被多个章节引用，那么这些章节的修改意见都要汇总起来，才能决定这个观点要怎么改。这就是 All-reduce 在做的事。

通信机制

图书馆里有一个专门的“传话员”系统，负责在抄书员之间传递稿件。在行并行中，GPU 之间的通信也有一套标准化的“传话”流程，这就是集合通信原语。

行并行涉及三种核心通信原语。第一种是 Broadcast（广播）：把输入 X 从一张卡复制到所有卡，就像主编把原书复印发给大家。第二种是 All-gather（全收集）：每张卡把自己算好的部分输出发给所有其他卡，最后每张卡都拥有完整的输出 Y，就像抄书员写完后互相传阅，让每个人都拿到完整的书。第三种是 All-reduce（全归约）：每张卡把自己的局部梯度贡献出来，先“求和”再分发回所有卡，就像每个抄书员提出修改意见后，主编汇总所有人的意见，再把汇总结果通知所有人。

通信时机也很有讲究。前向传播中，输入 X 广播是一次通信，输出 Y 的 All-gather 是一次通信。反向传播中，dY 需要类似切分的操作（实际上不需要显式通信，因为 Y 被 All-gather 过，梯度按切片逻辑本地分割即可），然后 dX 的 All-reduce 是一次通信。所以每个 Transformer 层大约有两次通信。

1. 广播 Broadcast：假设数据在 GPU0 上，需要被复制到所有卡

2. 全收集 All-gather：每张卡持有部分输出 $Y^{(i)}$，通信后每张卡都持有完整 $Y$

3. 全归约 All-reduce：每张卡持有部分梯度贡献，通信后每张卡都持有梯度总和

通信就像图书馆的“传话员”系统。广播是群发消息，全收集是收齐大家的部分拼成完整版再群发回去，全归约是把大家意见汇总后再通知所有人。传话员虽然高效，但毕竟要花时间跑腿，这就是通信开销的来源。

与列并行配合

行并行不是孤军奋战的，它有一个好搭档——列并行。如果说行并行是把书“横着撕”，列并行就是把书“竖着切”。两者交替使用，构成了 Megatron 张量并行的核心循环。

为什么要交替使用？关键在于避免不必要的通信。考虑一个 Transformer 层：前一个线性层的输出马上要作为后一个线性层的输入。如果两层都用行并行，那么第一层输出的 All-gather 结果会被第二层再次广播——这中间多了一次完整的通信往返。但如果第一层用行并行、第二层用列并行，情况就不同了：行并行的输出（各卡持有部分行）恰好就是列并行需要的输入分布，不需要重排数据。

图书馆的比喻在这里同样贴切：如果一个抄写流程是“按章节分给抄书员，抄完后的章节正好是下一个环节”排版“环节需要的版面布局”，那就不需要重新分发稿件，省去了中间的传阅时间。行并行和列并行就是这样默契配合，让数据“无缝衔接”。

通信开销分析

通信是有代价的。抄书员之间传一次稿子需要时间，GPU 之间搬运数据也需要时间。理解通信开销，就是要回答一个问题：行并行到底能省多少时间，又额外花了多少时间在通信上？

对于一个形状为 $(d_{out} \times d_{in})$ 的权重矩阵，行并行把计算量均分到 TP 张卡上，每张卡只需要做原来的 $1/TP$ 计算量。但通信量呢？前向传播中，输入 X 需要广播（通信量正比于输入大小 $d_{in}$），输出 Y 需要 All-gather（通信量正比于输出大小 $d_{out}$）。假设 batch size 为 B，序列长度为 L，那么一次前向的通信量大约是：

这里乘以 2 是因为有前向和反向各一次（近似），乘以 4 是因为浮点数占 4 字节。这个公式告诉我们：通信量随模型宽度（隐藏层维度）线性增长，而不是随卡数增长。也就是说，卡越多，每卡计算量下降，但通信量不变——这就是为什么并行度不能无限增加。

回到图书馆：如果 100 个人抄书，传稿子的时间不会因为人多了而变少，每次传稿子还是要走同样的路。如果书足够厚、每人的抄写时间足够长，那么传稿子的时间占比就小；如果书很薄但人很多，传稿子的时间可能比抄写时间还长，那就得不偿失了。这就是为什么行并行通常用在“单层足够大”的场景。

工程实践要点

了解了原理和代价，我们来看看实际使用行并行时需要注意的工程细节。这里不是烧脑的数学，而是前辈们在实践中踩过的坑。

第一个要点是并行度的换算。假设你有 16 张卡，数据并行度 DP=4，张量并行度 TP=4，那么实际采用的 batch size 要相应调整。这是因为数据并行“复制模型、切分数据”，而张量并行“切分模型、共享数据”，两者相反。在 Megatron 框架中，你可以通过 --tensor-model-parallel-size 参数指定 TP 大小，框架自动帮你切分权重矩阵。

第二个要点是通信体积的优化。有时候权重矩阵太大，通信量惊人。一个常用技巧是批量通信合并——把多个小的通信合并成一个大通信，减少通信启动次数。这就像把多个小包裹打成一个大包裹一次性寄出，省去了多次跑邮局的时间。

第三个要点是通信与计算重叠。现代 GPU 有独立的通信单元（如 NVLink 和 RDMA），可以边计算边通信。在 Megatron-LM 的实现中，前向传播的 All-gather 可以与当前层的矩阵乘法并行执行（称为“通信计算重叠”）。这就像抄书员一边写新的章节，一边让传话员把上一章送去审阅——两件事同时进行，总时间大幅缩短。

1. 并行度换算：总卡数 $N$ 分解为数据并行度 $DP$ 和张量并行度 $TP$ 的乘积

2. 通信合并：将 $k$ 个小通信合并为一个大通信，减少启动开销 $O(k)$ 次为 $O(1)$ 次

3. 计算重叠：设计算时间为 $T_{\text{calc}}$，通信时间为 $T_{\text{comm}}$，重叠后总时间

回到图书馆的比喻：选择 16 本书由 4 人负责管理、每本书 4 人抄写，这是并行度换算；把几章的内容打包一起传阅，这是通信合并；边抄写边送审，这是计算通信重叠。掌握了这三招，你就能在分布式训练的“图书馆”里，调度得高效从容。

AI 应用场景

Python 演示

import torch def row_parallel_forward(X, W, tp_size=2):    """    简化版行并行前向传播（模拟 2 卡场景）    参数说明:      X: 输入矩阵, 形状 (batch, in_features)      W: 权重矩阵, 形状 (in_features, out_features)      tp_size: 张量并行度（切成几份）    返回:      Y_full: 完整输出, 形状 (batch, out_features)    """    batch, in_features = X.shape    in_f, out_f = W.shape    assert in_features == in_f, "输入维度与权重维度不匹配"     # 1. 把权重矩阵 W 按行切成 tp_size 块    rows_per_gpu = out_f // tp_size  # 每卡分到的行数    W_chunks = []    for i in range(tp_size):        start = i * rows_per_gpu        end = start + rows_per_gpu        W_chunks.append(W[:, start:end])  # 每块形状: (in_features, out_f/tp_size)     # 2. 每张卡用完整的 X 和属于自己的 W 行切片做矩阵乘法    partial_outputs = []    for i in range(tp_size):        chunk = W_chunks[i]              # 第 i 张卡的权重切片        part = torch.matmul(X, chunk)    # 部分输出 (batch, out_f/tp_size)        partial_outputs.append(part)     # 3. 拼接所有部分输出得到完整结果（模拟 All-gather）    Y_full = torch.cat(partial_outputs, dim=1)    return Y_full  def row_parallel_backward(dY, W, X, tp_size=2):    """    简化版行并行反向传播（模拟 2 卡场景）    计算权重梯度 dW 和输入梯度 dX    """    batch, in_features = X.shape    in_f, out_f = W.shape    rows_per_gpu = out_f // tp_size     # 1. 把 dY 按列切成 tp_size 块（对应输出拼接的逆操作）    dY_chunks = []    for i in range(tp_size):        start = i * rows_per_gpu        end = start + rows_per_gpu        dY_chunks.append(dY[:, start:end])     # 2. 每张卡独立计算权重梯度 dW_chunk = X^T @ dY_chunk    dW_chunks = []    dX_chunks = []    for i in range(tp_size):        chunk = dY_chunks[i]        dW_chunk = torch.matmul(X.T, chunk)      # 本卡权重梯度        dW_chunks.append(dW_chunk)        dX_chunk = torch.matmul(chunk, W[:, i*rows_per_gpu:(i+1)*rows_per_gpu].T)        dX_chunks.append(dX_chunk)               # 本卡对输入梯度的贡献     # 3. 拼接 dW（各卡独立更新权重）    dW_full = torch.cat(dW_chunks, dim=1)     # 4. All-reduce: 输入的梯度 = 所有卡贡献之和    dX_full = torch.stack(dX_chunks, dim=0).sum(dim=0)    return dX_full, dW_full  # 冒烟测试: 验证行并行结果与标准矩阵乘法一致if __name__ == "__main__":    torch.manual_seed(42)    batch, in_f, out_f = 4, 8, 6   # 测试维度: 输出 6 列,可被 2 整除    tp_size = 2     X = torch.randn(batch, in_f)    W = torch.randn(in_f, out_f)     # 标准矩阵乘法    Y_standard = torch.matmul(X, W)     # 行并行前向    Y_row_parallel = row_parallel_forward(X, W, tp_size)     # 验证    assert torch.allclose(Y_standard, Y_row_parallel, atol=1e-6), \        "行并行结果与标准结果不一致!"    print("前向传播验证通过: 行并行结果与标准矩阵乘法一致")    print("Y_standard 形状:", Y_standard.shape)    print("Y_row_parallel 形状:", Y_row_parallel.shape)     # 测试反向传播    dY = torch.randn(batch, out_f)    dX_parallel, dW_parallel = row_parallel_backward(dY, W, X, tp_size)     # 验证 dW: 标准梯度 = X^T @ dY    dW_standard = torch.matmul(X.T, dY)    assert torch.allclose(dW_parallel, dW_standard, atol=1e-6), \        "权重梯度不一致!"    print("反向传播验证通过: 权重梯度与标准反向传播一致")    print("dX 形状:", dX_parallel.shape, " dW 形状:", dW_parallel.shape)

1. W_chunks.append(W[:, start:end])这是行并行的核心操作，把权重矩阵按行切片，第 i 张卡拿到属于自己的那部分列（PyTorch 中“行”维度的切片用列索引表示）
2. partial_outputs.append(torch.matmul(X, chunk))每张卡用完整的 X 与自己持有的权重切片独立做矩阵乘法，互不干扰
3. Y_full = torch.cat(partial_outputs, dim=1)模拟 All-gather，把部分输出在列方向拼接成完整输出
4. torch.matmul(X.T, chunk)反传权重梯度的局部计算，每张卡算出属于自己的那部分梯度
5. torch.stack(dX_chunks, dim=0).sum(dim=0)模拟 All-reduce，对输入梯度做跨卡求和

常见误区

面试题