手写 mini-sglang 阶段四 · 更快、更大

张量并行

模型放不进一张卡，或者一张卡的带宽撑不起想要的延迟，就要把每一层的权重切开分到多张卡上：这就是张量并行（TP）。前面各章的线性层、嵌入层、KV 池里那些 get_tp_info().size 在这一章终于派上用场。我们在 CPU 上用 gloo 做通信，TP=2、TP=4 的输出与单卡的 Hugging Face 完全一致。

自测：能答上来就可以跳过本章

1. 一个 decoder 层做张量并行，需要几次 all-reduce？分别在哪里？
2. 列并行和行并行分别切权重的哪一维？为什么 MLP 要"先列后行"？
3. Qwen3-0.6B 有 8 个 KV 头，TP=16 时每个 rank 怎么分 KV 头？
4. 词表并行的嵌入层用 all-reduce，输出层用 all-gather，为什么不同？

自测参考答案（先自己答，再展开对照）

1. 两次：注意力的 o_proj（行并行）之后一次，MLP 的 down_proj（行并行）之后一次。
2. 列并行按输出维切（每个 rank 算一部分输出），行并行按输入维切（每个 rank 得到部分和，需要 all-reduce）。先列后行：列并行输出的正好是行并行需要的那部分输入，中间的激活逐元素计算，不需要通信，整个 MLP 只要最后一次 all-reduce。
3. 8 个 KV 头不够分给 16 个 rank：每个 KV 头复制到 2 个 rank 上，每个 rank 放 1 个 KV 头（和它对应的 query 头）。
4. 嵌入层按词表切分：每个 rank 只查自己那段词表里的 token，查不到的输出 0，all-reduce 求和就得到完整的嵌入；输出层按词表切分时每个 rank 算出一部分词表的 logits，要 all-gather 拼成完整的词表分布才能采样。

本章要回顾的文件：layers/linear.py、layers/embedding.py、models/weight.py 中的 shard_tensor、distributed/、engine/engine.py 中的 _init_communication。

到这一章为止的文件：73 个（新建 3，修改 0，提前引入 0）

minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/  （4 个文件，不变）
├── distributed/
│   ├── __init__.py  ← 本章新建
│   ├── info.py  ← 本章新建
│   ├── impl.py  ← 本章新建
├── engine/  （5 个文件，不变）
├── kernel/  （2 个文件，不变）
├── kvcache/  （5 个文件，不变）
├── layers/  （8 个文件，不变）
├── llm/  （2 个文件，不变）
├── message/  （5 个文件，不变）
├── models/  （7 个文件，不变）
├── moe/  （2 个文件，不变）
├── scheduler/  （9 个文件，不变）
├── server/  （4 个文件，不变）
├── tokenizer/  （4 个文件，不变）
├── utils/  （8 个文件，不变）

这一步的 main：examples/ch16_tp.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

动画：张量并行怎样切一层、通信几次（约 1.5 分钟）

Megatron 式切分

一个 decoder 层有两个"矩阵乘夹着逐元素运算"的结构：注意力（qkv_proj → 注意力 → o_proj）和 MLP（gate_up_proj → SiLU × up → down_proj）。Megatron-LM 的切法（原理见推理系统手册的张量并行一章）：

• 第一个矩阵按输出维切（列并行）：每个 rank 算出一部分输出通道。注意力按头切：每个 rank 负责若干个头，头与头之间的注意力计算互不相关，所以注意力本身不需要任何通信；MLP 按中间维切，SiLU × up 是逐元素的，也不需要通信。
• 第二个矩阵按输入维切（行并行）：每个 rank 用自己那部分中间结果乘自己那部分权重，得到完整形状的部分和，最后一次 all-reduce 相加。

所以每层两次 all-reduce（o_proj 和 down_proj 之后），其余全是本地计算。

张量并行下一个 decoder 层的数据流

minisgl/layers/linear.py · LinearQKVMerged

