DeepSeek V4 源码剖析

第15章 Tensor / Expert 并行：ColumnParallel 与 RowParallel

作者 杨艺韬 · 7,512 字 · 发布于 · 更新于

1.6T 的模型活在多台机器上，每一次 forward 都是一场需要对拍子的合奏。

V4 的并行策略不是事后补丁，而是写在每个 Linear 类里的"出生属性"。

15.1 引子：1.6T 模型的部署算式

V4 Pro 总参 1.6T，权重占多大？这个数不用估——权重被切成 64 个 safetensors 分片，把每个分片的字节数加起来就是：

64 个 model-000XX-of-00064.safetensors，合计 864,721,029,744 字节
= 864.7 GB = 805.3 GiB

（这些文件在镜像里是 LFS 指针，但指针文本里带 size 字段，所以尺寸是可核实的；张量清单则确实拿不到。）

反推一下每参数字节数：864.7e9 / 1.6e12 ≈ 0.54 字节 / 参数——非常接近 FP4 的理论值（0.5 字节本体 + 沿 K 每 32 个元素 1 字节 scale = 0.53125 字节 / 参数），说明 routed expert 确实占了参数量的绝大多数，走 FP8 的那部分（1 字节本体 + 128×128 一个 scale）只是零头。

805 GiB 意味着：

  • 单节点 8 卡 H100（80 GB × 8 = 640 GB）装不下——光权重就超了
  • 单节点 8 卡 H200（141 GB × 8 = 1128 GB）装得下——按 1/8 均摊每卡约 108 GB 权重；但参考实现里 wq_a、wkv、shared expert、Compressor 等不切分的权重每卡各存一份，lm_head、Compressor 权重还在加载时提升成 FP32，按 pin 版 inference/model.py 在 meta 设备上实例化，TP=8 时每卡约 117 GB，只剩二十来 GB 给 KV cache 与激活。按 max_seq_len=1048576 的部署设定，1M context 下每卡的 KV 约 9.6 GiB / 序列（见 §15.13），所以能跑，但并发上不去
  • 16 卡（2 × 8 卡 IB 互联）：均摊每卡约 54 GB 权重（同样口径实例化约 63 GB），H100 剩十几 GB、H200 剩七十多 GB——这才是舒服的起点
  • 32 卡 / 64 卡：给大 batch、长上下文或低延迟需求

H100 这一档，V4 Pro 的单节点部署是不成立的。

无论哪种部署，模型必须在张量层面切分——不能简单复制。V4 用两个核心机制：

  1. Tensor Parallel (TP)：把每层 Linear 的权重沿 row / column 切到不同卡
  2. Expert Parallel (EP)：把 384 个 routed expert 分配到不同卡，每卡持有一部分

V4 源码 (inference/model.py) 通过 world_size、rank 这对全局变量，把切分逻辑直接写在每个并行类里。本章拆这 4 个类。

15.2 全局并行状态

V4 在文件顶部声明：

# inference/model.py:15（节选）
world_size = 1
rank = 0
block_size = 128

其中 world_size、rank 在 Transformer.__init__ 被覆盖（block_size 是量化块大小，不变）：

def __init__(self, args: ModelArgs):
    global world_size, rank, default_dtype, scale_fmt, scale_dtype
    world_size = dist.get_world_size() if dist.is_initialized() else 1
    rank = dist.get_rank() if dist.is_initialized() else 0
    ...

V4 用 PyTorch 的 torch.distributed 做通信——world_size = 总 GPU 数、rank = 当前 GPU 编号。world_size 在所有进程里相同，rank 每个进程各不相同——inference/README.md 用 torchrun 启动，由 launcher 设置。

V4 让这两个变量是全局可变，而不是每个 module 持有自己的引用——这种"全局可变变量"的设计违反传统软件工程纪律，但在 V4 这种"模块嵌套深 + 配置一次贯穿全程"的场景下能大幅简化代码。

15.3 ColumnParallelLinear：输出维度切分

