课程 05 · 并行

两张卡，一次前向

nano-vllm 的张量并行只有 156 行 linear.py 加一个共享内存通道——但它把 TP 的核心取舍全部暴露出来了：切什么、在哪同步、怎么传指令。

约 55 分钟 源码：layers/linear.py · layers/embed_head.py · model_runner.py 前置：04

把维度变化拆到每一步

1. 先确认数学权重[in,out]与PyTorch weight[out,in]的方向。
2. 输出列分片需要拼接；输入行分片的同形partial需要相加。

打开元素级图解：矩阵、轴、索引与可点击反例 →

1为什么要切

一张卡放不下 70B 模型时，最直接的办法是把权重切开。但切开的方式不止一种，代价也不一样：

方式	切什么	通信时机	适合
张量并行 TP	每层的权重矩阵	每层都要通信	单机多卡、低延迟
流水并行 PP	按层切成若干段	只在段边界通信	跨机、模型极深
专家并行 EP	MoE 的专家分给不同卡	token 路由时 all-to-all	MoE 模型

nano-vllm 选 TP：每张卡都跑完整的层序列，但每层只持有 1/N 的权重。好处是延迟低（不需要等别的卡的层），坏处是每层末尾都要同步一次。

2进程模型：一个主控 + N 个工人

llm_engine.py · L24–31ctx = mp.get_context("spawn")
for i in range(1, config.tensor_parallel_size):
    event = ctx.Event()
    process = ctx.Process(target=ModelRunner, args=(config, i, event))
    process.start()
    self.ps.append(process)
    self.events.append(event)
self.model_runner = ModelRunner(config, 0, self.events)   # rank 0 留在主进程

rank 0 不另起进程，它既当调度器又当一张卡。其它 rank 在 __init__ 末尾直接进入 loop()[L61–89]，等指令：

def loop(self):
    while True:
        method_name, args = self.read_shm()     # 阻塞等事件
        self.call(method_name, *args)
        if method_name == "exit": break

def read_shm(self):
    self.event.wait()                                       # ① 等事件
    n = int.from_bytes(self.shm.buf[0:4], "little")         # ② 读长度
    method_name, *args = pickle.loads(self.shm.buf[4:n+4])  # ③ 反序列化
    self.event.clear()                                      # ④ 复位
    return method_name, args

这条通道很有意思：用共享内存传数据，用 Event 传「有新数据」这个事实。没有 socket、没有 RPC 框架，1 MB 的 shared_memory 加一个 multiprocessing Event 就够了。

为什么序列要自定义 __getstate__

sequence.py · L72–83def __getstate__(self):
    last_state = self.last_token if not self.is_prefill else self.token_ids
    return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens,
            self.num_scheduled_tokens, self.block_table, last_state)

worker 侧只需要这 6 个字段。如果直接 pickle 整个 Sequence，token_ids 列表会跟着走——decode 时那是几千个 int，每条序列每步都传一遍，共享内存会被撑爆（1 MB），序列化开销也会盖过计算收益。自定义 __getstate__ 是分布式推理里最常见的优化之一。

3四种切法，四种权重加载

linear.py 的类不是随便分的，每个都对应一种「矩阵怎么切、权重怎么从 HF checkpoint 里取」。

类	切法	通信	典型用途
ReplicatedLinear	不切，每卡一份	无	小矩阵、bias
ColumnParallelLinear	按输出维切	无（输出天然分片）	QKV、gate/up
RowParallelLinear	按输入维切	all-reduce	attention 输出、down
VocabParallelEmbedding	按词表切	all-reduce（mask 后）	embedding

linear.py · L131–156class RowParallelLinear(LinearBase):
    def __init__(self, input_size, output_size, bias=False):
        tp_size = dist.get_world_size()
        super().__init__(divide(input_size, tp_size), output_size, bias, 1)

    def forward(self, x):
        y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None)
        if self.tp_size > 1:
            dist.all_reduce(y)          # 每卡算出部分和，加在一起
        return y

为什么 RowParallel 必须 all-reduce？因为它按输入维切：每张卡拿到的 x 是输入的一部分，算出的是完整输出的部分和。只有加起来才是真正的结果。而 ColumnParallel 按输出维切，每张卡的输出本来就是最终结果的一段，不需要通信。

4一次前向里的通信量

以 Qwen3-0.6B、2 卡为例，每层有两处 all-reduce：attention 的 o_proj 之后、MLP 的 down_proj 之后。每处传的是 [tokens, hidden] 的激活值。

• prefill 时 tokens 很多（比如 4096），通信量大但计算量更大，通信能被计算掩盖。
• decode 时 tokens 很少（比如 32），通信量小但延迟敏感——每层都要等 all-reduce 完成才能继续。

这解释了 TP 的一个反直觉现象：decode 阶段 TP 的加速比往往达不到线性，因为每步固定要做 2L 次 all-reduce，而计算量只有 O(bs)。这也是为什么大模型 decode 更偏好 EP 或 PD 分离。

5练习 A · 算一次 all-reduce 的代价

Qwen3-0.6B：28 层，hidden = 1024，激活用 bf16。TP=2，decode 时 batch = 16。每次 all-reduce 要传多少字节？一轮 decode 总共传多少？

展开：答案

每处 all-reduce 的张量是 [16, 1024] bf16 = 16 × 1024 × 2 = 32 KiB。

每层 2 处，28 层 → 56 处 → 56 × 32 KiB ≈ 1.75 MiB 每轮 decode。

按 4070 Ti 的 NVLink 缺失、走 PCIe 的情况（约 10 GB/s 有效），光通信就要约 0.18 ms——而 decode 一步的总时间通常在 10–30 ms 量级，所以 1% 左右，可以接受。但换成 4096 token 的 prefill，激活是 4096×1024 的百倍量级，通信就会变成瓶颈。

6练习 B · 闭卷测验

7练习 C · 代码寻宝

1. ColumnParallelLinear.weight_loader 为什么要 narrow(tp_dim, start_idx, shard_size)？

HF 的 checkpoint 存的是完整权重，每张卡只该拿自己那一片。按 tp_rank * shard_size 起 offset 取一段，就得到本 rank 的分片——注意切的是输出维（tp_dim=0）。

2. 为什么 RowParallelLinear 的 bias 只在 rank 0 上加？

bias 是完整向量，不是部分和。如果每张卡都加一遍，all-reduce 之后 bias 会被加 N 次。所以只在 rank 0 加，再让 all-reduce 把它摊开——这是「bias 只加一次」的标准做法。

3. VocabParallelEmbedding.forward 里 mask 的作用是什么？

每张卡只持有词表的一段。不在本段里的 token id 要屏蔽掉（置 0），否则会越界取到别的 embedding。屏蔽后 all-reduce，各卡的贡献拼成完整结果——这是词表并行的标准 pattern。

8面试复述卡

「TP 和 PP 怎么选？」

展开：30 秒版本

「TP 切的是每一层的矩阵，每层都要 all-reduce，所以对带宽敏感、适合单机 NVLink；它的优点是延迟低，所有卡同时在算同一批 token。PP 切的是层，卡与卡之间只有边界激活要传，通信少但会出现流水线气泡，适合跨机和超深模型。实际的大模型推理通常是 TP + PP + EP 组合：机内 TP、机间 PP、MoE 用 EP。」