张量并行和流水线并行原理深入理解与思考

文章目录

[一张量并行------横向切分（层内并行）](#一张量并行——横向切分（层内并行）)
- [1.1 核心思想：算子级的拆分](#1.1 核心思想：算子级的拆分)
- [1.2 Megatron-LM 的主流实现方案](#1.2 Megatron-LM 的主流实现方案)
- - [1.2.1 MLP 层（前馈网络）的拆分](#1.2.1 MLP 层（前馈网络）的拆分)
  - [1.2.2 Self-Attention 层的拆分](#1.2.2 Self-Attention 层的拆分)
- [1.3 张量并行的优缺点和使用](#1.3 张量并行的优缺点和使用)
- [1.4 与其他并行技术的区别](#1.4 与其他并行技术的区别)
[二流水线并行 ------ 纵向切分（层间并行）](#二流水线并行 —— 纵向切分（层间并行）)
- [2.1 核心思想：层间的纵向拆分](#2.1 核心思想：层间的纵向拆分)
- [2.2 核心挑战：流水线气泡（Bubble）](#2.2 核心挑战：流水线气泡（Bubble）)
- [2.3 流水线并行的优缺点和使用](#2.3 流水线并行的优缺点和使用)
- [2.4 TP 与 PP 的深度对比](#2.4 TP 与 PP 的深度对比)
[三混合并行的实战 ------ 3D 并行思维](#三混合并行的实战 —— 3D 并行思维)
- [3.1 混合并行举例](#3.1 混合并行举例)
- [3.2 数据流的全程追踪 (Forward Pass)](#3.2 数据流的全程追踪 (Forward Pass))
- [3.3 配置逻辑分析](#3.3 配置逻辑分析)
- [3.4 性能瓶颈分析](#3.4 性能瓶颈分析)
[四注意力头数与切分的数学约束](#四注意力头数与切分的数学约束)
- [4.1 Transformer 的多头注意力（MHA）](#4.1 Transformer 的多头注意力（MHA）)
- [4.2 Num_Heads被TP Size整除](#4.2 Num_Heads被TP Size整除)
- [4.3 流水线并行不关心头数](#4.3 流水线并行不关心头数)
- [4.4 实战建议](#4.4 实战建议)

一张量并行------横向切分（层内并行）

在深度学习模型规模（尤其是 LLM）呈指数级增长的今天，单张 GPU 的显存容量往往无法容纳整个模型的参数。张量并行（Tensor Parallelism, TP） 是一种将单个算子（Operator）的计算及参数拆分到多个设备上的并行技术，是分布式训练和推理的核心黑科技之一。

1.1 核心思想：算子级的拆分

与数据并行（Data Parallelism）每个 GPU 运行完整模型不同，张量并行将模型的权重矩阵（Weights） 直接切开。它的核心逻辑是：既然这一层的矩阵乘法太大，一张卡算不动，那我们就把矩阵拆开，多张卡一起算。
简单来说，如果一个矩阵乘法是 Y = X A Y = XA Y=XA，张量并行会把矩阵 A A A 拆分成 A 1 A_1 A1 和 A 2 A_2 A2，分别放在两个 GPU 上处理。这样，每个 GPU 只需维护一部分参数，从而突破单卡显存限制。

列并行 ：把权重矩阵 W W W 按列切分。 GPU 0 拿到 W W W 的左半部分，算出 Y Y Y 的一部分。GPU 1 拿到 W W W 的右半部分，算出 Y Y Y 的另一部分。此时，每个 GPU 算出的结果拼起来，正好是完整的 Y Y Y。
- 注意：输入 X X X 是完整的，必须被复制（广播）到所有 GPU 上。
行并行 ：把权重矩阵 W W W 按行切分。GPU 0 处理 X X X 的上半部分数据，GPU 1 处理下半部分。每个 GPU 算出一个局部结果，但这些局部结果相加才等于最终的 Y Y Y。算完后，必须执行一次 All-Reduce 操作，把所有 GPU 的局部结果加起来，并同步到所有 GPU 上。

张量并行的通信极其频繁。Transformer 的每一层（包括 Attention 和 MLP）通常需要一次 All-Reduce（甚至多次，取决于具体算子融合情况），极度依赖 GPU 之间的高速互联（如 NVLink），适合单层参数量极大的情况，如 GPT-3 的 175B 参数，必须切分到多张卡上才能放下权重。

1.2 Megatron-LM 的主流实现方案

目前工业界最常用的张量并行方案主要基于 NVIDIA 的 Megatron-LM 论文。它针对 Transformer 架构中的两个核心组件进行了优化：多头注意力机制（MHA） 和 前馈网络（FFP/MLP）。

1.2.1 MLP 层（前馈网络）的拆分

MLP 通常由两个线性层组成： A A A 层（升维）和 B B B 层（降维）。

行拆分（Column Parallelism）： 对第一层权重 A A A 进行列切分。GPU 0 计算 X A 1 XA_1 XA1，GPU 1 计算 X A 2 XA_2 XA2。
列拆分（Row Parallelism）： 对第二层权重 B B B 进行行切分。
减少通信： 巧妙之处在于，两层连续计算时，只有在 MLP 结束时才需要一次 All-Reduce 通信，极大地降低了同步开销。

减少通信的原因如下：在Transformer的实际张量并行实现中混合使用行列拆分。具体来说，标准的 Transformer 块（包括 Attention 和 MLP）采用的是 "列并行 + 行并行"交替 的策略，减少了最小化通信次数。如果我们只用一种，效率会大打折扣。

核心逻辑：利用矩阵乘法的结合律
假设有一个两层的神经网络，计算过程是：
Y = X ⋅ W 1 ⋅ W 2 Y = X \cdot W_1 \cdot W_2 Y=X⋅W1⋅W2
根据结合律，我们可以先算 A = X ⋅ W 1 A = X \cdot W_1 A=X⋅W1，再算 Y = A ⋅ W 2 Y = A \cdot W_2 Y=A⋅W2。
为了省事，我们希望在整个层计算结束后，只做一次通信 ，而不是每层都通信。这就引出了特定的切分顺序： 第一层 ( W 1 W_1 W1) 用列并行 、第二层 ( W 2 W_2 W2) 用行并行。

实战演练：Transformer 的 MLP 层。Transformer 的 MLP 层结构通常是：Linear1 -> GeLU -> Linear2。假设输入是 X X X，中间维度扩大 4 倍，最后再变回去。

第一步：Linear 1 ------ 列并行

操作： Y 1 = X ⋅ W 1 Y_1 = X \cdot W_1 Y1=X⋅W1。
输入 X X X：完整复制到所有 GPU。
切分方式 ：权重 W 1 W_1 W1：按列切分。GPU 0 拿左半，GPU 1 拿右半。
结果：每个 GPU 算出了 Y 1 Y_1 Y1 的一部分（也就是中间隐藏状态的一部分）。此时 Y 1 Y_1 Y1 在不同 GPU 上是独立的。GPU 0 只有前半截特征，GPU 1 只有后半截特征。
状态：无需通信，直接进入下一步。

第二步：GeLU 激活。每个 GPU 只对自己那部分 Y 1 Y_1 Y1 做 GeLU。这是独立的，无需通信。

第三步：Linear 2 ------ 行并行

操作： Y f i n a l = Y 1 ⋅ W 2 Y_{final} = Y_1 \cdot W_2 Yfinal=Y1⋅W2
输入 Y 1 Y_1 Y1：保持在各自的 GPU 上不动（即保持特征维度的切分状态）。
切分方式 ：权重 W 2 W_2 W2：按行切分。GPU 0 拿上半部分行，GPU 1 拿下半部分行。
计算：
- GPU 0 算： P a r t ( Y 1 ) ⋅ P a r t ( W 2 ) r o w s Part(Y_1) \cdot Part(W_2)_{rows} Part(Y1)⋅Part(W2)rows
- GPU 1 算： P a r t ( Y 1 ) ⋅ P a r t ( W 2 ) r o w s Part(Y_1) \cdot Part(W_2)_{rows} Part(Y1)⋅Part(W2)rows
数学原理 ：
- 根据矩阵乘法性质，当权重按行切分时，结果是由两部分相加得到的。
- 即： Y f i n a l = ( Y 1 ⋅ W 2 _ t o p ) + ( Y 1 ⋅ W 2 _ b o t t o m ) Y_{final} = (Y_1 \cdot W_{2\top}) + (Y_1 \cdot W{2\_bottom}) Yfinal=(Y1⋅W2_top)+(Y1⋅W2_bottom)
通信：此时，GPU 0 和 GPU 1 各自拿到了一个"不完整的结果"（部分和）。必须执行一次 All-Reduce 操作，把两个 GPU 的结果相加，得到最终的完整 Y f i n a l Y_{final} Yfinal。相加后，所有 GPU 都拥有了完整的 Y f i n a l Y_{final} Yfinal。

为什么要这样组合？你可能会问：为什么不能两个都用列并行，或者两个都用行并行？

如果两个都用列并行：Linear 1 算完是切分的。 Linear 2 如果也要列并行，它需要完整的输入（列并行要求输入复制）。
后果：你在两层之间必须做一次 All-Gather ，把切分的 Y 1 Y_1 Y1 拼完整再发给所有 GPU。这增加了通信量。
如果两个都用行并行 ：Linear 1 如果用行并行，要求输入 X X X 按行切分（也就是切分 Batch）。这通常不符合我们想保持数据并行的初衷，或者需要在 Linear 1 后立刻 All-Reduce，导致通信提前。
最佳组合（列 -> 行） ：Linear 1 (列) 产生的"切分状态"，正好是 Linear 2 (行) 所需要的"输入状态"。中间完全不需要通信！。通信被推迟到了 MLP 的最后，只需要一次 All-Reduce。

1.2.2 Self-Attention 层的拆分

在多头注意力机制中，张量并行利用了其天然的结构优势：

按"头"拆分： 将Query (Q)、Key (K)、Value (V)的多个注意力头分发到不同的 GPU 上。
独立计算： 每个 GPU 计算自己那一组头的注意力结果。
最后合并： 仅在输出线性层（Linear Projection）之后进行一次通信合并。

Attention层也是采用先列拆分后行拆分的方式。
- Attention 的 Q , K , V Q, K, V Q,K,V 投影用的是列并行（类似于 MLP 的 Linear 1）。
- Attention 的 Output 投影用的是行并行（类似于 MLP 的 Linear 2）。

1.3 张量并行的优缺点和使用

维度	表现	备注
显存利用	极高	将模型参数均匀分布，是加载千亿参数模型的先决条件。
通信频率	高	每次 Forward 和 Backward 都需要进行多次跨卡同步。
硬件要求	苛刻	通常要求设备间有极高的带宽（如 NVLink），否则通信延迟会成为瓶颈。
颗粒度	细	属于算子级并行，比流水线并行（Pipeline Parallelism）更精细。

TP的使用时机：

模型太大： 单个模型的 Layer 权重超过了 80GB（A100/H100 显存上限）。
低延迟推理： 需要极快的响应速度，TP 可以利用多卡的算力并行加速单个 Token 的生成。
拥有 NVLink： 如果你的服务器卡间通信带宽较低（仅走 PCIe），TP 的性能提升可能会被通信延迟抵消。

1.4 与其他并行技术的区别

数据并行 (DP/DDP)： DP 每个 GPU 都有完整参数，TP 只有一部分。
流水线并行 (PP)： PP 是按"层"拆（第一层在 GPU 0，第二层在 GPU 1），而 TP 是在"层内"拆。

** 在超大规模集群中，通常采用 3D 并行 策略，即结合 TP + PP + DP。通常在同一个机柜（NVLink 高带宽）内部使用 TP，而在机柜之间使用 PP 或 DP。

二流水线并行 ------ 纵向切分（层间并行）

2.1 核心思想：层间的纵向拆分

与张量并行（层内拆分）不同，流水线并行（Pipeline Parallelism, PP） 采取的是"横向切分"的策略。流水线并行将模型的不同层（Layers） 分配到不同的设备上。总 G P U 数 = t e n s o r − p a r a l l e l − s i z e × p i p e l i n e − p a r a l l e l − s i z e = 4 × 4 = 16 张 G P U 总GPU数 = tensor-parallel-size × pipeline-parallel-size = 4 × 4 = 16张GPU 总GPU数=tensor−parallel−size×pipeline−parallel−size=4×4=16张GPU
S t a g e 1 : [ G P U 0 , G P U 1 ] − 处理模型前 1 / 4 层，每层在 4 个 G P U 间做张量并行。 S t a g e 2 : [ G P U 2 , G P U 3 ] − 处理模型中 1 / 4 层，每层在 4 个 G P U 间做张量并行。 S t a g e 3 : [ G P U 4 , G P U 5 ] − 处理模型次后 1 / 4 层，每层在 4 个 G P U 间做张量并行。 S t a g e 4 : [ G P U 6 , G P U 7 ] − 处理模型最后 1 / 4 层，每层在 4 个 G P U 间做张量并行。 Stage 1: [GPU0, GPU1] - 处理模型前1/4层，每层在4个GPU间做张量并行。\\ Stage 2: [GPU2, GPU3] - 处理模型中1/4层，每层在4个GPU间做张量并行。 \\ Stage 3: [GPU4, GPU5] - 处理模型次后1/4层，每层在4个GPU间做张量并行。\\ Stage 4: [GPU6, GPU7] - 处理模型最后1/4层，每层在4个GPU间做张量并行。\\ Stage1:[GPU0,GPU1]−处理模型前1/4层，每层在4个GPU间做张量并行。Stage2:[GPU2,GPU3]−处理模型中1/4层，每层在4个GPU间做张量并行。Stage3:[GPU4,GPU5]−处理模型次后1/4层，每层在4个GPU间做张量并行。Stage4:[GPU6,GPU7]−处理模型最后1/4层，每层在4个GPU间做张量并行。

数据流动： GPU 0 完成前 16 层的计算后，将中间激活值（Activations）传给 GPU 1。
通信代价：相比 TP，PP 的通信量小得多，因为它只在相邻 GPU 之间传输激活值，不需要像 All-Reduce 那样全局汇总。
解耦： 每个设备只负责维护自己那一组层的权重，极大地节省了单卡的显存。
致命弱点：流水线气泡。当 GPU 1 在处理第一个 micro-batch 的后半段时，GPU 0 可能已经闲下来了（因为 GPU 1 还没腾出空来接收下一个 batch）。这种 GPU 空转等待的时间就叫"气泡"。

2.2 核心挑战：流水线气泡（Bubble）

最简单的流水线模式（Naive PP）存在巨大的效率问题：当 GPU 1 在计算时，GPU 0 必须等待数据传回或处理下一个 Batch，导致硬件利用率低下。这种空闲时间被称为 "气泡"（Bubble）。

为了解决这个问题，业界引入了 Micro-batch（微批次） 机制：将一个大的 Batch 拆分成多个小的 Micro-batch。

1F1B 策略（One Forward, One Backward）： Megatron-LM 采用的经典方案。一旦流水线填满，每个 GPU 在处理一个微批次的前向计算后，立即处理另一个微批次的反向计算。这样可以保持硬件持续忙碌，缩小气泡比例。

2.3 流水线并行的优缺点和使用

维度	表现	备注
通信量	较低	只需在"切分点"传输激活值，通信频率远低于 TP。
显存压力	中等	虽然权重分开了，但为了反向传播，需要缓存多个微批次的激活值。
设备利用率	受限	无法完全消除气泡，且受限于最慢的那一层的计算速度。
硬件要求	友好	对带宽要求不高，适合在没有 NVLink 的跨机柜（网络）环境中使用。

流水线并行使用时机：

跨节点扩展： 当你的模型大到无法装进单台服务器（通常 8 张卡）时，必须使用 PP 在多台机器间切分模型。
带宽受限： 在没有 NVLink 或 InfiniBand 的普通千兆/万兆网环境下，PP 是比 TP 更稳健的选择。
负载均衡： 需要注意每一段流水线的层数和计算量要尽可能对齐，否则最慢的一块卡会拖累整个系统。

2.4 TP 与 PP 的深度对比

特性	张量并行 (TP)	流水线并行 (PP)
切分对象	算子/矩阵内部 (Intra-layer)	模型层级之间 (Inter-layer)
通信内容	矩阵乘法的中间部分结果	层与层之间的激活值与梯度
通信频率	极高 (每层多次同步)	较低 (每组层一次同步)
主要瓶颈	通信延迟 (Latency)	流水线气泡与负载均衡

三混合并行的实战 ------ 3D 并行思维

在实际的大规模 AI 集群中，单打独斗的并行策略已经无法满足千亿甚至万亿参数模型的需求。3D 并行（3D Parallelism） 是将三种维度交织在一起的终极方案：TP（张量并行，横向切算子） 、PP（流水线并行，纵向切层） 、以及 DP（数据并行，数据切样本）。如在训练 GPT-4 或 Llama 3 这种巨型模型时，通常会同时开启：

TP： 解决单卡内算力与显存的矛盾。
PP： 解决跨机器部署的规模问题。
DP/ZeRO： 解决海量数据训练的吞吐量问题。

3.1 混合并行举例

假设模型层数为48，总GPU数为6，具体分配为tensor_parallel_size = 2, pipeline_parallel_size = 3,data_parallel_size=1。6张GPU可以构成一个逻辑上的三维坐标系（ 2 × 3 × 1 2\times 3\times 1 2×3×1）。假设：
- TP Size = 2（每 2 张卡组成一个 TP 组，通过 NVLink 交换权重碎片）
- PP Size = 3（模型被切成 3 段，逐级传递激活值）
- DP Size = 1（由于只有 6 张卡，刚好填满 TP 和 PP 的乘积，此时暂无额外冗余卡做数据并行。若有 12 张卡，则 DP Size = 2）
计算逻辑流程：

Stage 2 (Layers 32-47)
Stage 1 (Layers 16-31)
Stage 0 (Layers 0-15)
TP: NVLink
TP: NVLink
TP: NVLink
PP: P2P Communication
PP: P2P Communication
PP: P2P Communication
PP: P2P Communication
PP: P2P Communication
PP: P2P Communication
PP: P2P Communication
PP: P2P Communication
GPU0
GPU1
GPU2
GPU3
GPU4
GPU5

Stage 0 (流水线第一段): [GPU 0, GPU 1] 负责模型的前 1/3 层（比如 Layer 0-15）。
- 内部关系 ：GPU 0 和 GPU 1 是张量并行关系。它们通过 NVLink 高速互联，共同计算 Layer 0 的矩阵乘法。
- 输出：计算完 Layer 15 后，由 GPU 1（通常是 Rank 0 负责收集）将激活值传给下一阶段的 GPU 2。
Stage 1 (流水线第二段): [GPU 2, GPU 3] 负责模型的中间 1/3 层（比如 Layer 16-31）。
- 内部关系：GPU 2 和 GPU 3 互为 TP 拍档，共同计算每一层。
- 输入：接收来自 GPU 1 的数据。
Stage 2 (流水线第三段): [GPU 4, GPU 5] 负责模型的最后 1/3 层（比如 Layer 32-47）。
- 内部关系：GPU 4 和 GPU 5 互为 TP 拍档。
- 输出：计算 Loss，并开始反向传播。

TP 组内的"显微镜"协作 (Intra-node)：在 Stage 0 内部，GPU 0 和 GPU 1 是"连体婴儿"。

计算细节：当模型运行到第 5 层的一个全连接层时，GPU 0 负责计算左半部分矩阵，GPU 1 负责右半部分。
通信瓶颈 ：它们之间会频繁发生 All-Reduce 通信。因为必须保证带宽极高，这两张卡通常位于同一个物理节点 内，并由 NVLink 直连。

PP 阶段间的"接力棒" (Inter-node/Inter-GPU)：当 Stage 0 完成了前 16 层的计算后，它需要把结果交给 Stage 1。

通信细节 ：这是一种 P2P (Point-to-Point) 通信。GPU 0/1 会将激活值（Activations）打包发送给 GPU 2/3。
流水线气泡管理 ：为了不让后方的 GPU 4/5 干等着，系统会采用 1F1B (One Forward, One Backward) 调度策略，让不同 Micro-batch 的前向和后向计算在流水线上交错执行。

3.2 数据流的全程追踪 (Forward Pass)

一个数据包（Micro-batch）在 3D 架构中的流程：

步骤	物理位置	执行操作	通信负载
1	Stage 0 (GPU 0, 1)	计算第 0-15 层。每层内部进行 TP 切分计算。	极高 (NVLink 频繁同步)
2	传递 0 -> 1	将第 15 层的输出传给下一站。	低 (仅传输激活值张量)
3	Stage 1 (GPU 2, 3)	接收数据，计算第 16-31 层。	极高 (NVLink 频繁同步)
4	传递 1 -> 2	将第 31 层的输出传给下一站。	低
5	Stage 2 (GPU 4, 5)	计算第 32-47 层，计算 Loss 值。	极高 (NVLink 频繁同步)

3.3 配置逻辑分析

2x3的配置是基于硬件限制的最优权衡。

优先满足 TP ：TP 对延迟极其敏感，所以一定要把 TP 组放在同一个机器内。如果 tensor_parallel_size 设得太大（超过单机 8 卡），跨机的网络延迟会拖垮整个训练。
用 PP 填补显存 ：当单机显存塞不下 48 层模型时，通过 pipeline_parallel_size = 3 将模型纵向拉长，分布到更多机器上。
用 DP 提速 ：如果显存够了，但训练太慢，可以再增加 data_parallel_size，让多个 3D 并行的"副本"同时跑不同的数据。

3.4 性能瓶颈分析

在 3D 并行中，木桶的最短板通常来自：

通信比：如果网络带宽（如 InfiniBand）不够，PP 之间的传输会变成瓶颈。
负载不均：如果 Stage 0 包含 20 层，而 Stage 2 只有 10 层，那么 GPU 4/5 会长时间处于空闲状态。
显存峰值：PP 需要缓存多个 Micro-batch 的激活值以备反向传播，这可能导致某张卡的显存突然"爆表"。

四注意力头数与切分的数学约束

最容易出错的地方，也是大模型编程中的"细节魔鬼"。

4.1 Transformer 的多头注意力（MHA）

在 Transformer 的每一层，Attention 模块的输入是一个向量，维度通常是 hidden_size。多头注意力将这个维度切分为 num_attention_heads 个头。公式： H e a d _ D i m = H i d d e n _ S i z e N u m _ H e a d s Head\_Dim = \frac{Hidden\_Size}{Num\_Heads} Head_Dim=Num_HeadsHidden_Size例如 Qwen2.5-14B：
- Hidden_Size = 5120
- Num_Heads = 40
- Head_Dim = 128

这意味着 Q , K , V Q, K, V Q,K,V 的投影层权重矩阵形状大致为 (Hidden_Size, 3 * Hidden_Size)。而在内部逻辑中，这被看作 (Num_Heads, 3 * Head_Dim)。

4.2 Num_Heads被TP Size整除

这是张量并行在 Attention 层实现的物理限制。
场景 A：使用 TP = 2 ，在GPU 0和GPU 1上切分 Attention 层。通常做法是 按头切分（Column Parallelism 的变种）。总共 40 个头，GPU 0 负责计算第 0-19 个头。GPU 1 负责计算第 20-39 个头。结果：完美切分，负载均衡。
场景 B：使用 TP = 3 ，在 3 个 GPU 上切分 40 个头。 40 ÷ 3 = 13.33... 40 \div 3 = 13.33... 40÷3=13.33...。无法将一个完整的"头"再物理切开分配给不同的 GPU（虽然理论上可以做 further 的 finer-grained 切分，但主流实现如 Megatron-LM、HuggingFace 为了通信效率和算子对齐，要求 Head 是最小分配单元）。
如果强行分配：GPU 0 算 13 个头，GPU 1 算 13 个头，GPU 2 算 14 个头。这会导致显存占用不均，计算负载不均（GPU 2 会慢一点，拖慢整体节奏）。如果代码逻辑不支持非整除：程序直接报错 ValueError。

4.3 流水线并行不关心头数

流水线并行（PP）根本不在乎 Attention 头有多少 。PP 只是把"整个 Transformer Block"当做一个黑盒子搬来搬去。无论这个 Block 里有 32 个头还是 128 个头，对于 PP 来说，它只是一个待处理的层。所以，Num_Heads 只需要满足 Num_Heads % TP_Size == 0，与 PP_Size 无关。

4.4 实战建议

当准备训练一个大模型时，配置 GPU 集群必须满足以下硬性数学约束：
Num_Attention_Heads ( m o d Tensor_Parallel_Size ) = 0 \text{Num\_Attention\_Heads} \pmod{\text{Tensor\_Parallel\_Size}} = 0 Num_Attention_Heads(modTensor_Parallel_Size)=0

查表：在选模型和配 GPU 之前，先看模型的 config.json 中的 num_attention_heads。
常见头数 ：
- Llama-2-7B: 32 heads (适合 TP=1, 2, 4, 8, 16, 32)
- Llama-3-70B: 64 heads (适合 TP=1, 2, 4, 8, 16, 32, 64)
- Qwen-14B: 40 heads (适合 TP=1, 2, 4, 5, 8, 10, 20, 40)
KV Cache 的影响：在推理阶段，KV Cache 也是按头存储的。TP 切分头数同样能均匀分摊 KV Cache 的显存占用。