# inference/model.py:155（assert 提示信息从略）
class ColumnParallelLinear(Linear):
    """Shards output dim across TP ranks. No all-reduce needed on output."""
    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
        assert out_features % world_size == 0
        self.part_out_features = out_features // world_size
        super().__init__(in_features, self.part_out_features, bias, dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return linear(x, self.weight, self.bias)

ColumnParallelLinear 的语义：把 [in, out] 矩阵沿 out 维度切成 world_size 份。每个 rank 持有 [in, out / world_size] 的局部矩阵。

forward 时：每个 rank 独立计算局部输出 [B, S, out / world_size]——不需要通信。

输出维度被切了，怎么办？要么：

  • 后续操作能直接处理切分后的输出（如紧接 RowParallelLinear，正好把切分维度作为输入维度）
  • 或者外部显式 all_gather 重组

V4 的典型用法要看清楚哪些切、哪些不切：

  • Attention.wq_b（[q_lora_rank, n_heads × head_dim]）是 ColumnParallel——128 个 head 按 rank 切，8 卡时每 rank 拿 16 个
  • Attention.wo_a 也是 ColumnParallel（切 n_groups × o_lora_rank）
  • Indexer.wq_b 与 Indexer.weights_proj 是 ColumnParallel（切 index_n_heads=64 那一维）
  • Attention.wq_a 和 Attention.wkv 是普通 Linear——每个 rank 一份完整副本，不切

最后一条是 V4 的关键设计：wkv 输出的是 [B, S, head_dim=512] 的单份 KV（num_key_value_heads = 1，MQA 形态），所有 head 共享它。既然只有一份，就没什么可切的——每 rank 各存一份完整 KV，省掉了"KV 切分时的索引数学"，代价是 N 卡冗余 N 份 KV（详见 §15.13）。

15.4 RowParallelLinear：输入维度切分

# inference/model.py:166（assert 提示信息从略）
class RowParallelLinear(Linear):
    """Shards input dim across TP ranks. All-reduce on output to sum partial results."""
    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
        assert in_features % world_size == 0
        self.part_in_features = in_features // world_size
        super().__init__(self.part_in_features, out_features, bias, dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = linear(x, self.weight, None)
        if world_size > 1:
            y = y.float()
            dist.all_reduce(y)
        if self.bias is not None:
            y += self.bias
        return y.type_as(x)

RowParallelLinear 的语义：把 [in, out] 矩阵沿 in 维度切。每个 rank 持有 [in / world_size, out] 的局部矩阵。

forward 时：每个 rank 用局部 in 输入算局部输出，输出形状 [B, S, out]（完整 out 维度），但是数值上只是"部分和"。用 all_reduce 把所有 rank 的部分和加起来，得到完整输出。

注意 V4 的实现：先转 float32 再 all_reduce、加 bias、最后转回原 dtype。这是为了：

  • all_reduce 的精度敏感——不同 rank 的部分和量级可能不同，FP32 累加更稳
  • bias 在 reduce 之后加，避免被 reduce 重复 N 倍

ColumnParallel + RowParallel 的经典组合：

ColumnParallelLinear(in=D, out=H × heads)     # 输出按 head 维度切
   ↓ (输出: [B, S, head/world × D_h], 每个 rank 各持一部分 head)
ColumnParallel attention 内部计算
   ↓
RowParallelLinear(in=H × heads, out=D)        # 输入按 head 切，输出 reduce
   ↓ all_reduce 后得到完整 [B, S, D]

这套组合在 V4 的 Attention 类里直接可见——wq_b 是 ColumnParallel，wo_a 是 ColumnParallel，wo_b 是 RowParallel。

15.5 ColumnParallel + RowParallel 经典组合的张量流

V4 attention 内部的并行流转最典型地体现了这套组合：

flowchart LR
  X["x: [B,S,7168]<br/>每 rank 完整副本"] --> ColLin["wq_a（不切）→ ColumnParallel<br/>wq_b: 切 out_features"]
  ColLin --> Q["q: [B,S, n_heads/world × head_dim]<br/>每 rank 不同 head 切片"]
  Q --> Compute["attention 计算<br/>(rank 内独立)"]
  Compute --> O["o: [B,S, n_heads/world × head_dim]<br/>仍是 head 切片"]
  O --> RowLin["wo_a（Column）→ RowParallel<br/>wo_b: 切 in_features"]
  RowLin --> Partial["partial_output: [B,S,7168]<br/>每 rank 仅算了部分和"]
  Partial --> AllReduce[("all_reduce<br/>跨 rank 累加")]
  AllReduce --> Final["完整 x: [B,S,7168]<br/>每 rank 都有相同结果"]
  classDef parallel fill:#312e81,stroke:#a78bfa
  classDef sync fill:#7c2d12,stroke:#fb923c
  class ColLin,RowLin parallel
  class AllReduce sync

整段 attention 的主通信是最后一次 all_reduce——通信量 O(B×S×D)，远小于 attention 内部的 O(B×S×n_heads×head_dim) 计算量。唯一的例外是 30 个 compress_ratio=4 的层：Indexer 的 index_score 按 head 切在各 rank 上，求和前还要一次 dist.all_reduce(index_score)（inference/model.py:422-423），大小随上下文长度增长（§15.11）。

15.6 ParallelEmbedding：vocab 维度切分

# inference/model.py:83（节选）
class ParallelEmbedding(nn.Module):
    def __init__(self, vocab_size: int, dim: int):
        super().__init__()
        ...
        assert vocab_size % world_size == 0
        self.part_vocab_size = (vocab_size // world_size)
        self.vocab_start_idx = rank * self.part_vocab_size
        self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if world_size > 1:
            mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
            x = x - self.vocab_start_idx
            x[mask] = 0
        y = F.embedding(x, self.weight)
        if world_size > 1:
            y[mask] = 0
            dist.all_reduce(y)
        return y

ParallelEmbedding 的语义：把 vocab 切到不同 rank——rank 0 持有 token id [0, V/8)，rank 1 持有 [V/8, 2V/8) ......

forward 的关键技巧：

  • 每个 rank 算"自己负责的 token id"的 embedding——不属于本 rank 的 token id 被 mask 成 0
  • dist.all_reduce 把所有 rank 的部分 embedding 累加——因为每个 token id 只有一个 rank 真正算了，all_reduce 实际等于"取那个 rank 的输出"

这种"mask + all_reduce"的写法看起来浪费——为什么不直接用 dist.all_to_all 或者点对点通信？因为 TP 下每个 rank 后面都要用到全部 token 的完整 embedding，结果本来就得在所有 rank 上各有一份；all_reduce 一次完成"汇总 + 分发"，而 all_to_all 是给不同 rank 发不同数据，语义对不上。decode 时每 token 只有 7168 维，这次 all_reduce 的开销很小。

15.7 ParallelHead：lm_head 的 vocab gather

# inference/model.py:703
class ParallelHead(nn.Module):
    def __init__(self, vocab_size: int, dim: int, ...):
        super().__init__()
        ...
        self.part_vocab_size = (vocab_size // world_size)
        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32))

    def get_logits(self, x):
        return F.linear(x[:, -1].float(), self.weight)

    def forward(self, x, hc_fn, hc_scale, hc_base, norm):
        x = self.hc_head(x, hc_fn, hc_scale, hc_base)
        logits = self.get_logits(norm(x))
        if world_size > 1:
            all_logits = [torch.empty_like(logits) for _ in range(world_size)]
            dist.all_gather(all_logits, logits)
            logits = torch.cat(all_logits, dim=-1)
        return logits

ParallelHead 与 ParallelEmbedding 对偶：vocab 切到不同 rank，每个 rank 算"自己负责的部分 logits"，最后 all_gather 把所有 rank 的 logits 拼起来。

注意几个细节：

细节 1：只算最后一个 token 的 logits

return F.linear(x[:, -1].float(), self.weight)

x[:, -1] 取每条序列的最后一个 token——LLM 推理通常只关心下一个 token 的 logits。这避免了为序列中间 token 算 logits 的浪费。

细节 2：运行时的 weight 是 float32，但 checkpoint 里不是

源码这一行上面就挂着注释：

# lm_head in the checkpoint is stored in bf16, while the parameter here is
# stored in fp32 for easier computation of logits later.
self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32))

