什么是张量并行（Tensor Parallelism）？一种在层内切分模型以训练超大 Transformer 的模型并行方法

张量并行（Tensor Parallelism）是一种层内模型并行（intra-layer model parallelism）方法：它把 Transformer 单个层内部的权重张量切分到多张加速卡上，让每张卡只保存并计算其中一部分，再通过集合通信把结果拼合起来。它要解决的问题是显存约束——当模型参数量增长到数十亿级别时，单张 GPU 的显存无法容纳完整的模型副本，仅靠数据并行（每张卡保存一份完整模型）已经不可行，必须把模型本身拆开。

本词条主要依据 arXiv 论文《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》（Shoeybi 等，2019 年首次提交，2020 年修订至 v4）所描述的方法。该论文提出并实现了一种简单、高效的层内模型并行方案，用于训练参数量达数十亿的 Transformer 模型。

为什么重要

在张量并行被系统化描述之前，训练大模型主要依赖两类手段。一类是数据并行：把同一份模型复制到多张卡上，每张卡处理不同的数据批次，再同步梯度。这种做法实现简单，但显存开销随模型规模线性增长，每张卡都必须装下完整的参数、梯度和优化器状态，模型一大就撞上显存墙。另一类是流水线模型并行（pipeline model parallelism）：按层把模型切成若干段，不同段放在不同设备上，数据像流水线一样依次流过。它能缓解单卡显存压力，但会引入流水线气泡（bubble），设备在等待上下游时存在空闲，且切分粒度受层数限制。

论文指出，非常大的模型由于内存约束而相当难以训练。其提出的层内并行把切分点下沉到层内部，与流水线模型并行正交且互补：前者解决单层放不下的问题，后者解决层数太多、整模型放不下的问题，两者可以叠加使用。论文强调，该方法不需要新的编译器，也不需要修改库，只需在原生 PyTorch 中插入少量通信操作即可完整实现，这降低了把研究代码迁移到多卡环境的门槛。论文用 512 张 GPU 收敛了最高 83 亿参数的 Transformer 模型，整个应用维持 15.1 PetaFLOPs，相对单 GPU 强基线（39 TeraFLOPs，为峰值算力的 30%）取得 76% 的扩展效率。

工作机制

张量并行的核心做法是：对 Transformer 层中的矩阵乘法，按列或按行切分权重矩阵，使每张卡只持有一部分权重，并在必要的位置插入通信操作。论文给出的要点可以概括为：

按列切分前馈网络的第一层：把第一个线性层的权重按列分块，各卡独立计算得到部分激活，无需立即通信。
按行切分前馈网络的第二层：第二个线性层的权重按行切分，使各卡的部分结果可以直接相加；此时插入一次全归约（all-reduce），把各卡的部分和汇总为完整输出。
多头注意力同样切分：注意力中的 Q、K、V 投影按注意力头维度切分到各卡，每张卡负责一部分头，输出投影按行切分，同样以一次全归约收尾。
前向与反向各一次通信：上述结构使每个 Transformer 块在前向传播中只需两次全归约（前馈网络一次、注意力一次），反向传播中对应再各一次，通信量与切分方式匹配。
不引入新编译器或库改动：全部实现只是在原生 PyTorch 中插入少量通信算子，因此可以与流水线并行组合，形成更完整的模型并行方案。

从通用机制上说，这种「列切分—行切分」的配对设计之所以有效，是因为按列切分后的第一层输出天然是分块的，而按行切分的第二层可以把各卡的部分乘积直接相加，从而把通信压缩到每个块仅一次归约。这属于矩阵分块乘法的常规性质，论文的贡献在于把它落到 Transformer 的具体结构上，并给出可运行、可扩展的实现与实测数据。

典型例子

论文给出了若干真实训练结果，可作为该方法的直接例证：

使用 512 张 GPU 训练了最高 83 亿参数的 Transformer 模型，整个应用维持 15.1 PetaFLOPs，相对单 GPU 强基线取得 76% 的扩展效率。
训练了一个类似 GPT-2 的 83 亿参数 Transformer 语言模型，在 WikiText103 上取得困惑度 10.8（论文对比的此前最优为 15.8），在 LAMBADA 上取得 66.5% 准确率（论文对比的此前最优为 63.2%）。
训练了一个类似 BERT 的 39 亿参数模型，在 RACE 数据集上取得 90.9% 准确率（论文对比的此前最优为 89.4%）。
论文还报告，在类 BERT 模型中，层归一化（layer normalization）的放置位置对模型规模增大后的性能提升至关重要，这是该工作在工程实践层面给出的另一条观察。

需要说明的是，上述数字均出自该论文自身的实验设置与对比口径，属于论文作者报告的结果，而非独立第三方复现的结论。

边界与常见误解

张量并行常被误解为「只要把模型切到多卡上就能加速」。实际上它首先是一种显存换通信的手段：切分越细，单卡显存占用越低，但每个 Transformer 块都要执行全归约，卡间通信开销随之上升。论文的 76% 扩展效率是在特定硬件互联条件与特定模型规模下测得的，不能外推为任意集群、任意切分度都成立。

其次，它不等于流水线并行。张量并行切的是层内张量，流水线并行切的是层与层之间的段；论文明确把两者描述为正交且互补，实际大规模训练往往需要组合使用，而不是二选一。它也不解决数据并行所处理的问题——数据并行关心吞吐与梯度同步，张量并行关心单层能否装下。

第三，论文的方法依赖在原生 PyTorch 中手工插入通信操作，这意味着实现者需要理解每一处切分对应的通信位置；它并非「无需改动即可自动并行」的通用编译器方案。论文强调的是不需要新编译器或库改动，而不是不需要任何代码改动。

最后，该论文的结论有明确适用范围：实验对象是 Transformer 类语言模型，规模上限为 83 亿参数，硬件为 512 张 GPU。对于远小于此的模型，张量并行带来的通信开销可能超过收益；对于其他架构，切分方式需要重新设计。论文中关于层归一化放置的观察也仅针对类 BERT 模型报告，不宜直接套用到所有结构。

参考资料

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism