S17分布式服务·216 行

张量并行与流水线并行

先按列切、再按行切每个 matmul，每个子层只需要一次 all-reduce。流水线并行则是用延迟换容量。

按列/按行切分
all-reduce
按头切分
流水线阶段
micro-batch

为什么重要

问题

70B 模型 fp16 是 140 GB。H100 有 80 GB，而且你还得给 KV 缓存留地方。量化能帮你降一档模型规模；再往上，模型就必须住在多张 GPU 上了。

切分很容易。难的是切分之后别把收益全花在通信上，而这取决于 GPU 之间连着什么。

核心思路

解法

两种策略，成本结构不同。

张量并行把每个矩阵切到多张 GPU 上。每张 GPU 都参与每一层，模型因此不仅变小，也变快了。代价是每一层都要做一次集合通信，才能把输出重新拼出来。
流水线并行给每张 GPU 分配一段连续的层。通信量降到每个阶段边界一份激活，几乎可以忽略。但请求必须按顺序穿过各阶段，延迟不会改善，还会出现空泡。

张量并行之所以便宜，靠的是 Megatron 的切分模式：第一个 matmul 按列切，第二个按行切，整个子层恰好只需要一次 all-reduce。

Megatron 模式python

class ShardedFFN:
    def __init__(self, w1, w2, w3, rank, world):
        # W1 and W3: split the OUTPUT dimension -> each rank owns d_ff/world columns
        self.w1 = shard_columns(w1, rank, world)
        self.w3 = shard_columns(w3, rank, world)
        # W2: split the INPUT dimension -> each rank owns d_ff/world rows
        self.w2 = shard_rows(w2, rank, world)

    def forward(self, x):
        h = silu(x @ self.w1) * (x @ self.w3)   # sharded activation, no comms
        y = h @ self.w2                          # PARTIAL sum of the full output
        return all_reduce(y)                     # one collective per sublayer

注意力按头切分，而 GQA 给并行度设了上限

每张 GPU 拿走一部分注意力头，独立计算；输出投影按行切分，和 W2 一样。不过张量并行度必须整除 KV 头数。一个有 8 个 KV 头的模型，不复制 KV 头就用不了 TP16，搭建 Llama-3-70B 16 卡部署的人经常没料到这一点。

工作原理

示意图张量切分与流水线阶段

张量并行买到的是低延迟、付出的是带宽。流水线并行买到的是容量、付出的是延迟。互联决定了你负担得起哪一种。

把通信成本算清楚

N 张 GPU 上的 ring all-reduce，每张 GPU 要搬运 2(N−1)/N 倍的载荷。这里的载荷是激活：batch × d_model × 2 字节。每层有两次 all-reduce。

对一个 80 层、d=8192 的 70B 模型，batch 16 时每步约 70 MB。在单方向 450 GB/s 的 NVLink 上（H100 标称的 900 GB/s 是双向合计），这大约花 0.16 ms，相对几毫秒的计算可以忽略。batch 到 512 时载荷涨 32 倍，在 PCIe 或以太网上它会完全占主导。

于是就有了所有人最终都收敛到的那条规则：节点内用张量并行，节点间用流水线并行。

数据并行不是一回事

把整个模型在多张 GPU 上各复制一份，再在副本间做负载均衡，这叫数据并行。它以零通信换来吞吐倍增。但模型装不下时它帮不上忙，单请求延迟它也降不下来。

真实部署会把三者叠起来：节点内用 TP 让模型装得下、每个 token 更快；如果还装不下就跨节点加 PP；再往上用 DP 副本堆吞吐。最前面还有一个路由器，理想情况下它能保持前缀缓存的局部性。

动手观察

动手试试

模拟器并行成本模型先设 TP8，然后把互联切成 PCIe

模型规模70B

张量并行

流水线并行

互联带宽

batch16

GPU 数: 4
每 GPU 权重: 32.6 GB
每步耗时: 9.80 ms
扩展效率: 99%

一次 decode 的时间去向

compute

读取权重（有效计算）
all-reduce（张量并行的税）
流水线空泡（阶段空转）

每层两次 all-reduce，每步共计 60.0 MB。每一层的输出都必须先在所有 TP rank 之间求和，下一层才能开始。这是一次阻塞的、对延迟敏感的集合通信，而不是后台传输。

扩展曲线 · 每步耗时 vs TP 度数

TP1 · 38.9ms

TP2 · 19.5ms

TP4 · 9.8ms

TP8 · 4.9ms

算力每翻倍一次就减半，通信却不会。越过交叉点之后，加 GPU 只会让每个 token 更慢。切换互联带宽就能找到它。

流水线空泡

阶段0

没有流水线并行：只有一个阶段，没有空泡。把 PP 调大就能看到填充和排空的代价。

选一个 70B 模型，NVLink、batch 16、TP8：效率保持很高，步时下降接近 8 倍。现在把 batch 提到 512 并切到 PCIe。all-reduce 那一段迅速膨胀，直到 TP1 反过来胜过 TP8，扩展曲线面板会把这个交叉点画出来。之后再调高 PP：空泡出现，随着 micro-batch 增多又缩了回去。

亲手实现

实现

切分注意力块时有一个约束，值得写进断言。

code/s17_parallelism.py（节选）python

class ShardedAttention:
    def __init__(self, w, rank, world, n_heads, n_kv_heads):
        assert n_heads % world == 0, "TP degree must divide query heads"
        assert n_kv_heads % world == 0, (
            f"TP{world} needs n_kv_heads divisible by {world}, got {n_kv_heads}. "
            "Either lower the TP degree or replicate KV heads across ranks."
        )
        self.local_heads = n_heads // world
        self.local_kv_heads = n_kv_heads // world

        h = slice(rank * self.local_heads, (rank + 1) * self.local_heads)
        self.wq = w.wq[:, h]                  # column shard
        self.wk, self.wv = shard_kv(w, rank, world)
        self.wo = w.wo[h, :]                  # row shard -> needs all-reduce

每个 rank 必须采样出同一个 token

在精确算术下，最后一次 all-reduce 之后每个 rank 手里的 logits 完全相同。但浮点归约顺序在各 rank 之间可能差几个 ULP，各自独立采样时，接近平局的情况下就可能选出不同的 token。此后各 rank 生成的是不同的序列，KV 缓存也彼此不一致，输出在几百个 token 的尺度上缓慢劣化，而且哪里都不报错。要么在 rank 0 上采样再广播，要么让所有 rank 用完全相同的种子，并接受平局带来的风险。

在本地运行

在单进程内模拟张量并行与流水线并行，断言切分后的前向与单卡结果完全一致，并在不同互联带宽和 batch 大小下报告一个通信成本模型。

$ python code/s17_parallelism.py

预期输出： TP1–TP11 下切分与未切分路径的精确数值一致、一次演示 GQA 整除约束的断言失败，以及一张表：NVLink 上 TP8 获胜，而以太网 batch 512 时 TP8 输给 TP1。

只需要 NumPy — 查看环境准备.

生产实践

在生产环境中

vLLM —— tensor_parallel_size 和 pipeline_parallel_size；MultiProcExecutor 为每张 GPU 拉起一个 worker 进程，通过共享内存消息队列协调。
NCCL —— 底下那层集合通信库，带拓扑感知的 ring 和 tree 算法。
序列并行 / ring attention —— 第三个维度：把一条序列切到多张 GPU 上，让 FlashAttention 的滚动状态沿环传递。只有极长上下文才需要。
专家并行 —— 针对 MoE（S16），按专家而不是按张量切分。all-reduce 随之变成 all-to-all。

练习

1
算出你自己模型的 all-reduce 数据量，找出 TP8 不再胜过 TP4 的那个互联带宽。
2
实现「在 rank 0 采样再广播」的修法，然后刻意破坏它，测量各 rank 要多少个 token 才会出现可见的分歧。
3
为一套混合部署建模：节点内 TP8、跨两个节点 PP2、再往上 DP4。算出 GPU 总数、每卡显存和预期吞吐，并与纯 TP16 对比。

继续学习

接下来

现在你能把模型切到多张 GPU 上了。最后一个结构性想法是改为切分负载：S18 让 prefill 和 decode 跑在完全不同的机器上，因为这两个阶段想要的硬件正好相反。

自测

习题

先作答，再看解析。答错比答对更有价值，因为解析会指出你该回头重读哪一部分。

自测第 1 题 / 共 6

在 Megatron 模式中，为什么 FFN 的第一个 matmul 按列切、第二个按行切？

得分 0/6