张量并行

3秒看懂

张量并行(Tensor Parallelism, TP)是一种将模型中单个张量(通常是权重矩阵)沿特定维度切分到多个GPU上的分布式训练/推理策略。其核心特征是单层内跨设备并行,部分计算步骤需要设备间通信来整合结果,因此对GPU间高带宽、低延迟互连(如NVIDIA NVLink)依赖极重。

一句话:把一个大矩阵切到多张卡上一起算,关键步骤需高速通信同步中间结果。

典型场景:超大模型整体参数量与优化器状态超出单卡显存，通过张量并行切分单层权重以降低每卡显存占用，使大模型训练/推理可行。

3分钟产业解释

为什么需要张量并行

大模型整体参数与优化器状态远超单GPU显存容量。以GPT-3 175B为例，FP16参数约350GB，单卡HBM 80GB远无法容纳。数据并行每卡复制全模型,对显存无效;流水线并行切层导致流水线气泡;张量并行则直接切开单层权重,让多卡协作计算同一层，显著降低每卡显存压力。

产业价值

张量并行是大模型训练的基础使能技术。没有它,参数规模超千亿的模型无法在现有GPU上训练。也是推理时降低单卡显存压力、提升吞吐的核心手段。

核心权衡

• 优势:显存效率极高,单层可扩展至任意数量GPU。
• 代价:每次前向/反向传播需通信,通信量与模型并行度正相关,对节点内高速互连(NVLink、甚至NVL72级别的超带宽)构成硬需求。跨节点张量并行通常因带宽不足难以工程化。

产业现状

主流框架(NVIDIA Megatron-LM、DeepSpeed、PyTorch FSDP部分场景)均内置张量并行。GPT-4、Llama 3 400B+等万亿参数模型均依赖TP实现。NVIDIA DGX/DSS等超节点通过NVSwitch全互联,正是为张量并行提供物理基础。

15分钟专家深入

核心机制分解

张量切分维度选择(以Transformer自注意力为例):

• 列切分(Column-wise):按权重矩阵的列切分,每卡持有一部分列。输入需广播或全收集。
• 行切分(Row-wise):按权重矩阵的行切分,输入需在行维度对应切分,输出需部分求和。

通信模式:

• 列并行线性层:前向无需通信(每卡独立计算),但后接激活函数后进入行并行层时需AllReduce。
• 行并行线性层:前向需AllReduce汇聚各卡部分和,然后进入激活函数。
• Megatron-LM自注意力:组合列/行并行,巧妙安排通信时机,减少通信次数。

张量并行度(Tensor Parallel Size, TP Size):张量并行使用的GPU数。通常限制在单节点内(如8卡),因为跨节点带宽急剧下降。NVIDIA Hopper架构(NVLink 4.0, 900GB/s双向/卡)和未来NVL72(超节点级NVSwitch)在尝试突破节点限制。

通信原语:

• AllReduce:行并行输出汇聚。
• AllGather:列并行权重切分后,输入需在全卡可见时使用。
• ReduceScatter:反向传播时,梯度需分散回各卡。
• 注意:张量并行不涉及All-to-All(All-to-All多见于MoE的专家调度dispatching,是易混淆点)。

与其他并行策略的耦合(3D并行)

实际大模型训练采用3D混合并行:

• 数据并行(DP):外层,跨节点/机架,低通信频次,高延迟容忍。
• 流水线并行(PP):中间层,切分层间,减少通信量但需处理bubble。
• 张量并行(TP):内层,切分层内,高通信量,依赖节点内NVLink。

组合顺序:通常TP×PP×DP。例如,2048 GPU集群:TP=8(单节点),PP=8(层间),DP=32(数据副本)。

显存与通信数学

• 显存:每卡存储全量的 1/TP_Size ,切分维度参数线性压缩。
• 通信量:与TP Size正相关,与计算量之比恒定。例如,行并行层输出One AllReduce操作,数据量为batch_size × seq_len × hidden_dim。若TP Size过大,通信开销抵消并行收益。

技术原理

矩阵切分原理(以Transformer自注意力示例)

简化线性层定义:Y = X · W

列并行前向(Column Parallel Linear):

• 将权重W按列切为 [W_1, W_2],分别在GPU1和GPU2。
• 输入X通过AllGather或预先广播使两卡持有完整X。
• GPU1计算Y_1 = X · W_1,GPU2计算Y_2 = X · W_2。
• 此时两卡各持有部分输出列,可直接进入后续的列并行激活函数或行并行层。

行并行前向(Row Parallel Linear):

