张量并行：千亿大模型突破单卡显存限制的底层技术
摘要：千亿、万亿参数大模型体量巨大，完整模型甚至单层网络都无法装入单张 GPU 显存。数据并行、流水线并行都存在各自的局限，**张量并行（Tensor Parallelism，TP）** 通过拆分层内张量矩阵，把同一个网络层的计算分摊到多张 GPU 并行执行，解决单层显存爆炸难题。它是 Megatron‑LM 等主流大模型框架的核心能力，通常与流水线并行、数据并行共同组成大模型训练的三维并行架构。大语言模型参数规模飞速增长，千亿级别已经成为行业常态。以 FP16 精度为例，1600 亿参数仅参数权重就要占用 200GB 以上显存；再叠加训练过程中的梯度、优化器状态，显存需求会进一步放大数倍，远远超过单张 H100 80GB 的物理上限。很多人会疑惑：已经有多卡的数据并行，为什么还需要张量并行？本文对比三种主流并行方案，拆解张量并行的底层原理、矩阵切分方式、Transformer 工程实现，同时分析通信开销的约束，以及三维并行的落地思路。
一、现有并行方案为什么不够用
1.1 数据并行（DP/DDP）：切数据，不切模型
数据并行的逻辑：多张 GPU 保存一份完整的模型副本，每张 GPU 只处理一部分训练数据，训练结束后同步梯度完成参数更新Hugging Fa...。
✅优点：实现简单，工程改造成本低，主要用来提升训练吞吐。❌硬伤：每张显卡都需要存放完整模型。模型本身超过单卡显存时，无论怎么拆分数据，都无法解决 “模型装不下” 的核心矛盾。
1.2 流水线并行（PP，按层切分）：纵向切模型
流水线并行把 Transformer 的网络层做纵向切割，例如把前 10 层放到 GPU0，中间 10 层放到 GPU1，后 10 层放到 GPU2。数据像流水线一样，依次流经各个 GPU 完成计算。
✅优点：可以把整个大模型拆分到多张显卡，解决整体模型过大的问题。❌两大短板：
[*]存在流水线气泡：同一时刻只有部分 GPU 在运算，其余 GPU 空闲等待，硬件利用率下降；
[*]如果某一层本身规模就极大，比如 Transformer 中维度数万的前馈网络，单一层权重矩阵就占用几十 GB 显存，此时按层切分也会失效。
当单层网络就大到单卡放不下，无论数据并行还是流水线并行，都束手无策，张量并行正是为了解决这个场景诞生。
二、什么是张量并行：拆分层内部的矩阵运算
张量并行不再以 “完整网络层” 作为最小单元，而是深入网络层内部，直接拆分权重张量与矩阵乘法运算，将一份巨大的矩阵运算拆成多份子任务交给不同 GPU，各卡只保存一部分权重，计算完成后通过卡间通信合并最终结果飞桨Pad...。
核心本质：用 GPU 之间的通信开销换取显存占用下降。
深度学习网络中最核心的计算就是矩阵乘法 \(Y=XW\)，X 代表输入，W 代表模型权重矩阵。张量并行主要有两种切分策略：按列切（列并行）、按行切（行并行）。
2.1 按列切分（列并行）
假设权重矩阵 W 的尺寸为4096 × 8192，将矩阵沿着输出维度（列）切为两半，GPU0 保存前 4096 列，GPU1 保存后 4096 列。
[*]前向传播：输入 X 完整复制发送给两张 GPU；
[*]GPU0、GPU1 分别执行矩阵乘法，各自产出部分输出 \(Y_0=XW_0\)、\(Y_1=XW_1\)；
[*]通过AllGather通信算子，把两份局部结果拼接，得到完整输出矩阵 Y。
特点：前向传播阶段不需要拆分输入数据，只需要在计算末尾做一次结果归并。
2.2 按行切分（行并行）
同样的权重矩阵，沿着输入维度（行）切分，GPU0 保存上半部分行，GPU1 保存下半部分行。
[*]输入矩阵 X 也需要对应拆分成两份，分别送入两张 GPU；
[*]每张 GPU 独立计算局部矩阵乘法，得到局部结果；
[*]使用AllReduce将多张卡的局部结果做求和归约，输出完整矩阵 Y飞桨Pad...。
特点：输入需要分片，前向传播就会触发卡间求和通信；反向传播也对应一套不同的通信逻辑。
两种切分没有绝对优劣，工程上经常成对组合，目的是尽可能减少通信调用次数。
三、张量并行在 Transformer 中的工程实现（Megatron‑LM）
英伟达 Megatron‑LM 是张量并行的标杆实现，针对 Transformer 两大核心模块：多头自注意力、前馈网络（MLP）做了定制化切分逻辑。
[*]多头自注意力层对 QKV 投影权重矩阵，沿着注意力头维度切分。举例模型一共 64 个注意力头，使用 4 卡张量并行，每张 GPU 负责 16 个注意力头完整计算；各卡运算结束后，通过通信把多头输出拼接归约，得到完整注意力输出。
[*]前馈网络 MLP 层采用 “列切 + 行切” 的配对方案：第一层升维全连接使用列并行，第二层降维全连接使用行并行。巧妙之处：列并行输出分片格式刚好匹配行并行的输入分片，把多次通信合并为一次，大幅降低每一层的通信总开销。
关键约束：张量并行对硬件的要求
张量并行每一轮前向、反向传播，都需要频繁执行AllReduce、AllGather集合通信操作，对 GPU 之间带宽要求极高。
[*]最优场景：同一台服务器内部多张 GPU，依靠 NVLink、NVSwitch 高速互联；
[*]不适合跨多台机器：普通以太网网络延迟高，跨节点通信会带来巨大性能损耗，训练速度大幅下滑。
所以张量并行几乎只在单机内部使用，很少跨服务器部署。
四、三维并行：张量并行 + 流水线并行 + 数据并行协同工作
实际大模型训练几乎不会单独使用张量并行，通常三种并行策略组合，也就是业内常说的三维并行（3D Parallelism）：
[*]单机内部：张量并行 TP，解决单层网络显存爆炸；
[*]跨服务器节点：流水线并行 PP，把模型分层切分到不同机器，解决整体模型体量；
[*]最外层：数据并行 DP，复制整套模型流水线，切分训练数据，提升整体训练吞吐。
通俗理解分工：
[*]张量并行：把 “一层巨大的运算” 拆开分给本机多张卡；
[*]流水线并行：把 “整个模型” 按层分给不同服务器；
[*]数据并行：多组完整流水线同时跑不同批次的数据，加速训练。
这套组合策略，就是今天千亿、万亿参数大模型集群训练的标准范式。
五、总结
[*]数据并行：拆分数据，每张卡持有完整模型，提升吞吐，但解决不了模型过大；
[*]流水线并行：纵向按层切模型，受限于单层最大显存，存在流水线气泡；
[*]张量并行：横向拆分层内权重矩阵，解决单层显存超限，代价是增加卡间通信，适合单机高速互联 GPU 集群。
张量并行不是银弹，它是一种取舍：以通信开销换取显存降低。正是张量并行、流水线并行、数据并行三者互相配合，才让今天的超大规模大模型训练成为现实。理解张量并行，就看懂了大模型分布式训练最核心的底层逻辑。
拓展提示：除张量并行之外，FSDP/ZeRO 系列显存分片技术，也可以实现大模型显存降低；二者技术路线不同，张量并行修改模型内部矩阵运算，ZeRO 系列更多是参数、梯度、优化器状态分片，不需要改动网络计算逻辑。
页:
[1]