class LinearQKVMerged(_LinearTPImpl):
    """合并的 QKV 投影，按注意力头切分；KV 头数少于 TP 数时，KV 头在多个 rank 上复制。"""

    def __init__(self, hidden_size: int, head_dim: int, num_qo_heads: int, num_kv_heads: int,
                 has_bias: bool):
        tp_size = get_tp_info().size
        local_qo = div_even(num_qo_heads, tp_size)
        local_kv = div_even(num_kv_heads, tp_size, allow_replicate=True)
        full_osize = (num_qo_heads + 2 * num_kv_heads) * head_dim
        local_osize = (local_qo + 2 * local_kv) * head_dim
        super().__init__(hidden_size, full_osize, hidden_size, local_osize, has_bias)

minisgl/layers/linear.py · LinearRowParallel

class LinearRowParallel(_LinearTPImpl):
    """按输入维切分的行并行：每个 rank 算出部分和，再 all-reduce。"""

    def __init__(self, input_size: int, output_size: int, has_bias: bool):
        tp_info = get_tp_info()
        self._comm = DistributedCommunicator()
        self._tp_size = tp_info.size
        local_isize = div_even(input_size, tp_info.size)
        super().__init__(input_size, output_size, local_isize, output_size, has_bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.linear(x, self.weight, self.bias)
        return self._comm.all_reduce(y) if self._tp_size > 1 else y

合并的 qkv_proj 在切分时要小心：不能把合并后的大矩阵直接切成两半（那会把 q 的全部分给 rank 0、k 和 v 分给 rank 1），而要先分别切 q、k、v，再在每个 rank 上拼接。流式加载器正是这么做的：先 shard_tensor，再合并（第 3 章）。

KV 头不够分时复制

minisgl/models/weight.py · shard_tensor

def shard_tensor(key: str, value: torch.Tensor, rank: int, size: int,
                 num_kv_heads: int) -> torch.Tensor:
    """取出第 rank 个分片。"""
    if size == 1:
        return value
    if any(s in key for s in _SPLIT_DIM_0):
        is_kv = ".k_proj" in key or ".v_proj" in key
        if is_kv and num_kv_heads < size:  # KV 头比 rank 少：几个 rank 共用（复制）同一个 KV 头
            head_dim = value.shape[0] // num_kv_heads
            head = rank * num_kv_heads // size
            return value[head * head_dim:(head + 1) * head_dim].clone()
        return value.chunk(size, dim=0)[rank].clone()
    if any(s in key for s in _SPLIT_DIM_1):
        return value.chunk(size, dim=1)[rank].clone()
    if "lm_head" in key or "embed_tokens" in key:  # 词表并行：按行切，最后一段可能短一些
        per = div_ceil(value.shape[0], size)
        return value[rank * per:min((rank + 1) * per, value.shape[0])].clone()
    return value  # 其余（各种 norm、MoE 路由）每个 rank 一份完整的

GQA 模型的 KV 头比 q 头少。Qwen3-0.6B 有 16 个 q 头、8 个 KV 头：TP=2 时每个 rank 8 个 q 头、4 个 KV 头；TP=16 时每个 rank 1 个 q 头，而 KV 头只有 8 个——每两个 rank 共用同一个 KV 头（div_even(..., allow_replicate=True) 返回 1），各自保存一份副本。这会浪费一些显存，但保证了注意力仍然不需要通信。

词表并行

嵌入层的词表按行切，每个 rank 保存一段；查表时不在本段的 token 填 0，然后 all-reduce 相加——每个 token 恰好在一个 rank 上查到真实的向量。输出层同样按词表切，每个 rank 算出自己那段词表的 logits；采样需要完整的分布，所以用 all-gather 把各段拼起来（第 2 章的 ParallelLMHead.forward）。前者是"各 rank 的结果相加"，后者是"各 rank 的结果拼接"。

运行

examples/ch16_tp.py

"""第 16 章：张量并行下每个 rank 拿到的权重形状；两个进程（gloo）的前向与单卡对比，并统计通信次数。"""

import multiprocessing as mp

import torch

PATH = "models/Qwen3-0.6B"
PROMPT_IDS = [785, 6722, 315, 9625, 374]  # "The capital of France is"


def show_shapes() -> None:
    import safetensors
    from minisgl.models import shard_tensor

    names = ["model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight",
             "model.layers.0.self_attn.o_proj.weight", "model.layers.0.mlp.gate_proj.weight",
             "model.layers.0.mlp.down_proj.weight", "model.embed_tokens.weight",
             "model.layers.0.input_layernorm.weight"]
    with safetensors.safe_open(f"{PATH}/model.safetensors", "pt") as f:
        print(f"{'权重':<42} {'完整':<14} {'TP=2 每个 rank':<16} TP=16 每个 rank")
        for n in names:
            w = f.get_tensor(n)
            s2 = tuple(shard_tensor(n, w, 0, 2, num_kv_heads=8).shape)
            s16 = tuple(shard_tensor(n, w, 0, 16, num_kv_heads=8).shape)
            print(f"{n.removeprefix('model.'):<42} {str(tuple(w.shape)):<14} {str(s2):<16} {s16}")


def run(rank: int, size: int, out: mp.Queue) -> None:
    from minisgl.core import Batch, Req, SamplingParams
    from minisgl.distributed import DistributedCommunicator, DistributedInfo
    from minisgl.engine import Engine, EngineConfig

    engine = Engine(EngineConfig(model_path=PATH, tp_info=DistributedInfo(rank, size), dtype=torch.float32,
                                 max_running_req=2, num_page_override=256, max_seq_len_override=128,
                                 distributed_port=29600))
    counts = {"all_reduce": 0, "all_gather": 0}
    impl = DistributedCommunicator.plugins[-1]
    for op in counts:
        fn = getattr(impl, op)

        def wrapped(x, fn=fn, op=op):
            counts[op] += 1
            return fn(x)

        setattr(impl, op, wrapped)
    n = len(PROMPT_IDS)
    engine.page_table[0, :n] = torch.arange(n)
    req = Req(input_ids=torch.tensor(PROMPT_IDS, dtype=torch.int32), table_idx=0, cached_len=0,
              output_len=4, uid=0, sampling_params=SamplingParams(), cache_handle=None)
    batch = Batch(reqs=[req], phase="prefill")
    batch.padded_reqs = batch.reqs
    batch.input_ids, batch.positions = req.input_ids.clone(), torch.arange(n, dtype=torch.int32)
    batch.out_loc = engine.page_table[0, :n]
    engine.attn_backend.prepare_metadata(batch)
    with engine.ctx.forward_batch(batch):
        logits = engine.model.forward()
    qkv = engine.model.model.layers.op_list[0].self_attn.qkv_proj.weight
    out.put((rank, logits[0].numpy().copy(), tuple(qkv.shape), counts, tuple(engine.kv_cache.k_cache(0).shape)))
    engine.shutdown()


if __name__ == "__main__":
    show_shapes()
    mp.set_start_method("spawn")
    results = {}
    for size in (1, 2):
        q: mp.Queue = mp.Queue()
        procs = [mp.Process(target=run, args=(r, size, q)) for r in range(size)]
        for p in procs:
            p.start()
        results[size] = sorted((q.get() for _ in range(size)), key=lambda x: x[0])
        for p in procs:
            p.join()
    base = torch.from_numpy(results[1][0][1])
    for rank, logits, qkv_shape, counts, k_shape in results[2]:
        print(f"TP=2 rank {rank}: qkv_proj {qkv_shape}，KV 池每层 {tuple(k_shape)}，"
              f"一次前向 all_reduce {counts['all_reduce']} 次、all_gather {counts['all_gather']} 次，"
              f"logits 与 TP=1 最大误差 {(torch.from_numpy(logits) - base).abs().max().item():.2e}")

运行结果

权重                                         完整             TP=2 每个 rank     TP=16 每个 rank
layers.0.self_attn.q_proj.weight           (2048, 1024)   (1024, 1024)     (128, 1024)
layers.0.self_attn.k_proj.weight           (1024, 1024)   (512, 1024)      (128, 1024)
layers.0.self_attn.o_proj.weight           (1024, 2048)   (1024, 1024)     (1024, 128)
layers.0.mlp.gate_proj.weight              (3072, 1024)   (1536, 1024)     (192, 1024)
layers.0.mlp.down_proj.weight              (1024, 3072)   (1024, 1536)     (1024, 192)
embed_tokens.weight                        (151936, 1024) (75968, 1024)    (9496, 1024)
layers.0.input_layernorm.weight            (1024,)        (1024,)          (1024,)
TP=2 rank 0: qkv_proj (2048, 1024)，KV 池每层 (257, 1, 4, 128)，一次前向 all_reduce 57 次、all_gather 1 次，logits 与 TP=1 最大误差 1.14e-05
TP=2 rank 1: qkv_proj (2048, 1024)，KV 池每层 (257, 1, 4, 128)，一次前向 all_reduce 57 次、all_gather 1 次，logits 与 TP=1 最大误差 1.14e-05

- 权重形状：TP=2 时 q_proj、k_proj、gate_proj 按输出维减半，o_proj、down_proj 按输入维减半，词表减半，norm 不切；TP=16 时 k_proj 仍有 128 行——一个完整的 KV 头（复制）。 - 两个 rank 的 qkv_proj 都是 (8 + 2 × 4) × 128 = 2048 行；每个 rank 的 KV 池只存 4 个 KV 头。 - 一次前向 57 次 all-reduce（28 层 × 2 + 嵌入层 1）、1 次 all-gather（输出层），与分析一致。 - logits 与单卡的最大误差 1e-5 量级：all-reduce 改变了浮点加法的顺序，结果不再逐位相同，但贪心解码的 token 完全一致（本章的测试在服务层面验证了 TP=2、TP=4）。

通信的实现

minisgl/distributed/impl.py · TorchDistributedImpl

class TorchDistributedImpl:
    """用 torch.distributed 做集合通信。GPU 上走 NCCL，CPU 上走 gloo。"""

    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
        if dist.get_world_size() == 1:
            return x
        dist.all_reduce(x, op=dist.ReduceOp.SUM)
        return x

    def all_gather(self, x: torch.Tensor) -> torch.Tensor:
        """沿第 0 维拼接各 rank 的张量：[n, ...] -> [world_size * n, ...]。"""
        world_size = dist.get_world_size()
        if world_size == 1:
            return x
        x = x.contiguous()
        out = x.new_empty((world_size * x.shape[0],) + tuple(x.shape[1:]))
        if dist.get_backend() == "gloo":  # gloo 不支持 all_gather_into_tensor
            dist.all_gather(list(out.chunk(world_size, dim=0)), x)
        else:
            dist.all_gather_into_tensor(out, x)
        return out

DistributedCommunicator 通过插件列表选择实现，默认用 torch.distributed：GPU 上走 NCCL，CPU 上走 gloo。gloo 不支持 all_gather_into_tensor，我们在 CPU 上改用 all_gather 写进输出张量的各个切片。

引擎初始化时建立进程组：

minisgl/engine/engine.py · Engine._init_communication

def _init_communication(self, config: EngineConfig) -> Any:
    """TP>1 时建立进程组。GPU 上张量走 NCCL，另建一个 gloo 组做 CPU 侧的同步；CPU 上全用 gloo。"""
    if config.tp_info.size == 1:
        return None
    backend = "nccl" if is_cuda(self.device) else "gloo"
    dist.init_process_group(
        backend=backend,
        rank=config.tp_info.rank,
        world_size=config.tp_info.size,
        timeout=timedelta(seconds=config.distributed_timeout),
        init_method=config.distributed_addr,
    )
    return dist.new_group(backend="gloo") if backend == "nccl" else dist.group.WORLD

GPU 上另建一个 gloo 组，专门用于 CPU 上的控制信息（第 14 章的"广播条数"、启动时的显存同步）。

与官方的差异：PyNCCL

官方默认（use_pynccl=True）不用 torch.distributed 做 GPU 通信，而是通过 tvm-ffi 直接调用 NCCL（kernel/csrc/src/pynccl.cu），并预先分配一块通信缓冲区。好处是开销更小，且与 CUDA Graph 的配合更可控。我们只实现了 torch.distributed 这一种，PyNCCL 留作练习（需要 GPU）。

官方实现

• 线性层：layers/linear.py · LinearQKVMerged（第 71～88 行）、layers/linear.py · LinearRowParallel（第 109～127 行）
• 切分：models/weight.py · _shard_tensor（第 34～52 行）
• 通信：distributed/impl.py · DistributedCommunicator（第 63～70 行）
• 进程组：engine/engine.py · Engine._init_communication（第 112～137 行）

测试

tests/test_ch16_tp.py · test_tensor_parallel_matches_hf

@pytest.mark.parametrize("tp", [2, 4])
def test_tensor_parallel_matches_hf(tp):
    """CPU 上用 gloo 做张量并行：TP=2、TP=4 的贪心输出与单卡 HF 一致。"""
    proc, port = start_server("--tp", str(tp), "--num-pages", "2048")
    try:
        tok = AutoTokenizer.from_pretrained(QWEN3)
        refs = hf_greedy(QWEN3, PROMPTS[:3], 10)
        for prompt, ref in zip(PROMPTS[:3], refs):
            r = httpx.post(f"http://127.0.0.1:{port}/v1/chat/completions",
                           json={"model": "m", "prompt": prompt, "max_tokens": 10, "temperature": 0,
                                 "ignore_eos": True}, timeout=120).json()
            assert r["choices"][0]["message"]["content"] == tok.decode(ref)
    finally:
        os.killpg(proc.pid, signal.SIGKILL)

每个 TP 配置起一个完整的服务（多个调度器进程、tokenizer、API Server），3 个提示词的贪心输出与单卡 Hugging Face 逐字相同。这同时验证了第 14 章的多 rank 消息同步。

怎么讲清楚

讲张量并行：Megatron 切法——注意力按头切（qkv 列并行、o_proj 行并行），MLP 先列并行（gate/up 按中间维切）再行并行（down），每个 decoder 层只需要两次 all-reduce（注意力后一次、MLP 后一次）。合并的 qkv 要先分别切 q、k、v 再拼接，否则各 rank 拿到的头不对；KV 头数少于 TP 数时（Qwen3-0.6B 有 8 个 KV 头、TP=16）每个 KV 头复制到两个 rank。词表并行：嵌入层每个 rank 只查自己那段词表、查不到的置零，all-reduce 求和；输出层各 rank 算自己那段词表的 logits，all-gather 拼接。各 rank 与单卡只差浮点求和顺序带来的微小误差。

练习

1. 计算 Llama-3.1-70B（80 层、hidden 8192、bf16）在 TP=8、每轮 decode 批大小 64 时，每层两次 all-reduce 的数据量。NVLink 带宽 450 GB/s（单向）时，通信大约占多少时间？
2. 如果 TP 数不能整除 q 头数（例如 Qwen2.5-0.5B 有 14 个 q 头、TP=4），会在哪里报错？怎样支持？
3. 实现 PyNCCL 插件：用 ctypes 加载 libnccl.so，实现 all_reduce；在 GPU 上与 torch.distributed 的结果比较。

参考答案

1. 每次 all-reduce 的数据是 [64, 8192] 个 bf16，1 MiB；ring all-reduce 每个 rank 发送约 2 × (8−1)/8 × 1 MiB ≈ 1.75 MiB，约 4 µs 的带宽时间，再加上每次 all-reduce 几微秒的固定延迟。80 层 × 2 次 ≈ 160 次，合计约 1～2 ms——小批量时固定延迟占主导，这就是 decode 阶段 TP 规模不宜太大的原因。
2. LinearQKVMerged 里的 div_even(num_qo_heads, tp_size) 断言失败。可以给头数补齐（padding 出空头），或者只支持能整除的 TP 数；正式的引擎通常直接要求整除。
3. 参考官方 kernel/pynccl.py 与 SGLang 的 pynccl_wrapper.py：rank 0 用 ncclGetUniqueId 生成 ID，通过 CPU 进程组广播；各 rank ncclCommInitRank；all_reduce 调用 ncclAllReduce，stream 用当前的 CUDA stream。

小结

• Megatron 切法：先列并行（按头 / 按中间维）、再行并行，每层两次 all-reduce。
• 合并的 qkv 要先分别切再拼接；KV 头不够分时在多个 rank 上复制。
• 词表并行：嵌入层 all-reduce 求和，输出层 all-gather 拼接。
• 各 rank 的输出与单卡只差浮点求和顺序带来的微小误差，贪心解码结果一致。

本章练习题（1 道）

在浏览器里直接写代码、跑测试；也可以在本地用 python practice/judge.py 判题。

• 199. 张量并行的权重切分（含 KV 头复制）（中等）

发现错误或有更好的写法？在 GitHub 上编辑此页 · 提交问题