• 将权重W按行切为 [W_1; W_2],分别在GPU1和GPU2。
• 输入X也按列切分为 [X_1, X_2]。
• GPU1计算Y_1 = X_1 · W_1,GPU2计算Y_2 = X_2 · W_2。
• 输出Y = Y_1 + Y_2,需AllReduce跨卡求和,恢复完整输出。

Megatron-LM 自注意力张量并行机制(典型实现)

以多头自注意力为例:QKV投影和输出投影采用矩阵切分策略,使得通信最优。

输入 X (已切分或完整)
    |
    ├─> 列并行 QKV 线性层(无通信)
    |   对每个头独立计算 Q,K,V
    |
    ├─> 各卡独立计算注意力上下文(Head级)
    |   本次不跨卡
    |
    └─> 行并行输出线性层
        - 输入经注意力后仍为切分状态
        - 计算Y_1 = Context_1 · W_O_1, Y_2 = Context_2 · W_O_2
        - AllReduce 求和得最终输出

关键:仅输出线性层后需要一次AllReduce,前向通信量最小化。

关键参数与体系结构约束

• TP Size:通常2/4/8,受节点内GPU数量与NVLink拓扑限制。超节点(如NVL72)理论上可到72,但通信效率仍需验证。
• NVLink带宽:双向带宽越高的GPU(如H100 SXM, 900GB/s),张量并行效率越高。PCIe版本GPU(如H100 PCIe, 128GB/s)张量并行扩展性差。
• 计算通信重叠:通过异步通信和CUDA Stream,将AllReduce与后续计算重叠,隐藏通信延迟。

数值精度问题

张量并行切分后AllReduce求和,若使用FP16/BF16,需保证求和精度无损。主流框架采用BF16张量并行,因动态范围大,无需loss scaling。FP8张量并行(NVIDIA Transformer Engine)通过在线缩放进一步加速。

技术演进史

技术路线对比(量化表)

结论:TP是显存瓶颈的终极解但通信瓶颈重,仅当模型整体显存超出单卡时使用。

上下游

上游(依赖技术)

• 高速互连:NVLink(各代)、NVSwitch、InfiniBand(传输层),TP的物理基础。NVLink带宽决定了TP Size上限。
• GPU架构:Tensor Core计算能力、HBM容量/HBM带宽,HBM容量不决定TP Size,但越小,越早触发TP。
• 集合通信库:NCCL(AllReduce/AllGather implementations)的性能决定了TP通信开销。NCCL对NVLink/RDMA的利用效率是核心。

下游(受益应用)

• 超大模型训练:GPT-4、Gemini、Llama 3 405B+等,单层参数量跨卡,TP是必须组件。
• 低延迟推理:切分单层权重到多卡,提升推理吞吐/降低单卡延迟(计算并行化)。
• 科学计算:物理模拟中大型稀疏矩阵并行求解,TP可类比Message Passing矩阵切分方法。
• 搜索推荐模型:超大Embedding Table行并行切分,索引分布在各卡,前向AllReduce汇聚。

关键指标

• TP Communicaton Overhead: AllReduce耗时 / 计算耗时(层内)。理想<10%,取决于batch size与hidden dim。batch size越大,计算量上升,通信占比下降,TP扩展性改善。
• MFU(Model FLOPs Utilization):实际计算吞吐/理论峰值。TP的额外通信会降低MFU,通常8卡TP相比单卡,MFU从50%降至40%(HW依赖)。
• TP Breakeven Point:TP受益临界点。当单层参数>单卡HBM时,TP首次显存可行;但当TP通信时间>计算时间时,加速比负收益。通常TP Size≤8较为合理。
• 显存碎片率:切分不均衡导致部分卡存储不均匀,可能引发OOM。需框架支持自动平衡。

供需与市场数据

本部分受限于搜索失败,无具体行业出货量、营收数字,以定性为主。

• 需求侧:生成式AI训练集群规模化部署(万卡至十万卡),对TP的需求刚性。每一块H100/B200 DGX系统内,TP作为训练作业的标准组件启用。
• 供给/硬件渗透:NVIDIA DGX/HGX提供了支持8卡全互联的物理基座,是TP的主导硬件方案。AMD MI300X平台通过Infinity Fabric互连,也支持类TP切分(ROCm生态下)。
• 云端部署:AWS p5、Azure ND H100 v5、Google Cloud TPU v5p(TPU内有张量并行机制不同,采用ICI互联)等实例类型内置张量并行环境。
• 无具体市场规模数据[需引用可信行业报告,此处未获取]。

