张量并行Tensor Parallelism

张量并行（Tensor Parallelism，TP）就是把同一层神经网络中的张量（权重或计算）拆分到多张 GPU 上，并行完成一次层计算。

张量并行并不是 LLM 时代才提出的概念，在深度学习中很早就存在。不过早期模型通常只有几百万到几千万参数，单层参数规模远小于单卡显存容量，因此几乎没有使用张量并行的需求。

训练瓶颈

随着模型大小越来越大，一个 70B 的稠密模型单层就有 1B 左右，在 FP16+Adam 下面训练，则会存储这些东西

项目	字节
FP16 模型权重	2
FP16 梯度	2
FP32 master weight	4
FP32 Adam 一阶矩 $m$	4
FP32 Adam 二阶矩 $v$	4

有些实现会额外保留 FP32 梯度副本，1B 左右的参数需要 16-20G显存。

如果为了防止爆显存而让每个流水线 stage 只放一层 Transformer Block，那么虽然单卡能够容纳这一层，但会导致流水线 stage 数量过多，流水线 bubble 增大，吞吐下降，同时通信次数也会明显增加。

此外，这 20G 只是必须的参数量，还不含激活、临时 buffer、通信 buffer、CUDA context、显存碎片等。

推理瓶颈

如果是 prefill 阶段，模型可以使用 chunked prefill 方法，把输入变成多个 microbatch ，这样做流水线并行的话，几张卡可以一直有输入。

但是在 decode 阶段，模型一次通常只能生成一个 token，每一步计算都依赖上一步生成的结果，因此不同 token 无法像 prefill 那样拆成多个 microbatch。即使使用投机解码，也只能一次验证少量候选 token，很难形成足够深的流水线，因此流水线并行通常难以发挥作用，反而会增加跨 stage 通信，使单 token 延迟高于单卡。

为什么需要张量并行

在上面的思考中，我们发现：

• 流水线并行对吞吐量有收益，但是过多层延迟造成的问题不可忽视。
• 如果把每一层拆给几张卡并行计算，通信开销不大的情况下，对延迟改善应该有收益
• decode阶段流水线并行收益不大。

在考虑这些问题后，自然愿意尝试去做张量并行。

对哪些层

考虑 LLM 典型的推理过程

      ┌──────────────┐
      │  Embedding   │
      └──────┬───────┘
             │
┌────────────▼────────────┐
│ Transformer Block × N   │
│                         │
│  RMSNorm                │
│  Self-Attention         │
│  KV Cache               │
│  MLP                    │
└────────────┬────────────┘
             │
      ┌──────▼──────┐
      │   LM Head   │
      └──────┬──────┘
             │
        logits
             │
      Sampling / Greedy
             │
       Next Token

如果是训练过程的话，后面得到 logits 后过程如下

logits
      ↓
Softmax
      ↓
Cross Entropy
      ↓
Loss
      ↓
Backward

因此，我们主要关注的就是 Embedding,Self-Attention,MLP,Cross Entropy 四部分。以下实际上讲的是 Megatron-LM 的做法。

Embedding

词表嵌入

首先将每个模型 tokenize ，变成一段 $(n,1)$ 向量，然后再在 $(N,d)$ 词表寻找，得到一个 $(n,d)$ embedding。

这个时候，假设有 $G$ 张显卡，每张显卡负责 $\frac{N}{G}$ 个单词，分得 $(\frac{N}{G},d)$ 的矩阵。

对于每个 GPU 来说，判断每个 token_id :

• 判断它是否属于当前 GPU 负责的词表区间；
• 如果属于，就用本地 embedding 表查出来；
• 如果不属于，就置为 0；
• 最后对 TP group 做 all-reduce sum，得到完整 embedding。

最终，每张显卡得到一个 $(n,d)$ 的嵌入。

位置编码

Position embedding 通常不按词表切分，一般是每个 rank 保存一份，或者在本地计算后加到 hidden states 上。

如果是 RoPE，也是在每个 GPU 本地对自己的 attention heads 应用，不需要额外通信。

Self-Attention

多头注意力

多头注意力天生适合用来做 TP ，对于第 $i$ 个查询头来说，他只需要关注

$$O_i = Attention(Q_i, K_i, V_i)$$

我们把 QKV 合并成一个大矩阵,并按照输出维度切分

$$W_QKV: [H, 3H]\\ W_QKV_i: [H, 3H/t]$$

每个 GPU 负责的头数即 local_num_heads = num_heads / t

随后，每个 head 得到一个输出：

$$H_i \in \mathbb{R}^{S\times\frac{d}{t}}$$

拼接起来：

$$H=\operatorname{Concat}(H_1,H_2,\ldots,H_{32}), H\in\mathbb{R}^{S\times d}$$

然后经过

$$Y=HW_O,W_O\in\mathbb{R}^{d\times d}$$

这个 $W_O$ 叫 Output Projection。

最后和输入进行残差连接

$$X=X+Y$$

单头注意力

只能对隐藏层 $d$ 进行切分，如果每张卡保留Q，则分片意义不大，如果 QKV 都做切分，则计算的时候还需要聚合，没有什么收益。

MLP

对于一个 MLP 层，从最简单的 RELU 开始，然后我们逐步推广到目前的 SwiGLU

Megatron 提供拆分方法如下，我们结合 FFN 层实际来说

非门控函数

考虑

$$Y = \operatorname{act}(XA)B$$ $$A = [A_1, A_2, \dots, A_t]$$

将 $B$ 沿输入维切分：

$$B = \begin{bmatrix} B_1 \\ B_2 \\ \vdots \\ B_t \end{bmatrix}$$

每个 rank 本地计算：

$$Y_i = \operatorname{act}(XA_i)B_i$$

最后对所有 rank 的结果做 all-reduce sum：

$$Y = \sum_{i=1}^{t} Y_i$$

门控函数

$$Y = \operatorname{act}(XA_u,XA_g)B$$

对于门控函数，切分思路类似。gate projection 和 up projection 都沿输出维切分：

$$A_g = [A_{g,1}, A_{g,2}, \dots, A_{g,t}]$$ $$A_u = [A_{u,1}, A_{u,2}, \dots, A_{u,t}]$$

down projection 仍然沿输入维切分：

$$B = \begin{bmatrix} B_1 \\ B_2 \\ \vdots \\ B_t \end{bmatrix}$$

每个 rank 本地计算：

$$Y_i = \left( \operatorname{XLU}(XA_{g,i}) \odot XA_{u,i} \right) B_i$$

最后同样通过 all-reduce sum 得到完整输出：

$$Y = \sum_{i=1}^{t} Y_i$$

Cross Entropy