也就是说：磁盘上是 BF16，加载时提升到 FP32。区别很实际——算显存时按 BF16 算权重文件大小、按 FP32 算运行时占用。同样的模式还出现在 RMSNorm.weight 与 Compressor.wkv/wgate 上，注释都写了"checkpoint 里是 bf16"。

源码注释给的理由只是"方便后面算 logits"：get_logits 里 x 也先转成 float，整段 logits 都在 FP32 里算。

细节 3：all_gather 而非 all_reduce

vocab 在不同 rank 上是互不相交的——每个 rank 算的 logits 是 vocab 不同切片的真实结果，不是部分和。all_gather 把这些"真实切片"拼起来；如果用 all_reduce，会把不同切片错误相加。

15.8 Expert Parallel：384 expert 怎么切到 8 卡

V4 的 MoE 类里：

class MoE(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        ...
        assert args.n_routed_experts % world_size == 0
        self.n_routed_experts = args.n_routed_experts
        self.n_local_experts = args.n_routed_experts // world_size
        self.experts_start_idx = rank * self.n_local_experts
        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
        ...
        self.experts = nn.ModuleList([
            Expert(...) if self.experts_start_idx <= i < self.experts_end_idx else None
            for i in range(self.n_routed_experts)
        ])
        self.shared_experts = Expert(...)   # 每 rank 都有完整 shared expert

8 卡 TP 部署 V4 Pro：384 / 8 = 48 个 expert / rank。每 rank 的 nn.ModuleList 里：

  • 自己持有的 48 个位置是真实 Expert 实例
  • 其他 336 个位置是 None

forward 时：

# inference/model.py:635
for i in range(self.experts_start_idx, self.experts_end_idx):
    if counts[i] == 0:
        continue
    expert = self.experts[i]
    idx, top = torch.where(indices == i)
    y[idx] += expert(x[idx], weights[idx, top, None])

if world_size > 1:
    dist.all_reduce(y)

y += self.shared_experts(x)

每 rank 只跑自己持有的 48 个 expert，输出贡献到 y。最后 all_reduce 把所有 rank 的部分输出累加。

通信成本：每层 MoE 一次 all_reduce on [B*S, D]，而且注意源码里 y = torch.zeros_like(x, dtype=torch.float32)——这个 all_reduce 是 FP32 的，每元素 4 字节。对于 V4 Pro 在 1M context + batch=8 的极端假设下：

8 × 1048576 × 7168 × 4 bytes = 240 GB / 层

显然不可能——单层 all_reduce 传 240 GB，61 层就是十几 TB。

这说明的是：这段参考实现里的 all_reduce 只对得起"一次喂一小块"的用法。生产部署必须分块 prefill，把 B*S 压到几千的量级；decode 时 B*S 就是 batch 大小，[1, 7168] FP32 = 28 KB，完全不是问题。

第 16 章会深入这部分通信优化——DeepEP 给 MoE 提供的是"按 expert 路由 token"的 all-to-all，与这里"每 rank 算完全部 token 再求和"的 all_reduce 是两种不同的并行切法：前者是真正的 Expert Parallel（token 跑到 expert 那边去），后者更接近"expert 分片 + 结果求和"，通信量与 token 数成正比而与激活的 expert 数无关。参考实现选了后者，写起来只要一次 dist.all_reduce。

15.9 一段 Attention 的并行流转

把 V4 的 Attention 在 8 卡 TP 下的并行流转走一遍：

flowchart TB
  X["x: [B, S, 7168]<br/>每 rank 都有完整副本"] --> ColLin1["wq_a（普通 Linear，7168→1536）<br/>→ ColumnParallelLinear (wq_b)<br/>全量 [1536, 128×512]，每 rank 持 out 的 1/8"]
  ColLin1 --> Q["q: [B, S, 16 local heads, 512]<br/>128/8 = 16 个 head / rank"]
  X --> Lin["Linear (wkv)<br/>普通 Linear，每 rank 完整副本"]
  Lin --> KV["kv: [B, S, 512]<br/>每 rank 完整副本 (num_key_value_heads=1)"]
  Q --> SparseAttn["sparse_attn<br/>每 rank 只算自己的 head"]
  KV --> SparseAttn
  SparseAttn --> O["o: [B, S, 16 local heads, 512]"]
  O --> ColLin2["wo_a (ColumnParallel)<br/>全量 [128×512/16, 16×1024]<br/>每 rank 持 out 的 1/8 = 2048"]
  ColLin2 --> OLora["o_lora: [B, S, 2 local groups, 1024]<br/>o_groups=16, 16/8 = 2"]
  OLora --> RowLin["wo_b (RowParallel)<br/>全量 [16×1024, 7168]<br/>每 rank 持 in 的 1/8 = 2048"]
  RowLin --> Reduce["all_reduce"]
  Reduce --> XOut["x_out: [B, S, 7168]<br/>每 rank 都有完整结果"]

整个 attention 的通信主要是最后一次 all_reduce on [B, S, 7168]（带 Indexer 的层另有一次 index_score 的 all_reduce，图中未画）——其他都是 rank 内独立计算。这是 TP 在 attention 上的高效之处：通信量 O(B × S × D)，远小于参数量 O(D²)。

15.10 与 Megatron-LM 的对比

V4 的并行类与 Megatron-LM 的 ColumnParallelLinear / RowParallelLinear 有同源思路（Megatron-LM 是这套方案的工业化先驱）：

维度 Megatron-LM V4 inference/model.py
ColumnParallel 完整支持 完整支持
RowParallel 完整支持 完整支持
Embedding VocabParallelEmbedding ParallelEmbedding (同名不同实现)
LM Head ColumnParallelLinear（output_layer） ParallelHead (含 hc_head 处理)
Sequence Parallel 支持 不支持（推理不需要，见 §15.18）
Pipeline Parallel 支持 不支持（推理代码无 PP）
代码量 数千行 4 个类合计几十行

V4 的并行实现非常简洁——只覆盖推理需要的部分。训练时的 Pipeline Parallel、Sequence Parallel 等更复杂机制由训练框架（不公开的内部代码）处理，与 inference/model.py 解耦。

这种"训练 / 推理代码分开"的设计让公开的推理代码极其简洁——读者不需要被训练栈的复杂性淹没。

15.11 通信开销估算

把 V4 Pro 在 8 卡 TP 下的通信量估算一下（decode 阶段，batch=1，每 token）：

通信操作 张量大小 频率 总通信量 / token
Attention all_reduce（wo_b） [1, 7168] FP32 = 28 KB（RowParallelLinear 先 y.float()） 每层 1 次 × 61 层 1.7 MB
MoE all_reduce [1, 7168] FP32 = 28 KB（源码里 y 是 fp32） 每层 1 次 × 61 层 1.7 MB
LM Head all_gather 每 rank [129280/8] = [16160] FP32 = 63 KB 1 次，每 rank 收 8 份 505 KB
Embedding all_reduce [1, 7168] BF16 = 14 KB 1 次 14 KB
Indexer all_reduce（index_score） [1, T/4] BF16（T 为当前上下文长度），T=1M 时 512 KB 30 个 compress_ratio=4 层各 1 次 随 T 增长，T=1M 时 15 MB
总计 - - 约 4 MB / token + Indexer 那一项

（attention 那行的 FP32 来自 RowParallelLinear.forward 里的 y = y.float()，MoE 那行来自 y = torch.zeros_like(x, dtype=torch.float32)；Indexer 的 index_score 由 BF16 的 q、KV 与 weights_proj 算出，保持 BF16。）

H100 之间走的是 NVLink 4，单卡双向总带宽 900 GB/s（600 GB/s 是上一代 A100 的 NVLink 3，别拿错代次）。按 900 GB/s 做最乐观的估算，短上下文下约 4 MB 的纯传输时间在个位数微秒量级。

但这个数字只是下界，不该拿去做容量规划——真实的 all_reduce 开销由两部分主导：NCCL 每次 collective 的固定启动与同步开销（小消息上它比传输本身更贵），以及 61 层 × 2 次（再加 30 次 Indexer）同步点带来的串行等待。相对 decode 每 token 的整体耗时，这部分仍然不是大头；但"通信量 ÷ 带宽"这个算法本身只在大消息上才接近真实。

prefill 就完全是另一回事了：上表每一项都要乘以序列长度 S。单层单次 attention all_reduce 在 S = 1M 时就是 1M × 28 KB ≈ 30 GB，61 层两轮下来是 TB 量级；Indexer 的 index_score 在 prefill 时是 [1, S, S/4]，S = 1M 时单层就有 512 GiB。这就是为什么 1M 的 prefill 必须分块跑（参考实现 start_pos > 0 时只按单 token 写 KV，本身不支持分块 prefill），以及为什么第 16 章的 DeepEP 要为 prefill 单独准备一套高吞吐 kernel。

15.12 动手实验：跑通最小 TP 推理

# 启动 2 卡 TP（单机多进程；没有 GPU 时自动退回 CPU + gloo）
# torchrun --nproc-per-node=2 --master-port=29500 minimal_tp.py

import os
import torch
import torch.distributed as dist
import torch.nn as nn

use_cuda = torch.cuda.is_available()
dist.init_process_group(backend='nccl' if use_cuda else 'gloo')
world_size = dist.get_world_size()
rank = dist.get_rank()
if use_cuda:
    device = torch.device(f'cuda:{int(os.environ["LOCAL_RANK"])}')
    torch.cuda.set_device(device)
else:
    device = torch.device('cpu')

# 所有 rank 用同一个种子，生成同一份完整权重与输入，再各取自己的切片
torch.manual_seed(0)
W1 = torch.randn(256, 128, device=device)   # 完整 [out, in]
W2 = torch.randn(128, 256, device=device)
x = torch.randn(4, 128, device=device)      # TP 下输入在各 rank 是同一份

class ColParallel(nn.Module):
    def __init__(self, full_weight):
        super().__init__()
        out_dim = full_weight.size(0)
        assert out_dim % world_size == 0
        part = out_dim // world_size
        self.weight = nn.Parameter(full_weight[rank * part:(rank + 1) * part].clone())
    def forward(self, x):
        return x @ self.weight.T                  # 输出是 out 维的一个切片，不通信

class RowParallel(nn.Module):
    def __init__(self, full_weight):
        super().__init__()
        in_dim = full_weight.size(1)
        assert in_dim % world_size == 0
        part = in_dim // world_size
        self.weight = nn.Parameter(full_weight[:, rank * part:(rank + 1) * part].clone())
    def forward(self, x):
        y = x @ self.weight.T                     # 部分和
        dist.all_reduce(y)                        # 跨 rank 求和
        return y

with torch.no_grad():
    y = RowParallel(W2)(ColParallel(W1)(x))
    ref = x @ W1.T @ W2.T                         # 单卡完整计算
print(f"Rank {rank}: y shape = {tuple(y.shape)}, 与单卡结果一致: {torch.allclose(y, ref, rtol=1e-4, atol=1e-3)}")
dist.destroy_process_group()

跑这段代码，两个 rank 都会打印 与单卡结果一致: True——所有 rank 用同一个种子生成同一份完整权重与输入，各取自己的切片，ColumnParallel + RowParallel 经一次 all_reduce 得到的结果与单卡完整计算一致。没有 GPU 时它会退回 CPU + gloo，同样能跑通。

15.13 V4 并行策略的"代码极简" 哲学

把 V4 的并行实现与 Megatron-LM / DeepSpeed 等成熟训练框架对比，最显著的差异是代码量——V4 的 4 个并行类（ColumnParallelLinear / RowParallelLinear / ParallelEmbedding / ParallelHead）加起来只有几十行，Megatron-LM 同等功能要数千行。这种极简来自几个设计选择：

选择 1：把训练复杂度外置

Megatron-LM 是训练框架，必须处理优化器状态分片、梯度累积、Pipeline Parallel、混合精度训练等复杂场景。V4 的 inference/model.py 是推理代码——所有训练复杂度被外置到不公开的训练栈，只保留推理需要的最小并行抽象。

选择 2：放弃通用 API

Megatron-LM 的 ColumnParallelLinear 支持丰富配置——gather_output / async_tensor_model_parallel_allreduce / sequence_parallel 等十几个开关。V4 的版本只有一种行为——按 out_features 切分、forward 不通信。这种"放弃通用性"让代码极简。

选择 3：全局变量代替依赖注入

world_size / rank 是文件级全局变量——不是每个 module 持有自己的 ProcessGroup 引用。这违反了"显式依赖" 的传统软件工程纪律，但避免了在一长串构造函数签名里都加 group 参数。

选择 4：MQA-style 让 KV 不切

V4 的 num_key_value_heads=1 让 KV 在 TP 多卡间不切——wkv 是普通 Linear，每 rank 持有完整 KV cache。这绕过了"KV 切分时的复杂索引数学"，代码大幅简化。

代价是显存冗余 N 倍。这份冗余有多贵，可以从 config 直接算出来。Attention.__init__ 里：

kv_cache_size = args.window_size + (args.max_seq_len // self.compress_ratio if self.compress_ratio else 0)

compress_ratios 有 62 个元素（61 层 + 1 个 MTPBlock），其中 31 层用 128、30 层用 4、最后一个 MTP 用 0。按 max_seq_len=1048576 的部署设定（ModelArgs 默认只有 4096），1M context 下逐层累加 kv_cache_size × head_dim(512) × 2 bytes：

  • 主模型 61 层：约 7.7 GiB / 序列
  • 那 30 个带 Indexer 的层还各有一份压缩 KV（max_seq_len/4 × index_head_dim(128) × 2）：约 1.9 GiB / 序列
  • 合计约 9.6 GiB / 序列

8 卡各存一份，就是 77 GiB——每卡一份 9.6 GiB。这在 §15.1 那个"每卡只剩二十来 GB"的账里不是小数目。所以"MQA 让 KV 冗余可以忽略"这个说法要打折：它省的是代码复杂度，不是显存。

这些选择让 V4 的并行代码"看一眼就懂"——这种可读性对开源项目极重要。读者不需要花一周读代码才能理解 V4 怎么工作，几小时就够了。

15.14 部署 V4 时的并行策略选择

把 V4 部署到生产时，TP / EP / DP 的配比是关键工程决策。给一个决策树（以下分档是经验性的粗略建议，不来自官方文档）：

决策点 1：序列长度

  • < 32K：可以用更激进的 TP（16）让单序列延迟最低（参考实现的 TP 上限就是 16，另有 vocab_size 与 n_routed_experts 的整除约束，见 §15.16）
  • 32K - 256K：8 卡 TP 是甜区——平衡延迟与吞吐
  • 256K - 1M：必须 8 卡 TP + 充足 KV 显存——可能需要把 KV cache 切到主机内存

决策点 2：并发量（batch size）

  • 低并发（batch < 4）：TP 比 EP 重要——让单个 sequence 跑得最快
  • 中并发（batch 4-16）：TP=8 + DP（多副本）——每个副本服务一部分用户
  • 高并发（batch > 16）：TP=8 + EP=16/32（跨节点）—— 单一大模型实例服务大量并发

决策点 3：硬件拓扑

先过一遍 §15.1 那道显存关：V4 Pro 的权重实测 864.7 GB，8 卡 H100（640 GB）装不下，8 卡 H200（1128 GB）装下了也不剩多少显存给 KV。所以：

  • 单节点 8 卡：只有单卡显存足够大的卡（B200 这一档）才谈得上；H100 这一档请改跑 V4 Flash
  • 双节点 16 卡：TP=8 + 跨节点 EP=2（每节点 192 expert）——V4 Pro 的现实起点
  • 多节点 32+ 卡：TP=8 + EP 跨节点 + DP 多副本

决策点 4：延迟 vs 吞吐 trade-off

  • 延迟优先：增加 TP（更多并行算 single token）
  • 吞吐优先：增加 batch + 保持 TP 适中

一句话总结：V4 Flash 才是"单节点 8 卡"这个甜区的主角，V4 Pro 的起步就是多节点（除非你的卡显存足够大）。跨节点带来的 all-to-all 成本，正是第 16 章 DeepEP（以及技术报告第 3.1 节 V4 自己的 MegaMoE 方案）要解决的问题。

15.15 V4 的并行抽象与 vLLM PagedAttention 的协同

V4 的并行类（ColumnParallel / RowParallel / ParallelEmbedding / ParallelHead）与 vLLM 的 PagedAttention 在工程层有微妙的协同关系。

协同点 1：KV cache 不切

V4 的 num_key_value_heads=1 让 KV 在所有 TP rank 上是完整副本——每 rank 都能直接访问完整 KV，attention 内部不需要跨 rank 通信。（注意这只是"不用切"，不是"不占显存"：每 rank 都得存一份，1M context 下是 9.6 GiB/序列/卡。）

协同点 2：Q 切分到 head

V4 的 wq_b 是 ColumnParallel——Q head 被切分到不同 rank。vLLM 的 PagedAttention 接收"本 rank 的 Q heads"，与"全副本的 KV" 做 attention 计算，输出 ColumnParallel 形式的 attention 输出。

协同点 3：O 投影做 reduce

V4 的 wo_b 是 RowParallel——attention 输出经过它后做 all_reduce 得到完整 hidden。这与 vLLM 侧的 row-parallel 输出投影是同一个模式。注意 V4 在 all_reduce 之前先 y.float()，这个细节移植时别漏。

协同点 4：MoE 的 expert parallel

V4 的 384 expert 切到不同 rank。vLLM v0.8.5 里"每 rank 持有哪些 expert"由 FusedMoE 层的 determine_expert_map（vllm/model_executor/layers/fused_moe/layer.py:327）算出，对应 V4 的 experts_start_idx / experts_end_idx；而且 v0.8.5 的 EP 同样不做 all_to_all，是"各 rank 算自己的专家、最后一次 all_reduce"，与 V4 参考实现的通信形态一致。

这些协同点意味着：四个并行类本身的移植量不大——它们的语义与 vLLM 已有的 tensor-parallel 抽象基本一一对应。但"移植量不大"不等于"整体适配很薄"：真正的工作量在 KV cache 的双段几何、Indexer 那一步额外的 top-k 打分、以及调度器怎么配合 Compressor 的状态机（见第 19 章）。本节只说并行这一层，不替整个适配下结论。

vLLM 一侧的 TP / EP 实现见《vLLM 推理内核深度解析》专栏第 14 章"分布式并行"。

15.16 V4 并行实现的"边界条件"清单

V4 并行实现里有几个重要的边界条件——必须满足才能正确工作。把它们整理成一份"启动前检查清单"：

边界 1：out_features % world_size == 0

ColumnParallelLinear 要求输出维度能被 world_size 整除。V4 的 out_features 都是 64 的倍数（如 128 head × 512 head_dim = 65536、Indexer 的 weights_proj 是 64），可以整除 8 / 16 / 32 等常见 world_size。但真正先卡住的是 Attention 里没有断言保护的 n_local_groups = o_groups // world_size：o_groups=16，world_size=32 时它变成 0，forward 里 o.view(bsz, seqlen, 0, -1) 直接报错（在 meta 设备上实例化 pin 版 Attention 验证过）——所以参考实现的 TP 上限是 16。

边界 2：in_features % world_size == 0

RowParallelLinear 同理。V4 的 in_features 也都是 128 倍数，对常见 world_size 都满足。

边界 3：vocab_size % world_size == 0

ParallelEmbedding / ParallelHead 要求 vocab 能被 world_size 整除。129280 = 2^8 × 5 × 101——它能被 8（16160/rank）、16（8080）、32（4040）、64（2020）、128（1010）、256（505）整除，常见的 TP 规模全都满足。会踩雷的是 world_size 里带了 5 和 101 之外的奇因子——最常见的就是因子 3，比如 TP=6、TP=12，那时 assert vocab_size % world_size == 0 会直接失败。

边界 4：n_routed_experts % world_size == 0

MoE 要求 expert 数能被 world_size 整除。V4 的 384 expert 整除 8（48 per rank）、16（24 per rank）、32（12 per rank）——常见配置都满足。

边界 5：world_size == 1 时不通信

V4 的并行类有 if world_size > 1: 守卫——单卡模式下完全跳过通信。这让你可以单卡跑 ModelArgs 默认的小配置做调试（inference/model.py 末尾的 __main__ 就是这么做的）+ 多卡跑 V4 Pro 做生产，同一份代码兼容两种模式。

边界 6：dist 必须先 init

world_size = dist.get_world_size() if dist.is_initialized() else 1——V4 的 Transformer.init 检查 dist 是否初始化。如果你忘了 dist.init_process_group(...)，V4 会默默退化到 single-rank 模式，可能导致部署不正确。

把这 6 个边界做成 deployment 的 pre-check，可以在 V4 部署的第一天避免大部分配置错误。

15.17 V4 与 ZeRO 优化器分片的协同

V4 的并行类是"Tensor Parallel + Expert Parallel" 的组合。生产训练通常还会叠加 ZeRO 优化器分片——把优化器状态切到不同 rank。把这套叠加的工程细节梳理一下。

ZeRO-1：仅切分优化器状态（如 Muon 的 momentum buffer）。每 rank 持有所有 weight 但只持有部分优化器状态。

ZeRO-2：切分优化器状态 + 梯度。反向传播时梯度先在本 rank 计算，再 reduce_scatter 到对应 rank。

ZeRO-3：切分优化器状态 + 梯度 + weight。每 rank 只持有部分 weight—— forward 时需要 all_gather 把 weight 拼起来（PyTorch FSDP 的 FULL_SHARD 就是这一档）。

V4 与 ZeRO 的叠加规则：

  • TP 切分维度（如 head 维度）与 ZeRO 切分维度（如 optimizer state 的 layer 维度）必须正交——避免冲突
  • TP rank 之间共享 ZeRO rank ID—— 比如 8 卡 TP + 4 节点，每节点 8 卡是同一个 ZeRO rank
  • Expert 参数同样可以做 ZeRO 切分，前提是不把一个逻辑上独立的矩阵拆开

V4 的实际做法见技术报告第 3.5.1 节：Muon 要拿完整的梯度矩阵才能算更新，与按元素切分的传统 ZeRO 冲突，所以 V4 设计了"混合 ZeRO 分桶"——稠密参数限制 ZeRO 并行规模，用背包算法把整个参数矩阵分配到各 rank；MoE 参数按专家独立优化，把所有专家的矩阵展平后均分到各 rank（不拆开任何一个矩阵），且不限制 ZeRO 规模。MoE 梯度在数据并行 rank 间同步时还被随机舍入到 BF16，用 all-to-all 加本地 FP32 求和代替常规的 reduce-scatter。

这部分配置大多数公司不会自己实现——直接用 DeepSpeed 或 FSDP 的现成支持。但理解原理让你能 debug 配置错误。

15.18 V4 与 Sequence Parallel / Pipeline Parallel 的关系

V4 的并行类只覆盖 Tensor Parallel + Expert Parallel。但生产训练通常还有 Sequence Parallel 和 Pipeline Parallel——把它们与 V4 的关系说清楚。

Sequence Parallel (SP)：

把 sequence 维度切到不同 rank。每 rank 处理 sequence 的一段——主要在 activation memory 上节省（不需要每 rank 存完整 sequence）。

V4 的并行类没有原生 SP 支持——inference/model.py 不切 sequence。训练侧，技术报告第 3.5.3 节写明 V4 用了切 sequence 维度的上下文并行（Context Parallelism），并为压缩注意力设计了两阶段通信——只是 inference 用不到。

为什么 inference 不用 SP：

inference 的显存大头是权重与 KV cache。decode 阶段 activation 只有几个 token，SP 省不下什么；长 prefill 的 activation 才可观，而那要靠分块 prefill 解决，参考实现没有处理。所以 V4 inference 路径不带 SP。

Pipeline Parallel (PP)：

把不同 layer 切到不同 rank——前半模型在 rank 0、后半在 rank 1。每 token 顺序经过 rank。

V4 的 inference/model.py 没有 PP——所有 layer 都在同一 rank（每 rank 持有全部 layer 的 TP 部分）。这是 inference 路径的设计选择；训练侧则用了 DualPipe 的 1F1B 流水线（技术报告第 3.5.2 节）。

为什么 inference 不用 PP：

PP 把每个 token 的前向拆成跨 rank 的串行段，单请求延迟不降反增；要靠多个 micro-batch 交错才能填满流水线，否则就是"流水线泡沫"（pipeline bubble）——某些 rank 在等其他 rank。参考实现面向低并发的最简推理，自然不带 PP（vLLM 等生产推理框架仍提供 PP，常用于跨节点部署）。

训练 vs 推理的并行差异：

维度 训练 推理（V4 inference）
TP ✅ ✅
EP ✅ ✅
SP ✅ ❌（不需要）
PP ✅（大模型必须） ❌（延迟敏感）
ZeRO ✅ ❌（推理无优化器状态）
FSDP ✅（部分） ❌

理解这种差异让你正确选并行配置——训练时该用 PP 就用，推理时不要照搬训练配置。

15.19 延伸阅读

  • Megatron-LM 论文（arXiv:1909.08053）：TP 的工业化先驱
  • DeepSpeed Ulysses（arXiv:2309.14509）：sequence parallelism
  • Tensor Parallelism in Distributed Training（NVIDIA 文档）：TP 的硬件视角
  • 本专栏第 16 章：DeepEP——MoE 的专用 all-to-all 通信库
  • 《vLLM 推理内核深度解析》专栏第 14 章：vLLM 中的 TP 实现

15.20 本章小结

  • V4 用 4 个并行类：ColumnParallel / RowParallel / ParallelEmbedding / ParallelHead
  • ColumnParallel 切 out 维度，无 reduce；RowParallel 切 in 维度，需 all_reduce
  • ParallelEmbedding / ParallelHead 切 vocab 维度，分别用 mask+all_reduce / all_gather
  • Expert Parallel 用稀疏 ModuleList + None 占位 + per-rank 循环的"代码简洁"实现
  • Attention 的 ColumnParallel + RowParallel 经典组合让通信量降到 O(B × S × D)
  • 与 Megatron-LM 同源，但 V4 推理代码极简（4 个类合计几十行）——训练复杂度解耦到不公开训练栈

第 16 章：DeepEP——DeepSeek 为 MoE all-to-all 打造的通信库。