代表公司与资本映射

由于搜索结果为空,以下基于公开知识定性列出,不做具体财务/股票推荐,不对标盲目投资。

• NVIDIA:TP的硬+软全栈主导者。硬件(NVLink/NVSwitch)提供TP带宽基础,软件(Megatron-LM,Transformer Engine,NeMo)定义TP实践范式。资本关联:GPU生态份额较高,但具体份额需补权威来源后再写入。
• Google(TPU):TPU设计中张量并行由ICI(Inter-Chip Interconnect)实现,虽术语不同,本质相同,TPU v5p支持大规模TP。
• AMD:MI300X等自研互连+ROCm/RCCL,开展TP生态构建,性能仍落后于NVIDIA生态,市场验证中。
• Microsoft/OpenAI:大模型训练重度使用TP,投资于NVIDIA硬件,自研刀片服务器优化TP通信。
• Meta:自研硬件+PyTorch生态内置TP支持,开源贡献大。
• Hugging Face/云厂商:提供TP开箱即用推理服务,按Token收费,TP支撑大模型推理API经济。

投资逻辑

(基于技术趋势推演,非确定性建议)

• 互连带宽升级周期是TP扩散的直接受益方向:未来3年,NVLink从900GB/s向更高(如1.8TB/s)演进,跨节点统一互连(如NVL72)将扩大TP可用域,拉动高端GPU与交换机需求。
• TP生态绑定加深NVIDIA护城河:Megatron-LM等库的TP实现与NCCL强绑定,竞品自研互连+集合通信库需长期打磨才可等效替代。
• 关注“TP效率稀释”风险:进入万卡级集群后,TP已非瓶颈,MFU主导因素转向DP/PP和数据加载,TP增速放缓。
• 推理侧TP增长快:大模型推理需求爆发,TP降低延迟,成为推理部署必备,催生高密度推理服务器(如GH200 NVL32)需求。

常见误读纠偏

误读1:张量并行与数据并行总是一起用,所以等效。

• 误解:TP和DP都是“多用卡加速”,实际DP每卡存全模型,TP每卡存部分模型;显存限制下,单靠DP无法运行大模型,TP是显存的破局点,二者互补非等效。

误读2:张量并行通信用All-to-All。

• 事实:标准Transformer张量并行通信原语是AllReduce(行并行)或AllGather/ReduceScatter(列并行变种)。All-to-All主要用于MoE的expert dispatch/combine,与TP无关。混淆将导致性能建模全错。

误读3:TP Size越大越好。

• 事实:TP增加通信开销。8卡TP后的加速比显著衰减。单节点内最大64/72卡TP(如NVL72)的理论有效性仍待工程验证,应用需通信计算重叠深度优化。

学习路径

1. 理论基础

   • 阅读《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》(Shoeybi et al., 2019),掌握列/行切分数学。
   • 理解集合通信原语(AllReduce/AllGather/ReduceScatter),推荐《NCCL官方文档》与Amini、Broadcast等图解。

2. 动手实践

   • 用PyTorch的torch.distributed.tensor.parallel或DeepSpeed的TP模块,对一个小型GPT模型进行张量并行切分并分析通信trace(Nsight Systems)。
   • 在NVIDIA NeMo或Megatron-LM仓库运行TP Size=2/4/8对比微批次吞吐。

3. 进阶研究

   • 解读Megatron-LM 3D并行码,关注TP与PP/DP的接口交互,tensor sharding策略。
   • 研究TP与序列并行的结合(sequence parallelism),以及FP8通信压缩,理解未来方向。

一句话总结

张量并行是解决大模型整体参数与优化器状态超出单GPU显存的必备技术,通过切分单层权重降低每卡显存占用，依赖极致节点内带宽(AllReduce)实现大模型单层多卡协同,但通信约束使其实践限制在单节点内,构成大模型3D并行的最内层核心。

延伸阅读与来源

• 核心论文:
  Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, Shoeybi et al., 2019.
  Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, Narayanan et al., 2021. (含3D并行)
• NVIDIA官方:
  NVIDIA Megatron-LM GitHub, NVIDIA Transformer Engine文档, NCCL官方文档。
• DeepSpeed:
  ZeRO-3 & Tensor Parallelism Integration docs。
• 行业观察:
  SemiAnalysis对NVL72及张量并行架构的分析(需订阅),但不作为数据来源。
• 说明:因实时检索未果,本页面中硬件规格(如NVLink 900GB/s)源自公开产品页面;无第三方引用数字已标[厂商公开],并无内部数据。