[{"content":"Motivation 如果你问一个 LLM Infra 工程师\u0026quot;大模型系统优化的本质是什么\u0026quot;，大概率会得到两个词：显存和通信。\n一个 70B 参数的模型，FP16 下仅参数就占约 140 GB——远超任何单张 GPU 的显存容量。即便模型能塞进一张卡，训练时的梯度、优化器状态、激活值会让显存需求再膨胀 4-8 倍。于是我们不得不把模型\u0026quot;拆\u0026quot;到多张 GPU 上，而\u0026quot;拆\u0026quot;就意味着通信——卡与卡之间需要频繁交换梯度、参数、激活值。\n显存决定了能不能跑，通信决定了跑得快不快。 这两个问题贯穿了本系列后续所有文章：数据并行要同步梯度（all_reduce），模型并行要交换激活值（all_gather / reduce_scatter），流水线并行要点对点传递中间结果（send / recv），专家并行要全对全分发 token（all_to_all）。\n本文是整个系列的第一篇，我们将：\nPart A：深入 GPU 硬件架构，理解从寄存器到 HBM 的内存层级，以及为什么 HBM 带宽是 LLM 的核心瓶颈 Part B：系统梳理 NCCL 的 8 种通信原语、3 类通信算法、4 种硬件拓扑 掌握这两个基石之后，后续文章中的各种并行策略就不再是\u0026quot;背公式\u0026quot;，而是自然推导的结论。\n前置知识 基础 PyTorch 使用经验（能写训练循环） 了解 Transformer 架构基本原理（Self-Attention、FFN） 知道 GPU 编程的基本概念（kernel、thread、block），但不要求写过 CUDA Part A：GPU 显存模型 GPU vs CPU：不同的设计哲学 在深入 GPU 内存层级之前，先理解 GPU 和 CPU 在设计上的根本区别。\nflowchart LR subgraph CPU[\u0026#34;CPU 设计哲学：延迟优先\u0026#34;] direction TB C0[\u0026#34;强核心 0\\n大 Cache | 乱序执行\u0026#34;] ~~~ C1[\u0026#34;强核心 1\\n大 Cache | 分支预测\u0026#34;] C2[\u0026#34;强核心 2\\n大 Cache | 乱序执行\u0026#34;] ~~~ C3[\u0026#34;强核心 3\\n大 Cache | 分支预测\u0026#34;] CINFO[\u0026#34;核心数 8~128\\n重点：单线程性能\u0026#34;] end subgraph GPU[\u0026#34;GPU 设计哲学：吞吐优先\u0026#34;] direction TB S0[\u0026#34;SM\u0026#34;] ~~~ S1[\u0026#34;SM\u0026#34;] ~~~ S2[\u0026#34;SM\u0026#34;] ~~~ S3[\u0026#34;SM\u0026#34;] S4[\u0026#34;SM\u0026#34;] ~~~ S5[\u0026#34;SM\u0026#34;] ~~~ S6[\u0026#34;SM\u0026#34;] ~~~ S7[\u0026#34;SM ...\u0026#34;] GINFO[\u0026#34;108 SM × 64 CUDA Core\\n重点：大规模并行\u0026#34;] end CPU 追求延迟最优（Latency-oriented）：少量强大的核心，配备大容量缓存、乱序执行引擎和复杂的分支预测器，目标是让单条指令流尽快执行完。\nGPU 追求吞吐最优（Throughput-oriented）：大量简单的核心，放弃复杂的控制逻辑，把晶体管预算全部花在算术逻辑单元（ALU）上，目标是在单位时间内完成尽可能多的浮点运算。\n这就是为什么深度学习天然适合 GPU——矩阵乘法本质上是大规模的、相互独立的乘加运算，完美契合 GPU 的设计哲学。\nSM、Warp 与执行模型 GPU 的基本计算单元是 SM（Streaming Multiprocessor）。以 NVIDIA A100 为例，一块 GPU 包含 108 个 SM，每个 SM 内部结构如下：\nflowchart TD subgraph SM[\u0026#34;SM (Streaming Multiprocessor)\u0026#34;] direction TB subgraph PB[\u0026#34;4 个处理块 (Processing Block)\u0026#34;] direction LR subgraph PB0[\u0026#34;处理块 0\u0026#34;] direction TB P0A[\u0026#34;16 FP32 Core\\n8 FP64 Core\\n1 Tensor Core\u0026#34;] P0B[\u0026#34;Warp Scheduler\\nRegister File\u0026#34;] end subgraph PB1[\u0026#34;处理块 1\u0026#34;] direction TB P1A[\u0026#34;16 FP32 Core\\n8 FP64 Core\\n1 Tensor Core\u0026#34;] P1B[\u0026#34;Warp Scheduler\\nRegister File\u0026#34;] end subgraph PB2[\u0026#34;处理块 2\u0026#34;] direction TB P2A[\u0026#34;16 FP32 Core\\n8 FP64 Core\\n1 Tensor Core\u0026#34;] P2B[\u0026#34;Warp Scheduler\\nRegister File\u0026#34;] end subgraph PB3[\u0026#34;处理块 3\u0026#34;] direction TB P3A[\u0026#34;16 FP32 Core\\n8 FP64 Core\\n1 Tensor Core\u0026#34;] P3B[\u0026#34;Warp Scheduler\\nRegister File\u0026#34;] end end SMEM[\u0026#34;Shared Memory / L1 Cache (192 KB, 可配置比例)\u0026#34;] end PB --\u0026gt; SMEM Warp 是 GPU 执行的基本调度单位，由 32 个线程组成。同一个 Warp 中的 32 个线程在同一时钟周期执行相同的指令（SIMT，Single Instruction Multiple Threads）。这意味着：\n如果 Warp 内出现分支（if-else），两个分支会被串行执行（warp divergence），性能损失严重 内存访问时，如果 32 个线程访问的地址连续（coalesced access），可以合并成一次内存事务；否则需要多次访问 Occupancy（占用率） 是另一个关键概念。每个 SM 可以同时驻留多个 Warp，当一个 Warp 在等待内存数据返回时（延迟可达数百个时钟周期），SM 的 Warp 调度器会切换到另一个就绪的 Warp 继续执行。这种 延迟隐藏（Latency Hiding） 机制是 GPU 高吞吐的关键——但前提是有足够多的活跃 Warp。占用率越高，延迟隐藏越充分，SM 利用率越高。\n内存层级：从寄存器到 HBM GPU 的内存层级从快到慢依次为：\nflowchart TD REG[\u0026#34;\u0026lt;b\u0026gt;Register File\u0026lt;/b\u0026gt;\\n~256 KB | ~20 TB/s | 0 cycles\\n作用域: per thread\u0026#34;] SMEM[\u0026#34;\u0026lt;b\u0026gt;Shared Memory / L1 Cache\u0026lt;/b\u0026gt;\\n192 KB per SM | ~19 TB/s | 1-2 cycles\\n作用域: per SM\u0026#34;] L2[\u0026#34;\u0026lt;b\u0026gt;L2 Cache\u0026lt;/b\u0026gt;\\n40 MB (A100) | ~5 TB/s | 20-30 cycles\\n作用域: 全局\u0026#34;] HBM[\u0026#34;\u0026lt;b\u0026gt;HBM (Global Memory)\u0026lt;/b\u0026gt;\\n80 GB (A100) | 2.0 TB/s | 200-400 cycles\\n作用域: 全局\u0026#34;] REG --\u0026gt;|\u0026#34;~10x 带宽下降\u0026#34;| SMEM --\u0026gt;|\u0026#34;~4x\u0026#34;| L2 --\u0026gt;|\u0026#34;~2.5x\u0026#34;| HBM style REG fill:#2d6a4f,color:#fff style SMEM fill:#40916c,color:#fff style L2 fill:#74c69d,color:#000 style HBM fill:#b7e4c7,color:#000 逐层说明：\n寄存器（Register File）：每个线程私有，访问零延迟。A100 每个 SM 有 256 KB 寄存器文件，分配给 SM 上所有活跃线程。寄存器是最宝贵的资源——每个线程使用的寄存器越多，SM 上能同时驻留的 Warp 就越少，占用率就越低。\n共享内存（Shared Memory）/ L1 缓存：SM 内所有线程共享，A100 上每个 SM 有 192 KB，可以在 Shared Memory 和 L1 Cache 之间灵活配置比例。共享内存是程序员显式管理的\u0026quot;软件缓存\u0026quot;，在矩阵乘法的 Tiling 优化中是核心工具——先把数据块从 HBM 搬到 Shared Memory，然后在 Shared Memory 上做多次计算，摊薄 HBM 访问开销。\nL2 缓存：全局共享，A100 上有 40 MB。对程序员不可直接控制，由硬件自动管理。\nHBM（High Bandwidth Memory）：这就是我们通常说的\u0026quot;显存\u0026quot;。A100 SXM 版本有 80 GB、带宽 2.0 TB/s。虽然 2 TB/s 的带宽看起来已经很高，但对比寄存器的 ~20 TB/s，存在一个量级的差距。任何频繁访问 HBM 的操作都可能成为瓶颈。\n一个直观的类比：寄存器像你桌面上随手可拿的便签，Shared Memory 像抽屉里的文件夹，L2 像同一层楼的文件柜，HBM 像隔壁楼的仓库。你肯定希望把最常用的数据放在桌面上，而不是每次都跑去仓库取。\n为什么 HBM 带宽是 LLM 的核心瓶颈 要判断一个操作是\u0026quot;算力瓶颈\u0026quot;还是\u0026quot;带宽瓶颈\u0026quot;，核心工具是 Roofline 模型和算术强度（Arithmetic Intensity）。\n算术强度 = 计算量（FLOPs） / 数据访问量（Bytes）\n以 A100 SXM 为例：\n峰值算力：312 TFLOPS（FP16 Tensor Core） HBM 带宽：2.0 TB/s 平衡点：312 / 2.0 = 156 FLOPs/Byte 也就是说，每从 HBM 读取 1 字节数据，至少需要做 156 次浮点运算，才能让计算单元不闲着。低于这个比值的操作就是 Memory-bound（带宽瓶颈），高于则是 Compute-bound（算力瓶颈）。\nflowchart LR subgraph Roofline[\u0026#34;Roofline Model — A100 SXM\u0026#34;] direction LR MB[\u0026#34;🔵 **Memory Bound**\\nAI \u0026lt; 156 FLOPs/Byte\\n性能 = AI × 2.0 TB/s\u0026#34;] BP[\u0026#34;⚖️ **平衡点**\\n156 FLOPs/Byte\u0026#34;] CB[\u0026#34;🔴 **Compute Bound**\\nAI ≥ 156 FLOPs/Byte\\n性能 → 312 TFLOPS\u0026#34;] MB --- BP --- CB end style MB fill:#fff3cd,stroke:#856404 style BP fill:#cce5ff,stroke:#004085 style CB fill:#d4edda,stroke:#155724 现在来看 LLM 中的关键操作：\n矩阵乘法（GEMM）：对于大尺寸的矩阵乘法 $C = A \\times B$，其中 $A \\in \\mathbb{R}^{M \\times K}$，$B \\in \\mathbb{R}^{K \\times N}$，计算量为 $2MKN$ FLOPs，数据量约为 $2(MK + KN + MN)$ Bytes（FP16）。当 $M$、$K$、$N$ 足够大时，算术强度可以很高，是 Compute-bound 的。这就是为什么训练时的 forward pass（大 batch、大矩阵乘法）通常可以充分利用 GPU 算力。\n逐元素操作（LayerNorm、GELU、Softmax 等）：每个元素只做几次运算，但都要从 HBM 读一次、写一次。算术强度极低（通常 \u0026lt; 10 FLOPs/Byte），严重 Memory-bound。这就是 FlashAttention 和 kernel fusion 优化的动机——减少对 HBM 的访问次数。\n自回归推理（Autoregressive Decoding）：这是 LLM 推理的核心瓶颈。每个 decode step 只生成一个 token，对应的矩阵乘法退化为矩阵-向量乘法（GEMV）：$y = W \\cdot x$，其中 $W \\in \\mathbb{R}^{d \\times d}$，$x \\in \\mathbb{R}^{d \\times 1}$。计算量为 $2d^2$ FLOPs，但需要读取整个权重矩阵 $2d^2$ Bytes（FP16），算术强度 = $2d^2 / 2d^2$ = 1 FLOPs/Byte。这比平衡点 156 低了两个数量级，意味着 GPU 算力利用率不到 1%，几乎全部时间在等 HBM 传数据。\n这就是为什么 LLM 推理优化中如此强调 KV Cache（减少重复计算的内存访问）、量化（减少权重的字节数）、Continuous Batching（增大 batch size 提高算术强度）。\n实战：GPU 显存 Profiling 理论之后我们来看实际数据。以下代码分析了一个 Transformer 模型在单卡上的显存占用分布：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 import torch import torch.nn as nn def profile_memory(model, batch_size, seq_len, vocab_size): \u0026#34;\u0026#34;\u0026#34;分析模型训练时的显存占用\u0026#34;\u0026#34;\u0026#34; device = torch.device(\u0026#34;cuda\u0026#34;) model = model.to(device) # 1. 模型参数 param_mem = sum(p.numel() * p.element_size() for p in model.parameters()) print(f\u0026#34;模型参数显存: {param_mem / 1024**3:.2f} GB\u0026#34;) # 2. 优化器状态 (Adam: 每个参数额外 2 份 FP32 拷贝) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) # 3. 前向传播，观察激活值 torch.cuda.reset_peak_memory_stats() input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=device) with torch.cuda.amp.autocast(dtype=torch.float16): output = model(input_ids) loss = output.sum() peak_after_forward = torch.cuda.max_memory_allocated() print(f\u0026#34;前向后峰值显存: {peak_after_forward / 1024**3:.2f} GB\u0026#34;) activation_mem = peak_after_forward - param_mem print(f\u0026#34; 其中激活值约: {activation_mem / 1024**3:.2f} GB\u0026#34;) # 4. 反向传播 loss.backward() peak_after_backward = torch.cuda.max_memory_allocated() grad_mem = peak_after_backward - peak_after_forward print(f\u0026#34;梯度显存增量: {grad_mem / 1024**3:.2f} GB\u0026#34;) # 5. 优化器 step optimizer.step() peak_after_optim = torch.cuda.max_memory_allocated() optim_mem = peak_after_optim - peak_after_backward print(f\u0026#34;优化器状态增量: {optim_mem / 1024**3:.2f} GB\u0026#34;) print(f\u0026#34;\\n总峰值显存: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB\u0026#34;) 运行结果大致如下（以一个 1.3B 模型为例）：\n模型参数显存: 2.60 GB (FP16) 前向后峰值显存: 8.50 GB 其中激活值约: 5.90 GB ← 激活值是大头！ 梯度显存增量: 2.60 GB (与参数等大) 优化器状态增量: 10.40 GB (Adam: 2 x FP32 = 4x 参数大小) 总峰值显存: 18.50 GB 这里体现了一个重要的经验法则——Adam 优化器的 4 倍法则：\n组件 大小 精度 模型参数 $\\Phi$ FP16 = $2\\Phi$ bytes 梯度 $\\Phi$ FP16 = $2\\Phi$ bytes Adam m（一阶矩） $\\Phi$ FP32 = $4\\Phi$ bytes Adam v（二阶矩） $\\Phi$ FP32 = $4\\Phi$ bytes 主权重（master weights） $\\Phi$ FP32 = $4\\Phi$ bytes 总计 $16\\Phi$ bytes 对于一个 $\\Phi$ 参数的模型，仅参数相关的显存就需要 $16\\Phi$ bytes，这还不包括激活值。以 70B 模型为例：$16 \\times 70 \\times 10^9 = 1120$ GB——至少需要 14 张 A100 80GB 才能放下，而且还没有为激活值留余量。\n完整代码见 code/01-gpu-memory-distributed/memory_profiling.py\nPart B：分布式通信 为什么需要多卡 上一节的计算已经给出了答案：单卡装不下大模型。更具体地说：\n模型规模 参数显存 (FP16) 训练总显存 (Adam) 最少需要 GPU 数 (A100 80GB) 7B 14 GB ~112 GB + 激活 2 13B 26 GB ~208 GB + 激活 4 70B 140 GB ~1120 GB + 激活 16+ 405B 810 GB ~6480 GB + 激活 100+ 当模型分布到多张 GPU 上后，各种并行策略（DDP、FSDP、TP、PP、EP）在不同阶段需要不同类型的通信。NCCL（NVIDIA Collective Communications Library）提供了这些通信操作的高效实现。\nNCCL 通信原语：从训练场景推导 理解通信原语最好的方式，不是逐个背定义，而是从实际的训练/推理场景出发，看每个场景自然需要什么样的数据搬运。NCCL（NVIDIA Collective Communications Library）提供了 8 种通信原语，每一种都对应着真实的工程需求。\n下面我们假设有 4 张 GPU（Rank 0-3, $P = 4$），数据大小为 $N$，从 5 个场景出发逐一推导。\n场景一：DDP — 每张卡算不同数据，梯度怎么同步？ DDP（Distributed Data Parallel） 是最基本的数据并行策略：每张卡持有完整的模型副本，各自处理不同的 mini-batch，然后同步梯度使参数更新一致。\n问题很清楚：每张卡算出了自己的局部梯度 $g_i$，我们需要让每张卡都拿到 $\\bar{g} = \\frac{1}{P}\\sum_i g_i$。这正是 All-Reduce 的定义。\nAll-Reduce 功能：所有卡上的数据做归约（通常是求和），结果存到每张卡。\nflowchart LR subgraph Before[\u0026#34; \u0026#34;] direction TB R0B[\u0026#34;Rank 0: a0 a1 a2 a3\u0026#34;] R1B[\u0026#34;Rank 1: b0 b1 b2 b3\u0026#34;] R2B[\u0026#34;Rank 2: c0 c1 c2 c3\u0026#34;] R3B[\u0026#34;Rank 3: d0 d1 d2 d3\u0026#34;] end Before --\u0026gt;|\u0026#34;all_reduce(sum)\u0026#34;| After subgraph After[\u0026#34;Σi = ai + bi + ci + di\u0026#34;] direction TB R0A[\u0026#34;Rank 0: Σ0 Σ1 Σ2 Σ3\u0026#34;] R1A[\u0026#34;Rank 1: Σ0 Σ1 Σ2 Σ3\u0026#34;] R2A[\u0026#34;Rank 2: Σ0 Σ1 Σ2 Σ3\u0026#34;] R3A[\u0026#34;Rank 3: Σ0 Σ1 Σ2 Σ3\u0026#34;] end 通信量：每张卡发送和接收约 $2N \\cdot \\frac{P-1}{P}$ 数据（Ring 算法）。当 $P$ 很大时趋近 $2N$——与 GPU 数量几乎无关，这是 Ring 算法的精妙之处。\n场景二：FSDP — 参数都切碎了，怎么算前向/反向？ DDP 的问题是每张卡都存完整参数——70B 模型根本放不下。FSDP（Fully Sharded Data Parallel） 的思路是：参数、梯度、优化器状态全部按卡切分，每张卡只存 $1/P$。\n但切碎之后要计算怎么办？\n前向传播：计算某一层时，需要该层的完整参数。每张卡只有一个分片，必须临时把所有分片拼起来 → All-Gather 反向传播：每张卡算出完整梯度后，需要把梯度归约并重新切分，每张卡只保留自己负责的那份 → Reduce-Scatter All-Gather 功能：每张卡贡献自己的一个分片，拼出完整数据到所有卡。\nflowchart LR subgraph Before[\u0026#34; \u0026#34;] direction TB R0B[\u0026#34;Rank 0: a0 _ _ _\u0026#34;] R1B[\u0026#34;Rank 1: _ b1 _ _\u0026#34;] R2B[\u0026#34;Rank 2: _ _ c2 _\u0026#34;] R3B[\u0026#34;Rank 3: _ _ _ d3\u0026#34;] end Before --\u0026gt;|\u0026#34;all_gather\u0026#34;| After subgraph After[\u0026#34; \u0026#34;] direction TB R0A[\u0026#34;Rank 0: a0 b1 c2 d3\u0026#34;] R1A[\u0026#34;Rank 1: a0 b1 c2 d3\u0026#34;] R2A[\u0026#34;Rank 2: a0 b1 c2 d3\u0026#34;] R3A[\u0026#34;Rank 3: a0 b1 c2 d3\u0026#34;] end 通信量：每张卡接收 $N \\cdot \\frac{P-1}{P}$ 数据。\nReduce-Scatter 功能：先归约（求和），再将结果的不同部分分散到不同卡。\nflowchart LR subgraph Before[\u0026#34; \u0026#34;] direction TB R0B[\u0026#34;Rank 0: a0 a1 a2 a3\u0026#34;] R1B[\u0026#34;Rank 1: b0 b1 b2 b3\u0026#34;] R2B[\u0026#34;Rank 2: c0 c1 c2 c3\u0026#34;] R3B[\u0026#34;Rank 3: d0 d1 d2 d3\u0026#34;] end Before --\u0026gt;|\u0026#34;reduce_scatter(sum)\u0026#34;| After subgraph After[\u0026#34;Σi = ai + bi + ci + di\u0026#34;] direction TB R0A[\u0026#34;Rank 0: Σ0 _ _ _\u0026#34;] R1A[\u0026#34;Rank 1: _ Σ1 _ _\u0026#34;] R2A[\u0026#34;Rank 2: _ _ Σ2 _\u0026#34;] R3A[\u0026#34;Rank 3: _ _ _ Σ3\u0026#34;] end 通信量：每张卡发送 $N \\cdot \\frac{P-1}{P}$ 数据。\nFSDP 的通信闭环 flowchart LR FWD[\u0026#34;**前向**\u0026#34;] --\u0026gt; AG1[\u0026#34;All-Gather\\n拼出完整参数\u0026#34;] --\u0026gt; COMP1[\u0026#34;计算\u0026#34;] --\u0026gt; FREE[\u0026#34;释放完整参数\u0026#34;] BWD[\u0026#34;**反向**\u0026#34;] --\u0026gt; AG2[\u0026#34;All-Gather\\n拼出完整参数\u0026#34;] --\u0026gt; COMP2[\u0026#34;计算梯度\u0026#34;] --\u0026gt; RS[\u0026#34;Reduce-Scatter\\n只保留梯度分片\u0026#34;] UPD[\u0026#34;**更新**\u0026#34;] --\u0026gt; LOCAL[\u0026#34;每卡用自己的\\n梯度分片更新\\n自己的参数分片\u0026#34;] 关键洞察：All-Reduce 在概念上等价于 Reduce-Scatter + All-Gather。DDP 用 All-Reduce 是因为每张卡需要完整梯度；FSDP 用 Reduce-Scatter 是因为每张卡只需要自己那份。NCCL 在实现 All-Reduce 时，内部也经常将其分解为这两步。\n场景三：Pipeline Parallel — 模型按层切开，激活值怎么传？ Pipeline Parallel（PP） 把模型按层划分为多个 stage，每个 stage 放在不同的 GPU 上。前向传播时，stage 0 算完要把激活值传给 stage 1；反向传播时，stage 1 要把梯度传回 stage 0。\n这不需要集合通信——就是两张卡之间直接传数据 → Send / Recv。\nSend / Recv（Point-to-Point） 功能：两张卡之间的点对点通信。\nflowchart LR R0[\u0026#34;Rank 0\u0026#34;] --\u0026gt;|\u0026#34;send(activations) — 前向\u0026#34;| R1[\u0026#34;Rank 1\u0026#34;] R1 --\u0026gt;|\u0026#34;send(gradients) — 反向\u0026#34;| R0 通信量：$O(N)$，仅涉及两张卡。\nSend/Recv 是唯一的非集合通信原语。Pipeline 调度算法（1F1B、Zero Bubble 等）的本质就是精心编排这些 Send/Recv 的时序，让不同 stage 尽量同时忙碌，减少流水线气泡。\n场景四：MoE Expert Parallel — token 路由到不同专家，怎么搬？ 在 MoE（Mixture-of-Experts） 模型中，每个 token 经过 gating network 被路由到一个或几个专家。当使用 Expert Parallel（EP） 时，不同的专家分布在不同的 GPU 上。\n问题是：每张卡上的 token 可能需要去任意一张卡上的专家。这不是\u0026quot;一对多\u0026quot;或\u0026quot;多对一\u0026quot;，而是每张卡都要向每张卡发送不同的数据 → All-to-All。\nAll-to-All 功能：每张卡向其他所有卡发送不同的数据块，同时接收来自所有卡的数据块。\nflowchart LR subgraph Before[\u0026#34;Before — 按行: 每卡的数据\u0026#34;] direction TB R0B[\u0026#34;Rank 0: a→0 a→1 a→2 a→3\u0026#34;] R1B[\u0026#34;Rank 1: b→0 b→1 b→2 b→3\u0026#34;] R2B[\u0026#34;Rank 2: c→0 c→1 c→2 c→3\u0026#34;] R3B[\u0026#34;Rank 3: d→0 d→1 d→2 d→3\u0026#34;] end Before --\u0026gt;|\u0026#34;all_to_all\u0026#34;| After subgraph After[\u0026#34;After — 按列: 每卡收集来自所有卡的数据\u0026#34;] direction TB R0A[\u0026#34;Rank 0: a→0 b→0 c→0 d→0\u0026#34;] R1A[\u0026#34;Rank 1: a→1 b→1 c→1 d→1\u0026#34;] R2A[\u0026#34;Rank 2: a→2 b→2 c→2 d→2\u0026#34;] R3A[\u0026#34;Rank 3: a→3 b→3 c→3 d→3\u0026#34;] end MoE 层的通信模式是：All-to-All（dispatch: token → expert）→ 专家计算 → All-to-All（combine: expert output → 原始卡），一前一后两次 All-to-All。\n通信量：每张卡发送和接收 $N \\cdot \\frac{P-1}{P}$ 数据。All-to-All 是最\u0026quot;重\u0026quot;的集合通信，因为它涉及全网状（full-mesh）的数据交换，对网络拓扑和带宽极为敏感。\n场景五：初始化与数据流 — 一些\u0026quot;胶水\u0026quot;原语 上面四个场景覆盖了训练中的核心通信需求。还有几个原语用于初始化和数据管理：\nBroadcast 功能：把一张卡上的数据广播到所有卡。\nflowchart LR subgraph Before[\u0026#34; \u0026#34;] direction TB R0B[\u0026#34;Rank 0: A A A A\u0026#34;] R1B[\u0026#34;Rank 1: . . . .\u0026#34;] R2B[\u0026#34;Rank 2: . . . .\u0026#34;] R3B[\u0026#34;Rank 3: . . . .\u0026#34;] end Before --\u0026gt;|\u0026#34;broadcast\\n(from Rank 0)\u0026#34;| After subgraph After[\u0026#34; \u0026#34;] direction TB R0A[\u0026#34;Rank 0: A A A A\u0026#34;] R1A[\u0026#34;Rank 1: A A A A\u0026#34;] R2A[\u0026#34;Rank 2: A A A A\u0026#34;] R3A[\u0026#34;Rank 3: A A A A\u0026#34;] end 典型场景：训练开始前，Rank 0 上初始化模型参数，通过 Broadcast 确保所有卡参数一致。DDP 启动时内部就会调用 Broadcast。\n通信量：每张卡发送或接收 $O(N)$ 数据。但全局总通信量取决于实现：朴素实现（Root 逐一发送）为 $O(N \\cdot P)$；实际的树形广播中，多张卡并行转发，总步数为 $\\log P$，全局通信量为 $O(N \\log P)$。\n视角 通信量 单卡（发送或接收） $O(N)$ 全局（朴素实现） $O(N \\cdot P)$ 全局（树形广播） $O(N \\log P)$ Scatter / Gather Scatter：从一张卡分发不同数据块到各卡。Gather：各卡的数据收集到一张卡（Scatter 的逆操作）。\nflowchart LR subgraph SCATTER[\u0026#34;Scatter (Rank 0 分发)\u0026#34;] direction LR S_IN[\u0026#34;Rank 0: d0 d1 d2 d3\u0026#34;] --\u0026gt;|\u0026#34;scatter\u0026#34;| S_OUT0[\u0026#34;Rank 0: d0\u0026#34;] S_IN --\u0026gt; S_OUT1[\u0026#34;Rank 1: d1\u0026#34;] S_IN --\u0026gt; S_OUT2[\u0026#34;Rank 2: d2\u0026#34;] S_IN --\u0026gt; S_OUT3[\u0026#34;Rank 3: d3\u0026#34;] end subgraph GATHER[\u0026#34;Gather (收集到 Rank 0)\u0026#34;] direction LR G_IN0[\u0026#34;Rank 0: d0\u0026#34;] --\u0026gt;|\u0026#34;gather\u0026#34;| G_OUT[\u0026#34;Rank 0: d0 d1 d2 d3\u0026#34;] G_IN1[\u0026#34;Rank 1: d1\u0026#34;] --\u0026gt; G_OUT G_IN2[\u0026#34;Rank 2: d2\u0026#34;] --\u0026gt; G_OUT G_IN3[\u0026#34;Rank 3: d3\u0026#34;] --\u0026gt; G_OUT end 典型场景：数据加载时，Rank 0 读取一个大 batch 然后 Scatter 到各卡；评估阶段 Gather 各卡的预测结果到 Rank 0 做汇总。\n全景回顾：从场景到原语 训练场景 通信需求 对应原语 通信量（每卡） DDP 梯度同步 每卡的梯度求和，结果给所有卡 all_reduce $2N \\cdot \\frac{P-1}{P}$ FSDP 前向（拼参数） 每卡贡献分片，拼出完整数据 all_gather $N \\cdot \\frac{P-1}{P}$ FSDP 反向（切梯度） 梯度归约后各卡只留自己的分片 reduce_scatter $N \\cdot \\frac{P-1}{P}$ Pipeline Parallel 相邻 stage 间传递激活值/梯度 send / recv $N$ Expert Parallel (MoE) token 全排列式路由到各专家 all_to_all $N \\cdot \\frac{P-1}{P}$ 参数初始化 一张卡的参数复制到所有卡 broadcast $N$（全局 $N \\log P$） 数据分发 / 结果收集 一对多分发或多对一收集 scatter / gather $N \\cdot \\frac{P-1}{P}$ 以下代码展示了如何在 PyTorch 中使用 all_reduce：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 import torch import torch.distributed as dist import os def main(): dist.init_process_group(backend=\u0026#34;nccl\u0026#34;) rank = dist.get_rank() world_size = dist.get_world_size() device = torch.device(f\u0026#34;cuda:{rank}\u0026#34;) # 每张卡创建自己的 tensor tensor = torch.ones(1024, 1024, device=device) * (rank + 1) print(f\u0026#34;[Rank {rank}] Before all_reduce: sum = {tensor.sum().item():.0f}\u0026#34;) # All-Reduce: 求和 dist.all_reduce(tensor, op=dist.ReduceOp.SUM) print(f\u0026#34;[Rank {rank}] After all_reduce: sum = {tensor.sum().item():.0f}\u0026#34;) # 期望结果: 每个元素 = 1+2+3+4 = 10 (4 卡) # 总和 = 10 * 1024 * 1024 = 10,485,760 dist.destroy_process_group() if __name__ == \u0026#34;__main__\u0026#34;: main() # 运行: torchrun --nproc_per_node=4 nccl_allreduce.py 完整代码（包含所有原语的示例）见 code/01-gpu-memory-distributed/nccl_allreduce.py\n集合通信算法 通信原语定义了\u0026quot;做什么\u0026quot;，通信算法决定了\u0026quot;怎么做\u0026quot;。同一个 all_reduce 操作，不同的算法在延迟和带宽利用上差异巨大。\nRing All-Reduce Ring All-Reduce 是最经典的带宽最优算法，分为两个阶段：\n阶段一：Reduce-Scatter（$P-1$ 步）\n4 张卡排成一个环，每张卡将数据分为 $P=4$ 个块。每一步，每张卡向下一个邻居发送一个块，同时从上一个邻居接收一个块并累加。经过 $P-1=3$ 步后，每张卡上有一个块包含了所有卡的归约结果。\nflowchart TD subgraph INIT[\u0026#34;初始状态\u0026#34;] direction LR I0[\u0026#34;Rank 0: a0 a1 a2 a3\u0026#34;] ~~~ I1[\u0026#34;Rank 1: b0 b1 b2 b3\u0026#34;] I2[\u0026#34;Rank 2: c0 c1 c2 c3\u0026#34;] ~~~ I3[\u0026#34;Rank 3: d0 d1 d2 d3\u0026#34;] end subgraph STEP1[\u0026#34;Step 1 — 每卡发送一个块给右邻居，接收左邻居的块并累加\u0026#34;] direction LR S1R0[\u0026#34;Rank 0: a0 · a1 · a2 · **a3+d3**\u0026#34;] S1R1[\u0026#34;Rank 1: **b0+a0** · b1 · b2 · b3\u0026#34;] S1R2[\u0026#34;Rank 2: c0 · **c1+b1** · c2 · c3\u0026#34;] S1R3[\u0026#34;Rank 3: d0 · d1 · **d2+c2** · d3\u0026#34;] end subgraph STEP3[\u0026#34;Step 3 (最终) — 每张卡上有一个完整归约的块\u0026#34;] direction LR S3R0[\u0026#34;Rank 0: · · · · · · **Σ3** ← 块 3 完整\u0026#34;] S3R1[\u0026#34;Rank 1: **Σ0** · · · · · · ← 块 0 完整\u0026#34;] S3R2[\u0026#34;Rank 2: · **Σ1** · · · · ← 块 1 完整\u0026#34;] S3R3[\u0026#34;Rank 3: · · **Σ2** · · ← 块 2 完整\u0026#34;] end INIT --\u0026gt; STEP1 --\u0026gt;|\u0026#34;继续传递\\n部分归约的块\u0026#34;| STEP3 阶段二：All-Gather（$P-1$ 步）\n同样在环上传递，但这次不做归约，只做拷贝。$P-1$ 步后每张卡拥有完整的归约结果。\n复杂度分析：\n总步数：$2(P-1)$ 每步每卡传输量：$N/P$ 总通信量（每卡）：$2 \\cdot \\frac{P-1}{P} \\cdot N$ 当 $P$ 很大时趋近于 $2N$，与 GPU 数量无关——这就是带宽最优的含义 Ring 的缺点是延迟为 $O(P)$：数据必须绕环一圈，每一步都有一次网络延迟。当消息较小时，延迟开销会超过传输时间，效率降低。\nTree All-Reduce Tree All-Reduce 用一棵二叉树组织通信：\ngraph TD R0[\u0026#34;Rank 0 (Root)\u0026#34;] R1[\u0026#34;Rank 1\u0026#34;] R2[\u0026#34;Rank 2\u0026#34;] R3[\u0026#34;Rank 3\u0026#34;] R0 --- R1 R0 --- R2 R1 --- R3 style R0 fill:#d4a574,color:#000 阶段一 (Reduce): 叶子向根归约 | 阶段二 (Broadcast): 根向叶子广播\n复杂度：\n延迟：$O(\\log P)$——远优于 Ring 的 $O(P)$ 带宽利用率：较差，非叶子节点成为带宽瓶颈 Tree 适合小消息、多节点的场景。实际中 NCCL 会根据消息大小自动选择算法。\nRecursive Halving-Doubling 这是一种折中方案，结合了 Ring 的带宽效率和 Tree 的低延迟：\nHalving 阶段：每轮将参与者分成两半，两半之间交换并归约各自缺少的部分 Doubling 阶段：反向传播完整结果 延迟为 $O(\\log P)$，带宽利用率接近最优。适合 GPU 数为 2 的幂次的情况。\nNCCL 的实际选择策略：NCCL 并不固定使用某一种算法，而是根据消息大小、GPU 数量、拓扑结构动态选择：\n小消息（\u0026lt; 256 KB）：偏向 Tree 大消息（\u0026gt; 数 MB）：偏向 Ring 特定拓扑下会使用更高效的变种 通信拓扑 通信算法跑在硬件拓扑之上。不同的硬件链路带宽差异巨大，直接影响了分布式训练的瓶颈位置。\nNVLink NVLink 是 NVIDIA GPU 之间的高速直连链路：\n世代 单链路带宽 GPU 间总带宽 典型配置 NVLink 3 (A100) 50 GB/s 600 GB/s (12 links) DGX A100 NVLink 4 (H100) 50 GB/s 900 GB/s (18 links) DGX H100 A100 机内 8 卡通过 NVLink 互联，每对 GPU 间带宽 600 GB/s——是 HBM 带宽的 30%，是 PCIe Gen4 x16 的 20 倍以上。\nNVSwitch NVSwitch 是 NVIDIA 的全交换芯片，实现了节点内 GPU 之间的全双工（Full Bisection Bandwidth） 互联：\nflowchart LR subgraph DGX[\u0026#34;DGX A100 — 任意 GPU 对: 600 GB/s, 聚合: 4.8 TB/s\u0026#34;] direction LR subgraph LEFT[\u0026#34; \u0026#34;] direction TB G0[\u0026#34;GPU 0\u0026#34;] ~~~ G1[\u0026#34;GPU 1\u0026#34;] ~~~ G2[\u0026#34;GPU 2\u0026#34;] ~~~ G3[\u0026#34;GPU 3\u0026#34;] end NVS[\u0026#34;NVSwitch\\n× 6\u0026#34;] subgraph RIGHT[\u0026#34; \u0026#34;] direction TB G4[\u0026#34;GPU 4\u0026#34;] ~~~ G5[\u0026#34;GPU 5\u0026#34;] ~~~ G6[\u0026#34;GPU 6\u0026#34;] ~~~ G7[\u0026#34;GPU 7\u0026#34;] end LEFT \u0026lt;--\u0026gt;|\u0026#34;NVLink\u0026#34;| NVS \u0026lt;--\u0026gt;|\u0026#34;NVLink\u0026#34;| RIGHT end 有了 NVSwitch，节点内 all_reduce 的带宽几乎不受 GPU 对数限制。\nPCIe PCIe 是 CPU 和 GPU 之间、以及没有 NVLink 的 GPU 之间的通信通道：\nPCIe Gen4 x16: 约 32 GB/s（双向） PCIe Gen5 x16: 约 64 GB/s（双向） 相比 NVLink 600 GB/s，PCIe 带宽低了一个数量级。在消费级 GPU 或部分云实例上，GPU 间通信可能退化到走 PCIe，此时通信会成为严重瓶颈。\nRDMA / InfiniBand 跨节点通信依赖网络互联，主流方案是 InfiniBand + GPUDirect RDMA：\nInfiniBand HDR: 200 Gb/s = 25 GB/s InfiniBand NDR: 400 Gb/s = 50 GB/s GPUDirect RDMA: GPU 显存直接通过网卡发送数据，绕过 CPU 和系统内存，减少一次拷贝延迟 flowchart LR subgraph A[\u0026#34;节点 A — NVLink 600 GB/s\u0026#34;] direction TB A1[\u0026#34;GPU\u0026#34;] \u0026lt;--\u0026gt;|NVLink| A2[\u0026#34;GPU\u0026#34;] A3[\u0026#34;GPU\u0026#34;] \u0026lt;--\u0026gt;|NVLink| A4[\u0026#34;GPU\u0026#34;] A1 \u0026lt;--\u0026gt;|NVSwitch| A3 NICA[\u0026#34;NIC\u0026#34;] end subgraph B[\u0026#34;节点 B — NVLink 600 GB/s\u0026#34;] direction TB B1[\u0026#34;GPU\u0026#34;] \u0026lt;--\u0026gt;|NVLink| B2[\u0026#34;GPU\u0026#34;] B3[\u0026#34;GPU\u0026#34;] \u0026lt;--\u0026gt;|NVLink| B4[\u0026#34;GPU\u0026#34;] B1 \u0026lt;--\u0026gt;|NVSwitch| B3 NICB[\u0026#34;NIC\u0026#34;] end NICA \u0026lt;====\u0026gt;|\u0026#34;InfiniBand RDMA\\n25-50 GB/s\u0026#34;| NICB 节点内外的带宽差距（约 10-20 倍）深刻影响了并行策略的设计：\n通信量大的并行策略（如 TP）通常放在节点内，充分利用 NVLink 通信量较小的并行策略（如 PP、DP）可以跨节点部署 这就是大规模训练中\u0026quot;节点内 TP、节点间 DP/PP\u0026quot;成为标准配置的原因 总结与下一步 本文覆盖了 LLM Infra 的两大基石：\n显存方面：\nGPU 采用吞吐优先的设计，通过大量简单核心和 Warp 级别的延迟隐藏实现高吞吐 内存层级从寄存器到 HBM，带宽跨越 4 个数量级 LLM 推理（自回归解码）是典型的 Memory-bound 问题，算术强度仅约 1 FLOPs/Byte 训练时，Adam 优化器让显存需求膨胀到参数量的 16 倍 通信方面：\n8 种 NCCL 通信原语各有其对应的分布式并行场景 Ring All-Reduce 带宽最优但延迟 $O(P)$，Tree 延迟最优但带宽差 节点内 NVLink (600 GB/s) 与节点间 InfiniBand (25-50 GB/s) 的带宽鸿沟，决定了混合并行的拓扑布局 在下一篇文章**《分布式并行策略全景》**中，我们将在这些基础之上，系统介绍 DDP、FSDP、TP、PP、SP、EP 等并行策略——你会看到，每种策略本质上都是在显存和通信之间做不同的 trade-off，而选择哪种 trade-off，取决于模型规模、硬件拓扑和训练/推理场景。\n参考资料 NVIDIA CUDA Programming Guide — docs.nvidia.com/cuda — SM、Warp、内存层级的权威文档 NVIDIA A100 Whitepaper — GPU 架构细节、Tensor Core 规格、NVLink/NVSwitch 拓扑 NCCL Documentation — docs.nvidia.com/deeplearning/nccl — 通信原语 API 和算法说明 Roofline Model — Williams, Waterman, Patterson, \u0026ldquo;Roofline: An Insightful Visual Performance Model for Multicore Architectures\u0026rdquo;, Communications of the ACM, 2009 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al., 2020 — 对优化器状态、梯度、参数的显存分析 Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — Narayanan et al., 2021 — 混合并行策略与通信拓扑的关系 PyTorch Distributed Overview — pytorch.org/tutorials — PyTorch 分布式通信接口 ","permalink":"https://mzf666.github.io/llm-infra/zh/posts/01-gpu-memory-distributed/","summary":"从 GPU 显存层级到 NCCL 通信原语，理解 LLM Infra 优化的两大基石。","title":"GPU 显存模型与分布式通信基础"},{"content":"Motivation 上一篇文章中，我们算过一笔账：一个 70B 参数的模型，仅 Adam 优化器相关的显存就需要 $16 \\times 70 \\times 10^9 \\approx 1120$ GB——至少 14 张 A100 80GB 才放得下，而且还没算激活值。单卡训练大模型从物理上就不可能，我们必须把计算分布到多张 GPU 上。\n但\u0026quot;分布到多卡\u0026quot;不是一句话那么简单。不同的并行策略在切什么（参数、梯度、激活值、序列）、怎么切（按层、按维度、按数据）、通信开销（all_reduce、all_gather、send/recv、all_to_all）之间做出了截然不同的取舍。选错策略，轻则浪费算力，重则根本跑不起来。\n本文将从最基础的 DDP 出发，逐步覆盖 FSDP、TP、PP、SP、EP 以及前沿的 Context Parallel 和混合并行方案。读完之后，你应该能根据模型规模、硬件拓扑和训练需求，为自己的项目选出合理的并行策略组合。\n前置知识 GPU 显存模型与分布式通信基础（第 1 篇）——特别是 NCCL 通信原语和硬件拓扑 PyTorch DDP 的基本使用经验（写过 torchrun 启动的训练脚本） 了解 Transformer 的基本结构：Self-Attention、FFN、LayerNorm 先给出全景图，后面逐一展开：\nflowchart TD subgraph DP[\u0026#34;Data Parallel\u0026#34;] DDP[\u0026#34;DDP\u0026#34;] FSDP[\u0026#34;FSDP / ZeRO\u0026#34;] end subgraph MP[\u0026#34;Model Parallel\u0026#34;] TP[\u0026#34;TP (Tensor)\u0026#34;] SP[\u0026#34;SP (Sequence)\u0026#34;] end subgraph PP[\u0026#34;Pipeline Parallel\u0026#34;] PP1[\u0026#34;1F1B\u0026#34;] PP2[\u0026#34;Zero Bubble\u0026#34;] end subgraph EP[\u0026#34;Expert Parallel\u0026#34;] EP1[\u0026#34;EP (MoE)\u0026#34;] EP2[\u0026#34;DeepSeek\u0026#34;] end DP \u0026amp; MP \u0026amp; PP \u0026amp; EP --\u0026gt; HYBRID[\u0026#34;混合并行 (3D/5D)\\nTP 节点内 + FSDP 跨节点 + PP 跨节点组\u0026#34;] 经典并行策略 DDP（Distributed Data Parallel） DDP 是最简单、最常用的并行策略，核心思想只有三个字：复制模型，分数据。\n工作原理：\n每张 GPU 持有模型的完整副本（参数、梯度、优化器状态都完整存在） 训练数据通过 DistributedSampler 均分到各卡——每张卡只看自己的 mini-batch 各卡独立完成前向和反向传播，计算各自的梯度 通过 all_reduce 对梯度求和（或平均），确保所有卡拥有相同的聚合梯度 各卡执行相同的优化器更新，模型参数保持一致 flowchart TD DATA[\u0026#34;Data Batch\u0026#34;] --\u0026gt;|DistributedSampler 均分| B0[\u0026#34;B₀\u0026#34;] \u0026amp; B1[\u0026#34;B₁\u0026#34;] \u0026amp; B2[\u0026#34;B₂\u0026#34;] \u0026amp; B3[\u0026#34;B₃\u0026#34;] B0 --\u0026gt; G0[\u0026#34;GPU 0\\n完整模型\\n→ grad₀\u0026#34;] B1 --\u0026gt; G1[\u0026#34;GPU 1\\n完整模型\\n→ grad₁\u0026#34;] B2 --\u0026gt; G2[\u0026#34;GPU 2\\n完整模型\\n→ grad₂\u0026#34;] B3 --\u0026gt; G3[\u0026#34;GPU 3\\n完整模型\\n→ grad₃\u0026#34;] G0 \u0026amp; G1 \u0026amp; G2 \u0026amp; G3 --\u0026gt; AR[\u0026#34;all_reduce(gradients)\\n唯一的通信操作\u0026#34;] AR --\u0026gt; OPT[\u0026#34;optimizer.step()\\n各卡执行相同更新\u0026#34;] 显存：DDP 不节省任何显存。每张 GPU 独立存储完整的参数（$2\\Phi$ bytes FP16）、梯度（$2\\Phi$ bytes）和优化器状态（$12\\Phi$ bytes for Adam），总计 $16\\Phi$ bytes per GPU——跟单卡训练一样。\n通信：唯一的通信操作是对梯度做 all_reduce。PyTorch DDP 使用了一个重要的优化——Bucketed All-Reduce：它不会等所有梯度都算完才通信，而是将梯度分成若干个 bucket（默认 25 MB），当一个 bucket 内所有梯度就绪后立即开始 all_reduce，与后续层的反向传播重叠执行。这大幅隐藏了通信延迟。\n什么时候用 DDP：模型能放进单张 GPU 的时候。DDP 的优势是简单、高效、几乎线性扩展。它的局限也很明显——模型必须整个塞进一张卡。\n以下是 DDP 的核心代码片段：\n1 2 3 4 5 6 7 8 9 10 # DDP wrapping — 注册梯度同步 hook model = DDP(model, device_ids=[local_rank]) # 训练循环与单卡完全一样 for input_ids, target_ids in dataloader: logits = model(input_ids) loss = criterion(logits.view(-1, logits.size(-1)), target_ids.view(-1)) optimizer.zero_grad() loss.backward() # ← DDP hook 在这里自动触发 all_reduce optimizer.step() # ← 各卡梯度相同，更新后参数一致 完整代码见 code/02-parallel-strategies/ddp_example.py\nFSDP / FSDP2（Fully Sharded Data Parallel） DDP 的致命问题是：每张卡都存了一份完整的参数 + 梯度 + 优化器状态。当模型大到单卡放不下时，DDP 就无能为力了。\nFSDP（Fully Sharded Data Parallel）的核心思想来自 DeepSpeed 的 ZeRO（Zero Redundancy Optimizer）系列论文：既然每张卡都存了冗余的参数和优化器状态，为什么不把它们切分（shard） 到各卡上，需要时再临时拼回来？\nZeRO Stage 1/2/3 ZeRO 将显存优化分为三个阶段，逐步增加切分范围：\nZeRO Stage 切分内容 每卡显存 FSDP 对应策略 Stage 1 优化器状态 $4\\Phi + \\frac{12\\Phi}{N}$ — Stage 2 优化器状态 + 梯度 $2\\Phi + \\frac{14\\Phi}{N}$ SHARD_GRAD_OP Stage 3 优化器状态 + 梯度 + 参数 $\\frac{16\\Phi}{N}$ FULL_SHARD 其中 $\\Phi$ 是参数数量，$N$ 是 GPU 数量。可以看到，Stage 3 实现了近乎线性的显存缩减——8 张卡就能把显存需求降到单卡的 1/8（加上激活值的开销）。\nFSDP 的前向/反向过程 FSDP (FULL_SHARD) 的工作流程如下：\nflowchart LR subgraph FWD[\u0026#34;FSDP Forward Pass — Layer i\u0026#34;] direction LR FS1[\u0026#34;Shard\\n1/N\u0026#34;] --\u0026gt;|\u0026#34;all_gather\u0026#34;| FP1[\u0026#34;Full Params\\n(临时拼出)\u0026#34;] FP1 --\u0026gt;|\u0026#34;compute\u0026#34;| FO1[\u0026#34;Output\u0026#34;] FP1 -.-\u0026gt;|\u0026#34;丢弃 (N-1)/N\u0026#34;| X1[\u0026#34; \u0026#34;] end subgraph BWD[\u0026#34;FSDP Backward Pass — Layer i\u0026#34;] direction LR BS1[\u0026#34;Shard\\n1/N\u0026#34;] --\u0026gt;|\u0026#34;all_gather\u0026#34;| BP1[\u0026#34;Full W\\n(临时)\u0026#34;] BP1 --\u0026gt;|\u0026#34;backward\u0026#34;| BG1[\u0026#34;Grad\\n(full)\u0026#34;] BG1 --\u0026gt;|\u0026#34;reduce_scatter\u0026#34;| BGS[\u0026#34;Grad Shard\\n1/N\u0026#34;] end FWD ~~~ BWD style X1 fill:none,stroke:none 通信开销对比：\n策略 前向通信 反向通信 总通信量 DDP 无 1x all_reduce (梯度) $2\\Phi$ FSDP (FULL_SHARD) all_gather (参数) all_gather (参数) + reduce_scatter (梯度) $3\\Phi$ FSDP 的通信量约为 DDP 的 1.5 倍——这就是用通信换显存的代价。但在大模型场景下，这个 trade-off 非常值得：没有 FSDP，根本跑不起来。\nFSDP2：PyTorch 2.2+ 的 Composable API PyTorch \u0026gt;= 2.2 引入了 FSDP2（torch.distributed._composable.fsdp），相比 FSDP1 的主要改进：\nPer-parameter sharding：不再要求整个 module 作为 FSDP 单元，可以对单个参数做 sharding Composable：可以与 TP、PP 等其他并行策略自由组合，不需要嵌套包装 更灵活的 sharding 粒度：不同的层可以使用不同的 sharding 策略 核心概念与 FSDP1 完全一致（all_gather/reduce_scatter 的通信模式不变），API 更现代、与 PyTorch 2.x 的编译器栈更兼容。\n以下是 FSDP 不同策略的显存对比代码片段：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy # 三种策略对比 strategies = [ (ShardingStrategy.NO_SHARD, \u0026#34;NO_SHARD (= DDP)\u0026#34;), (ShardingStrategy.SHARD_GRAD_OP, \u0026#34;SHARD_GRAD_OP (= ZeRO-2)\u0026#34;), (ShardingStrategy.FULL_SHARD, \u0026#34;FULL_SHARD (= ZeRO-3)\u0026#34;), ] for strategy, name in strategies: model = FSDP( TransformerLM(), sharding_strategy=strategy, device_id=local_rank, ) # ... 训练并测量显存 完整代码见 code/02-parallel-strategies/fsdp_example.py\nTP（Tensor Parallel）— Megatron 列/行切分 FSDP 将参数\u0026quot;切碎\u0026quot;再\u0026quot;拼回\u0026quot;，通信和计算是串行的。Tensor Parallel 更进一步：直接把每一层的权重矩阵按维度切开，分给不同 GPU 各算一部分，从根本上减少单卡的计算量和显存。\nTP 的核心思想来自 Megatron-LM，定义了两种基本的并行线性层：\nColumnParallelLinear：按输出维度切分 对于线性层 $Y = XW + b$，将权重 $W \\in \\mathbb{R}^{d \\times h}$ 按列切分：\nflowchart TD W[\u0026#34;完整权重 W: (d_model × dim_ffn)\u0026#34;] W --\u0026gt;|\u0026#34;按列切分到 2 张 GPU\u0026#34;| G0 \u0026amp; G1 G0[\u0026#34;GPU 0: W₀ = W[:, :dim_ffn//2]\\nY₀ = X @ W₀\u0026#34;] G1[\u0026#34;GPU 1: W₁ = W[:, dim_ffn//2:]\\nY₁ = X @ W₁\u0026#34;] NOTE[\u0026#34;每张 GPU 独立计算，无需通信！\\n(因为输入 X 在所有 GPU 上是相同的)\u0026#34;] style NOTE fill:#d4edda,stroke:#155724 关键优势：前向传播不需要任何通信。每张 GPU 得到输出的一个分片（chunk），可以直接送入后续的逐元素操作（如 GeLU）。\nRowParallelLinear：按输入维度切分 将权重 $W \\in \\mathbb{R}^{h \\times d}$ 按行切分：\nflowchart TD subgraph SPLIT[\u0026#34;按行切分到 2 张 GPU\u0026#34;] G0[\u0026#34;GPU 0: W₀ = W[:dim_ffn//2, :]\\nY₀ = X₀ @ W₀ (部分和)\u0026#34;] G1[\u0026#34;GPU 1: W₁ = W[dim_ffn//2:, :]\\nY₁ = X₁ @ W₁ (部分和)\u0026#34;] end SPLIT --\u0026gt;|\u0026#34;all_reduce\u0026#34;| RESULT[\u0026#34;Y = Y₀ + Y₁\u0026#34;] style RESULT fill:#fff3cd,stroke:#856404 关键约束：每张 GPU 只计算了输出的一个部分和（partial sum），必须通过 all_reduce 才能得到完整输出。\nMegatron FFN：Column + Row = 只需 1 次 all_reduce Megatron-LM 的天才设计在于将 Column 和 Row 配对使用：\nflowchart TD X[\u0026#34;Input X\\n(每张 GPU 上相同)\u0026#34;] X --\u0026gt; COL[\u0026#34;ColumnParallelLinear\\n(W1 按列切, 无通信)\u0026#34;] COL --\u0026gt; GELU[\u0026#34;GeLU\\n(逐元素, 无通信)\u0026#34;] GELU --\u0026gt; ROW[\u0026#34;RowParallelLinear\\n(W2 按行切)\u0026#34;] ROW --\u0026gt;|\u0026#34;all_reduce\\n合并部分和\u0026#34;| Y[\u0026#34;Output Y\\n(每张 GPU 上相同)\u0026#34;] style COL fill:#d4edda style GELU fill:#d4edda style ROW fill:#fff3cd 整个 FFN 块只需 1 次 all_reduce。\nAttention 的 TP 对于 Multi-Head Attention，TP 的做法同样优雅：\nQ、K、V 投影：使用 ColumnParallel，将 attention head 分给不同 GPU——每张 GPU 负责 $\\frac{n_heads}{TP}$ 个 head Output 投影：使用 RowParallel，将各 GPU 的 head 输出合并 这样整个 Attention 块也只需要1 次 all_reduce。一个 Transformer 层总共 2 次 all_reduce（FFN 1 次 + Attention 1 次）。\nTP 的适用条件 TP 通信频繁（每层 2 次 all_reduce），对带宽要求极高。因此：\nTP degree 通常为 2、4 或 8，部署在同一节点内（NVLink 600 GB/s） 跨节点做 TP 几乎不可行（InfiniBand 25-50 GB/s，太慢） hidden dimension 和 head 数必须能被 TP degree 整除 以下是 ColumnParallelLinear 的核心实现：\n1 2 3 4 5 6 7 8 9 10 11 12 13 class ColumnParallelLinear(nn.Module): def __init__(self, in_features, out_features, world_size, rank): super().__init__() self.out_features_per_rank = out_features // world_size # 每张 GPU 只存 1/world_size 的权重 self.weight = nn.Parameter( torch.empty(self.out_features_per_rank, in_features) ) self.bias = nn.Parameter(torch.empty(self.out_features_per_rank)) def forward(self, x): # 无通信！各 GPU 独立计算 return F.linear(x, self.weight, self.bias) 完整代码见 code/02-parallel-strategies/tensor_parallel.py\nPP（Pipeline Parallel）— 1F1B 与 Zero Bubble TP 把每一层切开，PP 的思路完全不同：按层划分，把模型的不同层分配到不同 GPU 上（称为 stage）。\nflowchart LR S0[\u0026#34;Stage 0 (GPU 0)\\nLayer 0-10\u0026#34;] --\u0026gt;|\u0026#34;send/recv\\nactivations\u0026#34;| S1[\u0026#34;Stage 1 (GPU 1)\\nLayer 11-21\u0026#34;] S1 --\u0026gt;|\u0026#34;send/recv\\nactivations\u0026#34;| S2[\u0026#34;Stage 2 (GPU 2)\\nLayer 22-31\u0026#34;] 通信方式: send/recv（点对点，仅相邻 stage 之间）。\n通信：PP 只需要相邻 stage 之间的 send/recv——传输的是中间激活值（前向）和梯度（反向），通信量远小于 TP 的 all_reduce。这使得 PP 非常适合跨节点部署。\n但 PP 有一个致命问题：Pipeline Bubble（流水线气泡）。\nNaive PP：巨大的 Bubble 最朴素的方式：整个 batch 依次通过各 stage，一个 stage 在计算时，其他 stage 全部闲置。\nflowchart TD subgraph NaivePP[\u0026#34;Naive Pipeline Parallel — Time →\u0026#34;] direction LR subgraph GPU0[\u0026#34;GPU 0\u0026#34;] direction LR F0[\u0026#34;Forward\u0026#34;] ~~~ B0[\u0026#34;Backward\u0026#34;] end subgraph GPU1[\u0026#34;GPU 1\u0026#34;] direction LR IDLE1a[\u0026#34;idle\u0026#34;] ~~~ F1[\u0026#34;Forward\u0026#34;] ~~~ IDLE1b[\u0026#34;idle\u0026#34;] ~~~ B1[\u0026#34;Backward\u0026#34;] end subgraph GPU2[\u0026#34;GPU 2\u0026#34;] direction LR IDLE2a[\u0026#34;idle\u0026#34;] ~~~ F2[\u0026#34;Forward\u0026#34;] ~~~ IDLE2b[\u0026#34;idle\u0026#34;] ~~~ B2[\u0026#34;Backward\u0026#34;] end end NOTE[\u0026#34;Bubble 占比 ≈ (P-1)/P\\n4 个 stage → 75% 时间浪费！\u0026#34;] style IDLE1a fill:#f8d7da,stroke:#dc3545 style IDLE1b fill:#f8d7da,stroke:#dc3545 style IDLE2a fill:#f8d7da,stroke:#dc3545 style IDLE2b fill:#f8d7da,stroke:#dc3545 style F0 fill:#d4edda,stroke:#28a745 style F1 fill:#d4edda,stroke:#28a745 style F2 fill:#d4edda,stroke:#28a745 style B0 fill:#cce5ff,stroke:#007bff style B1 fill:#cce5ff,stroke:#007bff style B2 fill:#cce5ff,stroke:#007bff style NOTE fill:#fff3cd,stroke:#856404 假设有 $P$ 个 stage，bubble 占比约为 $\\frac{P-1}{P}$——4 个 stage 意味着 75% 的时间在浪费！\nGPipe：Micro-batching GPipe 的解决方案：将一个 mini-batch 切成 $M$ 个 micro-batch，让多个 micro-batch 像流水线一样在各 stage 间流动。\nflowchart TD subgraph GPipe[\u0026#34;GPipe — M=4 Micro-batches, Time →\u0026#34;] direction LR subgraph G0[\u0026#34;GPU 0\u0026#34;] direction LR G0F[\u0026#34;F0 F1 F2 F3\u0026#34;] ~~~ G0GAP[\u0026#34;· · ·\u0026#34;] ~~~ G0B[\u0026#34;B3 B2 B1 B0\u0026#34;] end subgraph G1[\u0026#34;GPU 1\u0026#34;] direction LR G1F[\u0026#34;F0 F1 F2 F3\u0026#34;] ~~~ G1B[\u0026#34;B3 B2 B1 B0\u0026#34;] end subgraph G2[\u0026#34;GPU 2\u0026#34;] direction LR G2F[\u0026#34;F0 F1 F2 F3\u0026#34;] ~~~ G2B[\u0026#34;B3 B2 B1 B0\u0026#34;] end subgraph G3[\u0026#34;GPU 3\u0026#34;] direction LR G3F[\u0026#34;F0 F1 F2 F3\u0026#34;] --- G3B[\u0026#34;B3 B2 B1 B0\u0026#34;] end end NOTE[\u0026#34;Bubble: (P-1)/(M+P-1) — 当 M \u0026gt;\u0026gt; P 时, bubble → 0\u0026#34;] style G0F fill:#d4edda,stroke:#28a745 style G1F fill:#d4edda,stroke:#28a745 style G2F fill:#d4edda,stroke:#28a745 style G3F fill:#d4edda,stroke:#28a745 style G0B fill:#cce5ff,stroke:#007bff style G1B fill:#cce5ff,stroke:#007bff style G2B fill:#cce5ff,stroke:#007bff style G3B fill:#cce5ff,stroke:#007bff style G0GAP fill:#f8d7da,stroke:#dc3545 style NOTE fill:#fff3cd,stroke:#856404 Bubble 占比从 $\\frac{P-1}{P}$ 降到 $\\frac{P-1}{M+P-1}$。但 GPipe 的问题是：所有 micro-batch 的前向都做完后才开始反向，需要同时保存所有 micro-batch 的激活值，显存开销巨大。\n1F1B Schedule 1F1B（1 Forward 1 Backward）交错执行前向和反向，每做完一个 micro-batch 的前向后就尽快做反向，从而释放激活值显存：\nflowchart TD subgraph OneF1B[\u0026#34;1F1B Schedule — P=4, M=8, Time →\u0026#34;] direction TB G0[\u0026#34;GPU 0: F0 F1 F2 F3 | B0 F4 B1 F5 B2 F6 B3 F7 | B4 B5 B6 B7\u0026#34;] G1[\u0026#34;GPU 1: · F0 F1 F2 | B0 F3 B1 F4 B2 F5 B3 F6 | B4 F7 B5 B6 B7\u0026#34;] G2[\u0026#34;GPU 2: · · F0 F1 | B0 F2 B1 F3 B2 F4 B3 F5 | B4 F6 B5 F7 ...\u0026#34;] G3[\u0026#34;GPU 3: · · · F0 | B0 F1 B1 F2 B2 F3 B3 F4 | B4 F5 B5 ...\u0026#34;] end NOTE[\u0026#34;稳态阶段: 1 Forward + 1 Backward 交错\\n显存占用远小于 GPipe\u0026#34;] style NOTE fill:#d4edda,stroke:#155724 1F1B 在稳态阶段（warmup 结束后），每个 GPU 同时只保留有限数量的 micro-batch 的激活值，显存占用远小于 GPipe。Bubble 比例不变（仍为 $\\frac{P-1}{M+P-1}$），但显存显著改善。\nZero Bubble PP 2024 年提出的 Zero Bubble PP 进一步减少气泡。核心思想是：将反向传播分解为两部分——计算输入梯度（B） 和 计算权重梯度（W）。B 需要传递给前一个 stage，但 W 不需要通信，可以填充到 bubble 中：\nflowchart TD subgraph ZB[\u0026#34;Zero Bubble PP — Time →\u0026#34;] direction TB G0[\u0026#34;GPU 0: F · F · F · B · F · B · W · B · W · B · W\u0026#34;] G1[\u0026#34;GPU 1: · F · F · B · F · B · W · B · W · B · W\u0026#34;] end NOTE[\u0026#34;W = 权重梯度计算 (不需要通信)\\nW 填充了原本的 bubble → 气泡率 ≈ 0\u0026#34;] style NOTE fill:#d4edda,stroke:#155724 Zero Bubble PP 在理论上可以将 bubble 率降到接近零，代价是实现复杂度增加和更精细的调度。\nPP 小结：\n调度策略 Bubble 占比 激活值显存 实现复杂度 Naive $(P-1)/P$ 低 低 GPipe $(P-1)/(M+P-1)$ 高（所有 micro-batch） 中 1F1B $(P-1)/(M+P-1)$ 中（有限 micro-batch） 中 Zero Bubble $\\approx 0$ 中 高 SP（Sequence Parallel）— LayerNorm/Dropout 序列维度切分 TP 将线性层的权重按维度切分，但 Transformer 中还有不少操作是逐元素的，不涉及权重矩阵——比如 LayerNorm、Dropout、残差连接。这些操作在 TP 下的问题是：它们需要在完整的隐藏维度上执行，因此每张 GPU 都要持有完整的激活值，造成激活值显存的冗余。\nSequence Parallel（SP）解决这个问题的方式是：对于这些非 TP 操作，改为在序列维度上切分：\nflowchart TD LN1[\u0026#34;LayerNorm — SP: seq/N\u0026#34;] AG1[\u0026#34;all_gather: 拼出完整序列\u0026#34;] ATT[\u0026#34;Attention — TP: 按 head 切分\u0026#34;] RS1[\u0026#34;reduce_scatter: 切分到序列维度\u0026#34;] DR1[\u0026#34;Dropout + Residual — SP: seq/N\u0026#34;] LN2[\u0026#34;LayerNorm — SP\u0026#34;] AG2[\u0026#34;all_gather: 拼出完整序列\u0026#34;] FFN[\u0026#34;FFN — TP: Column + Row\u0026#34;] RS2[\u0026#34;reduce_scatter: 切回序列维度\u0026#34;] DR2[\u0026#34;Dropout + Residual — SP\u0026#34;] LN1 --\u0026gt; AG1 --\u0026gt; ATT --\u0026gt; RS1 --\u0026gt; DR1 --\u0026gt; LN2 --\u0026gt; AG2 --\u0026gt; FFN --\u0026gt; RS2 --\u0026gt; DR2 style LN1 fill:#d4edda style DR1 fill:#d4edda style LN2 fill:#d4edda style DR2 fill:#d4edda style ATT fill:#fff3cd style FFN fill:#fff3cd style AG1 fill:#cce5ff style AG2 fill:#cce5ff style RS1 fill:#cce5ff style RS2 fill:#cce5ff 关键洞察：TP 中的 all_reduce 被拆分为 reduce_scatter + all_gather，分别放在 TP 区域的出口和入口。这样通信量不变（all_reduce = reduce_scatter + all_gather），但 SP 区域的激活值显存降到了 $1/N$。\nSP 的价值在大 batch、长序列时尤为显著——此时激活值是显存的主要来源，SP 直接将这部分开销除以 TP degree。\nMoE 并行 EP（Expert Parallel） Mixture-of-Experts（MoE）在模型中引入了一组\u0026quot;专家\u0026quot;子网络，每个 token 只激活其中的 $k$ 个（通常 $k=1$ 或 $2$）。这使得 MoE 可以在显著增大参数量的同时保持计算量基本不变——但也带来了独特的并行挑战。\nMoE 层的结构：\nflowchart TD INPUT[\u0026#34;Input tokens\u0026#34;] ROUTER[\u0026#34;Router 决定每个 token 发给哪个 expert\u0026#34;] E0[\u0026#34;Expert 0 FFN\u0026#34;] \u0026amp; E1[\u0026#34;Expert 1 FFN\u0026#34;] \u0026amp; E2[\u0026#34;Expert 2 FFN\u0026#34;] \u0026amp; E3[\u0026#34;Expert 3 FFN ...\u0026#34;] COMBINE[\u0026#34;Combine outputs\u0026#34;] INPUT --\u0026gt; ROUTER ROUTER --\u0026gt; E0 \u0026amp; E1 \u0026amp; E2 \u0026amp; E3 E0 \u0026amp; E1 \u0026amp; E2 \u0026amp; E3 --\u0026gt; COMBINE Expert Parallel（EP） 将不同的 expert 分配到不同 GPU 上。如果有 64 个 expert 和 8 张 GPU，每张 GPU 负责 8 个 expert。\n核心通信操作是 all_to_all，执行两次：\nDispatch（分发）：Router 决定每个 token 要去哪个 expert 后，通过 all_to_all 将 token 从\u0026quot;按数据分片\u0026quot;的分布重新排列为\u0026quot;按 expert 分组\u0026quot;的分布——每张 GPU 收到所有发给它持有的 expert 的 token Combine（回收）：各 expert 计算完成后，通过 all_to_all 将结果送回原来的 GPU flowchart TD subgraph BEFORE[\u0026#34;Before: 按数据分片\u0026#34;] direction LR G0B[\u0026#34;GPU 0 tokens → 各 Expert\u0026#34;] G1B[\u0026#34;GPU 1 tokens → 各 Expert\u0026#34;] end BEFORE --\u0026gt;|\u0026#34;all_to_all (dispatch)\u0026#34;| AFTER subgraph AFTER[\u0026#34;After: 按 expert 分组\u0026#34;] direction LR G0A[\u0026#34;GPU 0 (E0,E1) 收到所有发往 E0,E1 的 token\u0026#34;] G1A[\u0026#34;GPU 1 (E2,E3) 收到所有发往 E2,E3 的 token\u0026#34;] end AFTER --\u0026gt;|\u0026#34;expert 计算\u0026#34;| COMPUTE[\u0026#34;各 GPU 运行自己的 expert\u0026#34;] COMPUTE --\u0026gt;|\u0026#34;all_to_all (combine)\u0026#34;| RESULT[\u0026#34;恢复原始数据分片\u0026#34;] EP 的通信开销取决于 token 的路由分布——如果 token 均匀分散到各 expert，all_to_all 的通信量最大；如果 token 集中在少数 expert，通信量较小但会导致负载不均。\nDeepSeek MoE：All-to-All Dispatch 与 Token Dropping DeepSeek 在 MoE 架构上做了几个重要改进：\n1. Fine-grained Experts（细粒度专家）\n传统 MoE（如 Switch Transformer）使用少量大 expert，DeepSeek 使用大量小 expert——例如 160 个 expert 每个 token 激活 6 个，而非 16 个 expert 每个 token 激活 2 个。更多更小的 expert 提供了更灵活的组合能力：\n$$\\binom{160}{6} \\gg \\binom{16}{2}$$\ntoken 可以组合出的 expert 组合数指数级增加，表达能力更强。\n2. Shared Experts + Routed Experts\nDeepSeek MoE 引入了\u0026quot;共享专家\u0026quot;——每个 token 都会经过的 expert，加上通过 router 选择的 expert：\nOutput = SharedExpert(x) + Σ Router_topk(RoutedExpert_i(x)) flowchart TD subgraph NON_EXPERT[\u0026#34;Non-expert layers (Attention, LN, Embedding)\u0026#34;] FSDP_SHARD[\u0026#34;FSDP across all GPUs shard params/grads/opt_state\u0026#34;] end subgraph MOE_LAYERS[\u0026#34;MoE layers\u0026#34;] direction LR EP_0[\u0026#34;Expert 0-7 → GPU 0\u0026#34;] EP_1[\u0026#34;Expert 8-15 → GPU 1\u0026#34;] EP_N[\u0026#34;... (EP)\u0026#34;] end COMM[\u0026#34;通信: FSDP: all_gather + reduce_scatter EP: all_to_all (dispatch/combine)\u0026#34;] NON_EXPERT --- MOE_LAYERS --- COMM PyTorch 2.x 的 FSDP2 和 DTensor 框架为这种混合并行提供了原生支持。\n前沿方案 Context Parallel：Ring Attention / Stripe Attention 随着 LLM 处理越来越长的上下文（128K、1M tokens），序列长度成为了新的显存瓶颈。Self-Attention 的显存和计算复杂度为 $O(S^2)$（$S$ 为序列长度），一条 128K 的序列在 FlashAttention 下仍然需要巨量的显存来存储 KV 。\nContext Parallel（CP）的解决方案：将长序列切分到多张 GPU 上，每张 GPU 只处理序列的一个片段，通过通信交换 KV 来完成完整的 attention 计算。\nRing Attention Ring Attention 的核心思想与 Ring All-Reduce 类似：将 GPU 排成一个环，每张 GPU 持有一段序列的 Q，同时 KV block 在环上循环传递：\nflowchart LR subgraph RING[\u0026#34;Ring Attention — 序列长度 S, 4 GPUs\u0026#34;] G0[\u0026#34;GPU 0 Q[0:S/4]\u0026#34;] --\u0026gt;|\u0026#34;KV\u0026#34;| G1[\u0026#34;GPU 1 Q[S/4:S/2]\u0026#34;] G1 --\u0026gt;|\u0026#34;KV\u0026#34;| G2[\u0026#34;GPU 2 Q[S/2:3S/4]\u0026#34;] G2 --\u0026gt;|\u0026#34;KV\u0026#34;| G3[\u0026#34;GPU 3 Q[3S/4:S]\u0026#34;] G3 --\u0026gt;|\u0026#34;KV\u0026#34;| G0 end KV 在环上循环传递，每步 send/recv 与 attention 计算重叠执行。\n关键优化：KV 的 send/recv 与 attention 计算可以重叠——当 GPU 在用当前 KV block 计算 attention 时，已经在传输下一个 KV block 了。\n显存：每张 GPU 只存 $S/N$ 长度的 Q 和对应的 KV，显存从 $O(S^2)$ 降到 $O(S^2/N)$（更准确地说，FlashAttention 下从 $O(S)$ 降到 $O(S/N)$）。\nStripe Attention Ring Attention 的一个问题是 causal mask 导致的负载不均衡：排在序列前面的 GPU 需要计算更少的 attention（因为 causal mask 屏蔽了后续位置），导致最后一个 GPU 的计算量最大。\nStripe Attention 通过交错分配序列位置来解决：\nflowchart TD subgraph RING[\u0026#34;Ring Attention (连续分配)\u0026#34;] direction TB RG0[\u0026#34;GPU 0: tokens 0,1,2,3 → 计算量最少 (causal)\u0026#34;] RG1[\u0026#34;GPU 1: tokens 4,5,6,7\u0026#34;] RG2[\u0026#34;GPU 2: tokens 8,9,10,11\u0026#34;] RG3[\u0026#34;GPU 3: tokens 12,13,14,15 → 计算量最多\u0026#34;] end subgraph STRIPE[\u0026#34;Stripe Attention (交错分配)\u0026#34;] direction TB SG0[\u0026#34;GPU 0: tokens 0,4,8,12 → 计算量均衡\u0026#34;] SG1[\u0026#34;GPU 1: tokens 1,5,9,13 → 计算量均衡\u0026#34;] SG2[\u0026#34;GPU 2: tokens 2,6,10,14 → 计算量均衡\u0026#34;] SG3[\u0026#34;GPU 3: tokens 3,7,11,15 → 计算量均衡\u0026#34;] end style RG0 fill:#d4edda,stroke:#155724 style RG3 fill:#f8d7da,stroke:#dc3545 style SG0 fill:#d4edda,stroke:#155724 style SG1 fill:#d4edda,stroke:#155724 style SG2 fill:#d4edda,stroke:#155724 style SG3 fill:#d4edda,stroke:#155724 flowchart TD IN[\u0026#34;输入: 每 GPU 持有 seq/N, 完整 heads\u0026#34;] QKV[\u0026#34;QKV 投影 (本地计算)\u0026#34;] A2A1[\u0026#34;all_to_all (seq/N, heads) → (seq, heads/N)\u0026#34;] ATT[\u0026#34;Attention 每 GPU: heads/N 的完整序列\u0026#34;] A2A2[\u0026#34;all_to_all (seq, heads/N) → (seq/N, heads)\u0026#34;] OUT[\u0026#34;Output 投影 (本地计算)\u0026#34;] IN --\u0026gt; QKV --\u0026gt; A2A1 --\u0026gt; ATT --\u0026gt; A2A2 --\u0026gt; OUT style A2A1 fill:#cce5ff style A2A2 fill:#cce5ff Ulysses vs Ring Attention：\n特性 Ring Attention Ulysses SP 通信模式 send/recv (P2P, 多步) all_to_all (两次) 通信量 $O(S \\cdot d)$ $O(S \\cdot d)$ 通信-计算重叠 容易（ring 结构天然重叠） 较难 对 head 数的要求 无 head 数必须被 SP degree 整除 适用场景 极长序列，head 数不够分 head 数足够时更带宽高效 两种方案在不同配置下各有优劣，实际系统中有时会结合使用。\n混合并行（Hybrid Parallelism） 现实中的大模型训练几乎不会只用单一并行策略——而是将多种策略组合，根据硬件拓扑分层部署。这就是所谓的 3D 并行甚至 5D 并行。\n经典 3D 并行 最基础的混合方案是 TP + PP + DP（或 FSDP）：\nflowchart TD subgraph N0[\u0026#34;Node 0 (8 GPUs) — NVLink 600 GB/s\u0026#34;] direction LR subgraph TP0A[\u0026#34;TP=4\u0026#34;] G0[\u0026#34;0\u0026#34;] ~~~ G1[\u0026#34;1\u0026#34;] ~~~ G2[\u0026#34;2\u0026#34;] ~~~ G3[\u0026#34;3\u0026#34;] end subgraph TP0B[\u0026#34;TP=4\u0026#34;] G4[\u0026#34;4\u0026#34;] ~~~ G5[\u0026#34;5\u0026#34;] ~~~ G6[\u0026#34;6\u0026#34;] ~~~ G7[\u0026#34;7\u0026#34;] end TP0L[\u0026#34;Stage 0 — DP/FSDP across\u0026#34;] end subgraph N1[\u0026#34;Node 1 (8 GPUs)\u0026#34;] direction LR TP1L[\u0026#34;Stage 1 — DP/FSDP across\u0026#34;] end N0 ==\u0026gt;|\u0026#34;PP: send/recv (InfiniBand)\u0026#34;| N1 style N0 fill:#f0f0f0 style N1 fill:#f0f0f0 TP: 节点内 (NVLink) | PP: 跨节点组 (InfiniBand) | FSDP: 跨节点 (InfiniBand)\n5D 并行 加上 SP（Sequence Parallel）和 EP（Expert Parallel），就形成了所谓的 5D 并行：\n$$\\text{Total GPUs} = TP \\times SP \\times PP \\times DP \\times EP$$\n每种并行在不同维度上切分：\n并行策略 切分维度 通信操作 通信量 适合的互联层级 TP 隐藏维度 all_reduce 高 节点内 (NVLink) SP 序列维度 reduce_scatter / all_gather 中 节点内 PP 层 send/recv 低 节点内或跨节点 DP/FSDP 数据 all_reduce / all_gather + reduce_scatter 中 跨节点 EP Expert all_to_all 取决于路由 跨节点 如何选择并行策略组合？ 实际选择取决于三个因素：模型规模、硬件拓扑、序列长度。\n以下是一个决策参考：\nflowchart TD Q1{\u0026#34;模型能放进单卡？\u0026#34;} Q1 --\u0026gt;|\u0026#34;Yes\u0026#34;| DDP[\u0026#34;DDP（最简单、最快）\u0026#34;] Q1 --\u0026gt;|\u0026#34;No\u0026#34;| Q2{\u0026#34;Adam 状态放不下\\n但参数放得下？\u0026#34;} Q2 --\u0026gt;|\u0026#34;Yes\u0026#34;| FSDP_GO[\u0026#34;FSDP (SHARD_GRAD_OP)\u0026#34;] Q2 --\u0026gt;|\u0026#34;No\u0026#34;| Q3{\u0026#34;参数都放不下？\u0026#34;} Q3 --\u0026gt;|\u0026#34;Yes\u0026#34;| FSDP_FS[\u0026#34;FSDP (FULL_SHARD)\u0026#34;] Q3 --\u0026gt;|\u0026#34;MoE 模型\u0026#34;| MOE[\u0026#34;EP for experts\\n+ FSDP for non-expert params\u0026#34;] FSDP_FS --\u0026gt; Q4{\u0026#34;仍然 OOM？\u0026#34;} Q4 --\u0026gt;|\u0026#34;Yes\u0026#34;| TP[\u0026#34;加 TP (通常 2/4/8 within node)\u0026#34;] TP --\u0026gt; Q5{\u0026#34;还是不够？\u0026#34;} Q5 --\u0026gt;|\u0026#34;Yes\u0026#34;| PP[\u0026#34;加 PP (跨节点)\u0026#34;] FSDP_FS --\u0026gt; Q6{\u0026#34;序列太长导致 OOM？\u0026#34;} Q6 --\u0026gt;|\u0026#34;Yes\u0026#34;| CP[\u0026#34;加 Context Parallel\u0026#34;] style DDP fill:#d4edda,stroke:#155724 style FSDP_GO fill:#d4edda,stroke:#155724 style FSDP_FS fill:#cce5ff,stroke:#004085 style MOE fill:#fff3cd,stroke:#856404 style TP fill:#cce5ff,stroke:#004085 style PP fill:#cce5ff,stroke:#004085 style CP fill:#cce5ff,stroke:#004085 一些经验法则：\n7B 模型：2-8 GPU DDP 或 FSDP 就够了 13B-70B 模型：FSDP + TP（节点内 TP=2 或 4） 70B+ 模型：FSDP + TP + PP，完整的 3D 并行 MoE 模型（如 Mixtral、DeepSeek）：EP + FSDP + TP 超长序列（128K+）：Context Parallel + TP + FSDP 配套代码 本文配套代码位于 code/02-parallel-strategies/：\nddp_example.py — DDP 完整训练循环，包含 DistributedSampler、bucketed gradient sync、吞吐量测量和参数一致性验证。运行方式：torchrun --nproc_per_node=2 ddp_example.py\nfsdp_example.py — 对比 NO_SHARD（=DDP）、SHARD_GRAD_OP（=ZeRO-2）、FULL_SHARD（=ZeRO-3）三种策略的实际显存差异。运行后可以直观看到 FULL_SHARD 的显存节省效果。运行方式：torchrun --nproc_per_node=2 fsdp_example.py\ntensor_parallel.py — 从零实现 Megatron 风格的 ColumnParallelLinear 和 RowParallelLinear，并组合成 TensorParallelFFN。包含正确性验证和显存分析。运行方式：torchrun --nproc_per_node=2 tensor_parallel.py\n所有代码支持 CPU 模式（自动使用 gloo backend），方便在没有 GPU 的环境下学习逻辑。但显存测量和性能数据仅在 CUDA 环境下有意义。\n总结与下一步 本文系统梳理了大模型训练中的所有主流并行策略。让我们回顾核心要点：\n经典并行策略：\nDDP：复制模型、分数据、all_reduce 梯度——最简单，但不省显存 FSDP/ZeRO：切分参数+梯度+优化器状态，用 all_gather/reduce_scatter 通信换显存——大模型训练的基石 TP：将权重矩阵按维度切开（Column+Row），每层 2 次 all_reduce——需要 NVLink 高带宽，适合节点内 PP：按层划分，send/recv 通信——通信量小但有 pipeline bubble，1F1B 和 Zero Bubble 在努力消除 SP：序列维度切分 LayerNorm/Dropout，与 TP 互补——降低激活值显存 MoE 并行：\nEP：将 expert 分布到不同 GPU，all_to_all 做 token dispatch——负载均衡是关键挑战 DeepSeek MoE：细粒度 expert + 共享 expert + token dropping 前沿方案：\nContext Parallel（Ring/Stripe Attention）：处理超长序列 Ulysses SP：all_to_all 做序列-head 维度转换 混合并行（3D/5D）：根据硬件拓扑分层组合 核心洞察：每种并行策略本质上都是在显存和通信之间做 trade-off。选择哪种组合，取决于你的模型有多大、卡间带宽有多快、序列有多长。没有银弹，只有工程上的最优权衡。\n下一篇文章**《LLM 推理系统架构》**将把视角从训练转向推理——当模型训好之后，如何高效地服务请求？我们将深入 PagedAttention、RadixAttention（SGLang 的核心创新）、Continuous Batching 等推理优化技术，看看推理系统如何解决一个全新的 Memory-bound 挑战。\n参考资料 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al., 2020 — FSDP 的理论基础，定义了 Stage 1/2/3 的显存切分策略 Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Shoeybi et al., 2020 — Tensor Parallel 的 Column/Row 切分方案 Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — Narayanan et al., 2021 — 3D 并行（TP+PP+DP）的系统设计 GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism — Huang et al., 2019 — Pipeline Parallel 的 micro-batching 方案 Zero Bubble Pipeline Parallelism — Qi et al., 2024 — 通过分离 B 和 W 计算消除 pipeline bubble Ring Attention with Blockwise Transformers for Near-Infinite Context — Liu et al., 2023 — 长序列的 Ring Attention 方案 DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model — DeepSeek AI, 2024 — 细粒度 MoE + 共享 expert 架构 DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale — Rajbhandari et al., 2022 — EP 与 FSDP 联合训练 Reducing Activation Recomputation in Large Transformer Models — Korthikanti et al., 2023 — Sequence Parallel 减少激活值显存 PyTorch FSDP Documentation — pytorch.org/docs/stable/fsdp — FSDP/FSDP2 官方文档 DeepSpeed ZeRO Tutorial — deepspeed.ai/tutorials/zero — ZeRO Stage 1/2/3 实践指南 ","permalink":"https://mzf666.github.io/llm-infra/zh/posts/02-parallel-strategies/","summary":"从 DDP 到混合并行，系统梳理大模型训练中的所有并行策略。","title":"分布式并行策略全景"},{"content":"Motivation 训练一个大模型很贵，但服务一个大模型可能更贵。\n训练是一次性开销——几千张 GPU 跑几周就结束了。但推理是 7×24 小时不间断的：每一次用户输入、每一个 API 调用，都需要模型实时响应。OpenAI 每天处理数十亿个 token，推理成本占运营总成本的大头。如果推理效率翻倍，意味着成本直接减半——或者同样预算下服务两倍的用户。\n推理和训练的优化方向截然不同。训练追求总计算吞吐——尽快把所有数据跑完；推理则需要同时优化两个相互矛盾的指标：\n吞吐量（Throughput）：每秒生成多少个 token？决定了能服务多少并发用户 延迟（Latency）：每个用户等多久？TTFT（Time To First Token，首 token 延迟）和 TPOT（Time Per Output Token，逐 token 延迟）是关键指标 更重要的是，推理的计算特性和训练完全不同。我们在第一篇中分析过：自回归解码（autoregressive decoding）本质上是矩阵-向量乘法，算术强度只有约 1 FLOPs/Byte，GPU 算力利用率不到 1%。推理不是算力瓶颈，是显存带宽瓶颈——而 KV Cache 是其中最大的那块显存。\n本文以 SGLang 为例，深入剖析推理系统的四个核心技术：\nKV Cache——推理显存的大头，为什么它会成为瓶颈 PagedAttention——借鉴 OS 虚拟内存思想，解决 KV Cache 的显存碎片问题 RadixAttention——SGLang 的核心创新，用 Radix Tree 实现 KV Cache 复用 Continuous Batching——迭代级调度，最大化 GPU 利用率 Chunked Prefill——解决长 prompt 阻塞解码的问题 读完之后，你应该能理解为什么 vLLM 和 SGLang 能比朴素实现快 2-10 倍，以及它们之间的设计差异。\n前置知识 GPU 显存模型与分布式通信基础（第 1 篇）——特别是 HBM 带宽瓶颈和 Roofline 模型 Transformer 注意力机制的基本理解（Self-Attention、Q/K/V 投影） 了解自回归生成的基本流程（逐 token 生成） KV Cache 基础 为什么需要 KV Cache 自回归语言模型在生成第 $t$ 个 token 时，Attention 的计算公式是：\n$$\\text{Attention}(Q_t, K_{1:t}, V_{1:t}) = \\text{softmax}\\left(\\frac{Q_t \\cdot K_{1:t}^T}{\\sqrt{d_k}}\\right) V_{1:t}$$\n关键点：当前 token 的 Query 需要和所有已生成 token 的 Key、Value 做内积。如果每次都从头计算所有 token 的 K 和 V，复杂度是 $O(t^2)$——生成 1000 个 token 就要做 $1 + 2 + \\cdots + 1000 = 500500$ 次 attention 计算。\nKV Cache 的核心思想：之前 token 的 K 和 V 不会变，只需要计算一次，缓存起来复用。\nflowchart TD subgraph NO[\u0026#34;自回归解码（无 KV Cache）\u0026#34;] N1[\u0026#34;Step 1: 输入 [t1] → 计算 K1,V1 → Attention → 输出 t2\u0026#34;] N2[\u0026#34;Step 2: 输入 [t1,t2] → 重新计算 K1,V1,K2,V2 → Attention\u0026#34;] N3[\u0026#34;Step 3: 输入 [t1,t2,t3] → 重新计算 K1,K2,K3... → Attention\u0026#34;] NN[\u0026#34;Step N: 输入 [t1,...,tN] → 重新计算所有 K,V → O(N²) 总计算\u0026#34;] N1 --\u0026gt; N2 --\u0026gt; N3 -.-\u0026gt; NN end subgraph YES[\u0026#34;自回归解码（有 KV Cache）\u0026#34;] Y1[\u0026#34;Step 1: 输入 [t1] → 计算 K1,V1 → 缓存 → Attention → 输出 t2\u0026#34;] Y2[\u0026#34;Step 2: 输入 [t2] → 计算 K2,V2 → 追加缓存 → Attention\u0026#34;] Y3[\u0026#34;Step 3: 输入 [t3] → 计算 K3,V3 → 追加缓存 → Attention\u0026#34;] YN[\u0026#34;Step N: 输入 [tN] → 计算 KN,VN → 追加缓存 → O(N) 总计算\u0026#34;] Y1 --\u0026gt; Y2 --\u0026gt; Y3 -.-\u0026gt; YN YNOTE[\u0026#34;每步只计算 1 个 token 的 K,V，复用缓存中的所有 K,V\u0026#34;] end NO --\u0026gt; YES style NO fill:#fff3cd,stroke:#856404 style YES fill:#d4edda,stroke:#155724 KV Cache 将总计算量从 $O(N^2)$ 降到了 $O(N)$——这是推理优化的第一性原理，所有现代推理引擎都必须实现它。\nKV Cache 有多大 每个 token 的 KV Cache 大小为：\n$$\\text{KV per token} = 2 \\times n_{\\text{layers}} \\times n_{\\text{heads}} \\times d_{\\text{head}} \\times \\text{dtype_size}$$\n其中因子 2 是因为要存 K 和 V 两份。以几个主流模型为例（FP16，2 bytes per element）：\n模型 $n_{\\text{layers}}$ $n_{\\text{heads}}$ $d_{\\text{head}}$ KV per token 2048 tokens LLaMA-7B 32 32 128 0.5 MB 1.0 GB LLaMA-13B 40 40 128 0.8 MB 1.6 GB LLaMA-70B 80 64 128 2.5 MB 5.0 GB GPT-3 175B 96 96 128 4.5 MB 9.2 GB 一个 LLaMA-70B 的请求如果生成 2048 个 token，仅 KV Cache 就占 5 GB。如果你想同时服务 16 个这样的请求，KV Cache 就需要 80 GB——一整张 A100 的显存。而且模型权重本身还要 140 GB（FP16）。\n这就是推理显存的核心矛盾：模型权重是固定开销，KV Cache 是动态开销，而 KV Cache 随并发请求数和序列长度线性增长，很容易成为瓶颈。\n能同时 batch 多少请求，取决于 KV Cache 能放多少——batch size 直接决定吞吐量。所以 KV Cache 的内存管理效率，就是推理系统的性能上限。\n朴素分配的浪费 最直观的做法是：为每个请求预分配 max_seq_len 大小的连续显存来存 KV Cache。问题是，你不知道一个请求到底会生成多少 token——可能是 10 个，也可能是 2000 个。\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 class NaiveKVCache: \u0026#34;\u0026#34;\u0026#34; Pre-allocated KV cache for a single sequence. Problem: if max_seq_len=2048 but the sequence only generates 50 tokens, we waste 2048-50 = 1998 slots of memory. Multiply by batch_size and num_layers, and you\u0026#39;re throwing away most of your GPU memory. \u0026#34;\u0026#34;\u0026#34; def __init__(self, max_seq_len: int, num_heads: int, head_dim: int): self.max_seq_len = max_seq_len self.num_heads = num_heads self.head_dim = head_dim self.current_len = 0 # Pre-allocate full tensors — this is the waste! self.k_cache = torch.zeros(max_seq_len, num_heads, head_dim, device=device) self.v_cache = torch.zeros(max_seq_len, num_heads, head_dim, device=device) 配套代码中的实验表明，当请求长度分布不均匀时（现实中几乎总是如此），朴素预分配的显存浪费率高达 60-80%：\n5 requests, max_seq_len=2048 Actual lengths: [37, 152, 8, 420, 91] Memory allocated: 160.0 MB Memory actually used: 22.1 MB Waste: 86.2% --\u0026gt; This is why we need PagedAttention! 显存浪费 86%——这意味着原本能同时处理 7 个请求的 GPU，实际上只能处理 1 个。浪费的不仅是内存，更是吞吐量和收入。\nPagedAttention 从操作系统借鉴的智慧 PagedAttention 是 vLLM（OSDI'23）的核心贡献，它的灵感来自一个成熟了 50 年的技术——操作系统的虚拟内存。\n在操作系统中，进程看到的是虚拟地址空间——连续的、独立的。但物理内存是以固定大小的页（page） 分配的，通过页表（page table） 将虚拟地址映射到物理地址。进程不需要拿到连续的物理内存，操作系统在幕后把碎片化的物理页拼成看似连续的虚拟空间。\nPagedAttention 把这个思想完美移植到 KV Cache 管理上：\nflowchart LR subgraph OS[\u0026#34;OS 虚拟内存\u0026#34;] A1[\u0026#34;进程\u0026#34;] A2[\u0026#34;虚拟页\u0026#34;] A3[\u0026#34;物理页帧\u0026#34;] A4[\u0026#34;页表\u0026#34;] A5[\u0026#34;按需分配（缺页中断）\u0026#34;] A6[\u0026#34;页大小固定（4KB）\u0026#34;] A7[\u0026#34;CoW（fork）\u0026#34;] A8[\u0026#34;页面置换（swap）\u0026#34;] end subgraph PA[\u0026#34;PagedAttention\u0026#34;] B1[\u0026#34;请求（sequence）\u0026#34;] B2[\u0026#34;逻辑 KV block\u0026#34;] B3[\u0026#34;物理 KV block\u0026#34;] B4[\u0026#34;block table\u0026#34;] B5[\u0026#34;按需分配（生成新 token 时）\u0026#34;] B6[\u0026#34;block 大小固定（16 tokens）\u0026#34;] B7[\u0026#34;CoW（beam search）\u0026#34;] B8[\u0026#34;preemption（换出请求）\u0026#34;] end A1 -.-\u0026gt; B1 A2 -.-\u0026gt; B2 A3 -.-\u0026gt; B3 A4 -.-\u0026gt; B4 A5 -.-\u0026gt; B5 A6 -.-\u0026gt; B6 A7 -.-\u0026gt; B7 A8 -.-\u0026gt; B8 style OS fill:#cce5ff,stroke:#004085 style PA fill:#d4edda,stroke:#155724 Block 分配机制 PagedAttention 将 GPU 上的 KV Cache 显存划分为一个固定大小的 block 池。每个 block 能存储固定数量的 token 的 KV（通常 16 或 32 个 token）。每个请求有一个 block table，记录其逻辑 block 到物理 block 的映射。\nflowchart TD subgraph POOL[\u0026#34;物理 Block 池（GPU 显存）\u0026#34;] B0[\u0026#34;B0\u0026lt;br/\u0026gt;Seq0\u0026#34;] B1[\u0026#34;B1\u0026lt;br/\u0026gt;Seq1\u0026#34;] B2[\u0026#34;B2\u0026lt;br/\u0026gt;Free\u0026#34;] B3[\u0026#34;B3\u0026lt;br/\u0026gt;Seq0\u0026#34;] B4[\u0026#34;B4\u0026lt;br/\u0026gt;Seq1\u0026#34;] B5[\u0026#34;B5\u0026lt;br/\u0026gt;Free\u0026#34;] B6[\u0026#34;B6\u0026lt;br/\u0026gt;Seq0\u0026#34;] B7[\u0026#34;B7\u0026lt;br/\u0026gt;Free\u0026#34;] end subgraph T0[\u0026#34;Seq 0 Block Table\u0026#34;] L0_0[\u0026#34;Logic 0 → B0\u0026#34;] L0_1[\u0026#34;Logic 1 → B3\u0026#34;] L0_2[\u0026#34;Logic 2 → B6\u0026#34;] end subgraph T1[\u0026#34;Seq 1 Block Table\u0026#34;] L1_0[\u0026#34;Logic 0 → B1\u0026#34;] L1_1[\u0026#34;Logic 1 → B4\u0026#34;] end L0_0 --\u0026gt; B0 L0_1 --\u0026gt; B3 L0_2 --\u0026gt; B6 L1_0 --\u0026gt; B1 L1_1 --\u0026gt; B4 NOTE[\u0026#34;物理 block 不需要连续，逻辑上连续即可\u0026#34;] style POOL fill:#cce5ff,stroke:#004085 style B2 fill:#f8f9fa,stroke:#6c757d style B5 fill:#f8f9fa,stroke:#6c757d style B7 fill:#f8f9fa,stroke:#6c757d 工作流程：\n新请求到达：分配一个空的 block table，不预分配任何 block 生成第 1 个 token：从 free pool 拿一个 block，存入 KV，更新 block table 当前 block 填满（写满 16 个 token 的 KV）：再从 free pool 拿一个新 block 请求结束：归还所有 block 到 free pool，立即可被其他请求复用 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 class BlockManager: \u0026#34;\u0026#34;\u0026#34; PagedAttention-style block manager. Key ideas: - Physical KV cache is a pool of fixed-size blocks. - Each sequence has a block table mapping logical -\u0026gt; physical blocks. - Blocks are allocated on demand as sequences grow. - Free blocks are recycled when sequences finish. - Copy-on-Write: shared blocks are copied only when modified. \u0026#34;\u0026#34;\u0026#34; def __init__(self, num_blocks, num_layers, num_heads, head_dim, block_size=16): self.num_blocks = num_blocks self.block_size = block_size # Physical block pool — the actual GPU memory self.k_pool = torch.zeros( num_layers, num_blocks, block_size, num_heads, head_dim, device=device ) self.v_pool = torch.zeros( num_layers, num_blocks, block_size, num_heads, head_dim, device=device ) # Track physical block metadata self.physical_blocks = [PhysicalBlock(block_id=i) for i in range(num_blocks)] self.free_block_ids: list[int] = list(range(num_blocks)) self.seq_tables: dict[int, SequenceBlockTable] = {} 从 86% 浪费到 4% 浪费 PagedAttention 的内存浪费只来自一个地方：每个请求最后一个 block 的内部碎片（internal fragmentation）。如果 block_size=16，最后一个 block 平均只用了一半，浪费约 $\\frac{1}{2} \\times \\frac{1}{N_{\\text{blocks}}}$ 的显存。\n配套代码中的对比实验结果：\nConfig: 32 heads, head_dim=128, max_seq=2048, block_size=16 10 requests with lengths: [47, 183, 12, 891, 256, 5, 1024, 73, 330, 15] Metric Naive Paged Actual ----------------------------------------------------------------------- KV cache memory (MB) 320.0 45.5 44.3 Memory utilization 13.8% 97.4% 100.0% Waste (MB) 275.7 1.2 0.0 Waste (%) 86.2% 2.6% 0.0% 朴素方式浪费 86%，PagedAttention 只浪费 2.6%。这意味着同样的 GPU 显存可以 batch 多 2-4 倍的请求，直接把吞吐量提高 2-4 倍。\nCopy-on-Write：高效的序列分叉 在 beam search 和 parallel sampling 场景下，一个请求会分叉出多个候选序列。这些序列共享前缀的 KV Cache——如果每个候选都复制一份前缀的 KV，显存开销会翻倍。\nPagedAttention 的解决方案是 Copy-on-Write（CoW），同样借鉴自操作系统的 fork() 系统调用：\nflowchart TD subgraph BEFORE[\u0026#34;fork 前\u0026#34;] S0a[\u0026#34;Seq 0: Block A → Block B → Block C (ref=1)\u0026#34;] end subgraph AFTER[\u0026#34;fork 后（Seq 0 + Seq 1）\u0026#34;] S0b[\u0026#34;Seq 0: Block A → Block B → Block C\u0026#34;] S1b[\u0026#34;Seq 1: Block A → Block B → Block C\u0026#34;] REF[\u0026#34;Block C ref=2（共享，无拷贝！）\u0026#34;] end subgraph WRITE[\u0026#34;Seq 1 写入新 token 到 Block C\u0026#34;] S0c[\u0026#34;Seq 0: Block A → Block B → Block C (ref=1)\u0026#34;] S1c[\u0026#34;Seq 1: Block A → Block B → Block C\u0026#39; (ref=1, 此时拷贝)\u0026#34;] NOTE2[\u0026#34;只有被修改的 block 才拷贝，前面的 block 仍共享\u0026#34;] end BEFORE --\u0026gt; AFTER --\u0026gt; WRITE style BEFORE fill:#cce5ff,stroke:#004085 style AFTER fill:#fff3cd,stroke:#856404 style WRITE fill:#d4edda,stroke:#155724 核心逻辑在 append_token 中：写入前检查 block 的引用计数，如果 \u0026gt;1（被共享），就先拷贝一份新 block，再写入：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 def append_token(self, seq_id, k, v, layer_idx): table = self.seq_tables[seq_id] logical_block_idx = token_idx // self.block_size slot_in_block = token_idx % self.block_size # Allocate new block if current one is full if logical_block_idx \u0026gt;= len(table.logical_to_physical): phys_id = self._allocate_block() table.logical_to_physical.append(phys_id) phys_id = table.logical_to_physical[logical_block_idx] # CoW check: if this block is shared, copy before writing if self.physical_blocks[phys_id].ref_count \u0026gt; 1: new_phys_id = self._allocate_block() self.k_pool[:, new_phys_id] = self.k_pool[:, phys_id] self.v_pool[:, new_phys_id] = self.v_pool[:, phys_id] self._free_block(phys_id) # Decrement old block\u0026#39;s ref count table.logical_to_physical[logical_block_idx] = new_phys_id phys_id = new_phys_id # Write KV to the physical block self.k_pool[layer_idx, phys_id, slot_in_block] = k self.v_pool[layer_idx, phys_id, slot_in_block] = v CoW 的收益在 beam search 中尤其显著：beam_size=4 时，4 个候选序列共享大量前缀 KV，只有最后一个 block（正在生成的部分）需要独立存储。显存开销接近 1 个序列而非 4 个。\nRadixAttention（SGLang 核心创新） 从 PagedAttention 到 Prefix Caching PagedAttention 解决了单个请求内部的显存碎片问题。但它忽略了一个更大的优化机会：不同请求之间的 KV Cache 复用。\n在实际的 LLM 应用中，大量请求共享相同的前缀：\nflowchart TD subgraph S1[\u0026#34;场景 1: System Prompt\u0026#34;] SP[\u0026#34;共享前缀: You are a helpful assistant...\u0026#34;] UA[\u0026#34;+ 用户 A: What is machine learning?\u0026#34;] UB[\u0026#34;+ 用户 B: Write a Python function for...\u0026#34;] UC[\u0026#34;+ 用户 C: Translate this text...\u0026#34;] SP --\u0026gt; UA SP --\u0026gt; UB SP --\u0026gt; UC end subgraph S2[\u0026#34;场景 2: Few-shot Learning\u0026#34;] FS[\u0026#34;共享前缀: Here are some examples: × 5 样例\u0026#34;] Q1[\u0026#34;+ Input: new query 1\u0026#34;] Q2[\u0026#34;+ Input: new query 2\u0026#34;] FS --\u0026gt; Q1 FS --\u0026gt; Q2 end subgraph S3[\u0026#34;场景 3: Multi-turn Chat\u0026#34;] MT[\u0026#34;共享前缀: System + Turn 1 + Turn 2\u0026#34;] T3[\u0026#34;+ User Turn 3（新消息）\u0026#34;] MT --\u0026gt; T3 end style S1 fill:#cce5ff,stroke:#004085 style S2 fill:#d4edda,stroke:#155724 style S3 fill:#fff3cd,stroke:#856404 如果 system prompt 有 500 个 token，每个请求都重新计算和存储这 500 个 token 的 KV Cache，那么 100 个并发请求就重复存了 100 份——完全相同的数据。\nPrefix Caching 的核心思想：如果两个请求的 token 序列有相同的前缀，它们的 KV Cache 在这些位置上完全一样，可以直接共享，无需重新计算。\nvLLM 后来也加入了 prefix caching 功能，但它的实现基于预定义的前缀匹配——你需要显式声明哪些请求共享前缀。SGLang 的 RadixAttention 提供了一种更灵活、更优雅的方案。\nRadix Tree：索引任意前缀的数据结构 SGLang（LMSYS，2024）的核心创新是用 Radix Tree（基数树） 来索引和管理所有已缓存的 KV Cache。\nRadix Tree 是一种压缩前缀树（compressed trie）。普通 trie 每个节点代表一个字符，radix tree 将只有一个子节点的路径压缩成单个节点，减少树的深度。在 SGLang 中，每个节点代表一段 token 序列，节点上存储着对应的 KV Cache block 指针。\ngraph TD ROOT[\u0026#34;Root\u0026#34;] SYS[\u0026#34;System Prompt\u0026lt;br/\u0026gt;KV blocks: B0, B1, B2\u0026#34;] ML[\u0026#34;User: What is ML?\u0026lt;br/\u0026gt;KV blocks: B3, B4\u0026#34;] PY[\u0026#34;User: Write Python\u0026lt;br/\u0026gt;KV blocks: B5, B6\u0026#34;] AST[\u0026#34;Asst: ML is...\u0026lt;br/\u0026gt;KV blocks: B7\u0026#34;] ROOT --\u0026gt; SYS SYS --\u0026gt; ML SYS --\u0026gt; PY ML --\u0026gt; AST REQD[\u0026#34;新请求 D: System Prompt + What is ML? + new query\u0026lt;br/\u0026gt;→ 匹配 B0-B4 直接复用，只需为 new query 计算新 KV\u0026#34;] style SYS fill:#cce5ff,stroke:#004085 style ML fill:#d4edda,stroke:#155724 style PY fill:#d4edda,stroke:#155724 style AST fill:#fff3cd,stroke:#856404 style REQD fill:#f8f9fa,stroke:#6c757d 工作流程 当一个新请求到达时，SGLang 的调度器执行以下步骤：\nflowchart TD subgraph FLOW[\u0026#34;RadixAttention 处理新请求\u0026#34;] A[\u0026#34;1. PREFIX MATCHING\u0026lt;br/\u0026gt;在 Radix Tree 中查找最长匹配前缀\u0026lt;br/\u0026gt;例: sys_prompt 的 KV blocks 可直接复用\u0026#34;] B[\u0026#34;2. REUSE CACHED KV\u0026lt;br/\u0026gt;匹配到的前缀部分跳过 prefill\u0026lt;br/\u0026gt;只需计算不匹配的后缀部分的 KV\u0026#34;] C[\u0026#34;3. INSERT INTO TREE\u0026lt;br/\u0026gt;请求完成后将新 KV 插入 Radix Tree\u0026lt;br/\u0026gt;后续相同前缀的请求可复用\u0026#34;] A --\u0026gt; B --\u0026gt; C end subgraph TIMELINE[\u0026#34;时间对比\u0026#34;] T1[\u0026#34;无 cache: ====== prefill 500 tokens ====== decode\u0026#34;] T2[\u0026#34;有 cache: == prefill 100 tokens == decode\u0026lt;br/\u0026gt;400 tokens 的 KV 从 cache 复用，跳过计算\u0026#34;] end FLOW --\u0026gt; TIMELINE style FLOW fill:#cce5ff,stroke:#004085 style T1 fill:#fff3cd,stroke:#856404 style T2 fill:#d4edda,stroke:#155724 在 system prompt 场景下（几乎所有请求共享同一个 system prompt），prefix caching 可以跳过大部分 prefill 计算，将 TTFT（Time To First Token）降低 50-90%。\nLRU 逐出策略 GPU 显存有限，不可能缓存所有见过的前缀。当 block 池用满时，需要逐出一些缓存来给新请求腾空间。\nSGLang 使用 LRU（Least Recently Used） 策略：最久没被访问的 KV Cache 节点优先被逐出。\nflowchart TD subgraph NODES[\u0026#34;Radix Tree 节点（按最近访问时间排序）\u0026#34;] N1[\u0026#34;sys_prompt + user_A — 10 sec ago ✅ 保留\u0026#34;] N2[\u0026#34;sys_prompt + user_B — 30 sec ago ✅ 保留\u0026#34;] N3[\u0026#34;sys_prompt + user_C — 120 sec ago ⚠️ 候选逐出\u0026#34;] N4[\u0026#34;old_prompt + old_query — 300 sec ago ❌ 优先逐出\u0026#34;] end subgraph EVICT[\u0026#34;逐出流程\u0026#34;] E1[\u0026#34;1. 从 LRU 链表尾部取最久未访问的节点\u0026#34;] E2[\u0026#34;2. 释放该节点的 KV blocks\u0026#34;] E3[\u0026#34;3. 从 Radix Tree 中删除该节点\u0026#34;] E4[\u0026#34;4. 若父节点变成无引用叶子，递归清理\u0026#34;] E1 --\u0026gt; E2 --\u0026gt; E3 --\u0026gt; E4 end style N1 fill:#d4edda,stroke:#155724 style N2 fill:#d4edda,stroke:#155724 style N3 fill:#fff3cd,stroke:#856404 style N4 fill:#f8d7da,stroke:#721c24 LRU 的合理性在于时间局部性：最近被访问的前缀往往会再次被请求到（比如同一个用户的多轮对话，或热门的 system prompt）。\nCache-Aware Scheduling 有了 Radix Tree 之后，调度器可以做一个非常聪明的优化：优先调度那些在 cache 中匹配前缀更长的请求。\nflowchart TD subgraph QUEUE[\u0026#34;等待队列\u0026#34;] RA[\u0026#34;Req A: 匹配 500 tokens 缓存 → prefill 仅需 100 tokens\u0026#34;] RB[\u0026#34;Req B: 匹配 0 tokens → prefill 需要 600 tokens\u0026#34;] RC[\u0026#34;Req C: 匹配 300 tokens 缓存 → prefill 仅需 200 tokens\u0026#34;] end subgraph SCHED[\u0026#34;调度策略对比\u0026#34;] FCFS[\u0026#34;朴素 FCFS: A → B → C\u0026#34;] CACHE[\u0026#34;Cache-aware: A → C → B（优先 cache 命中率高的）\u0026#34;] end subgraph BENEFIT[\u0026#34;好处\u0026#34;] B1[\u0026#34;1. 高 cache 命中的请求 prefill 更快 → TTFT 更低\u0026#34;] B2[\u0026#34;2. KV 留在 tree 中 → 后续相似请求也命中 → 正反馈循环\u0026#34;] B3[\u0026#34;3. 减少总 prefill 计算量 → 更多 GPU 时间给 decode\u0026#34;] end QUEUE --\u0026gt; SCHED --\u0026gt; BENEFIT style RA fill:#d4edda,stroke:#155724 style RB fill:#f8d7da,stroke:#721c24 style RC fill:#fff3cd,stroke:#856404 style CACHE fill:#d4edda,stroke:#155724 这种调度策略使 SGLang 在 few-shot learning、multi-turn chat 等共享前缀场景下性能远超不做 prefix caching 的系统。\n与 vLLM 的对比 vLLM（v0.3+）后来也加入了 prefix caching 功能，但设计思路不同：\n特性 SGLang (RadixAttention) vLLM (Prefix Caching) 数据结构 Radix Tree Hash Map 前缀匹配 自动匹配任意共享前缀 基于 token block hash 匹配粒度 token 级别 block 级别（block_size 对齐） 逐出策略 LRU on tree nodes LRU on blocks Cache-aware 调度 原生支持 后续版本加入 多轮对话 天然支持（树结构） 需要 hash 匹配 SGLang 的 Radix Tree 方案在灵活性上更胜一筹：\n任意前缀共享：不限于预定义的 prompt template，任何两个请求只要有共同前缀就能自动匹配 树结构天然适配多轮对话：对话的历史消息自然形成一棵树，不同分支代表不同的后续对话 渐进式匹配：可以匹配到 token 级别的精确前缀，而非 block 级别的近似匹配 vLLM 的 hash-based 方案在实现简单性和大规模部署稳定性上有优势，两者在实际性能上差距取决于具体的工作负载特征。\nContinuous Batching 静态 Batching 的问题 在理解了 KV Cache 的内存管理之后，我们来看另一个维度的优化：调度策略。\n最朴素的推理方式是 静态 batching：把一组请求凑成一个 batch，一起做 prefill，然后一起做 decode，等 batch 中所有请求都生成完毕后，才处理下一个 batch。\n问题在于：不同请求的输出长度差异巨大。\nflowchart TD subgraph STATIC[\u0026#34;静态 Batching — Batch = Req 0-3, 输出长度: 3, 8, 2, 5\u0026#34;] direction LR R0[\u0026#34;Req 0: ## ## ## .. .. .. .. ..\u0026lt;br/\u0026gt;第 3 步结束\u0026#34;] R1[\u0026#34;Req 1: ## ## ## ## ## ## ## ##\u0026lt;br/\u0026gt;最长，其他人都在等\u0026#34;] R2[\u0026#34;Req 2: ## ## .. .. .. .. .. ..\u0026lt;br/\u0026gt;第 2 步结束\u0026#34;] R3[\u0026#34;Req 3: ## ## ## ## ## .. .. ..\u0026lt;br/\u0026gt;第 5 步结束\u0026#34;] end WASTE[\u0026#34;## = 有效计算 · .. = GPU 空闲\u0026lt;br/\u0026gt;GPU 利用率: 18/32 = 56.25%\u0026lt;br/\u0026gt;浪费了 43.75% 的 GPU 算力\u0026lt;br/\u0026gt;Req 4, 5 必须等 Req 1 结束才能开始\u0026#34;] style STATIC fill:#fff3cd,stroke:#856404 style WASTE fill:#f8d7da,stroke:#721c24 当输出长度方差很大时（这在实际中几乎总是如此——有人问\u0026quot;是几点了\u0026quot;，有人要写一篇长文），静态 batching 的 GPU 利用率可以低至 20-50%。\n迭代级调度 Continuous Batching（Orca, OSDI'22）的核心改进是将调度粒度从 batch 级别 降到 iteration 级别：\n每一个 decode step 结束后，立即检查：\n哪些请求已经结束（生成了 EOS 或达到 max_length）？→ 驱逐，释放 KV Cache 等待队列中有没有新请求可以加入？→ 纳入，开始 prefill flowchart TD subgraph CB[\u0026#34;Continuous Batching — 迭代级调度\u0026#34;] direction LR R0[\u0026#34;Req 0: ## ## ## — step 3 结束\u0026#34;] R1[\u0026#34;Req 1: ## ## ## ## ## ## ## ## — step 8 结束\u0026#34;] R2[\u0026#34;Req 2: ## ## — step 2 结束\u0026#34;] R3[\u0026#34;Req 3: ## ## ## ## ## — step 5 结束\u0026#34;] R4[\u0026#34;Req 4: step 3 加入 → ## ## ## ##\u0026#34;] R5[\u0026#34;Req 5: step 4 加入 → ## ## ## ## ## ##\u0026#34;] end RESULT[\u0026#34;每步活跃数: 4 4 4 4 4 3 3 3 1\u0026lt;br/\u0026gt;GPU 几乎一直满载！结束的请求立刻被新请求替换\u0026#34;] style CB fill:#d4edda,stroke:#155724 style RESULT fill:#cce5ff,stroke:#004085 配套代码中的对比实验显示，continuous batching 在吞吐量上通常有 2-3 倍的提升：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 def run_continuous_batching(requests, max_batch_size, total_blocks=256, ...): \u0026#34;\u0026#34;\u0026#34; After EVERY decode step: 1. Remove finished sequences (free their KV blocks). 2. If there are free slots AND enough KV blocks, admit waiting requests. 3. If out of KV blocks, preempt the newest running sequence (LIFO). \u0026#34;\u0026#34;\u0026#34; # ... 每步检查并调度 ... for req in running: req.tokens_generated += 1 if req.tokens_generated \u0026gt;= req.target_output_length: req.status = SequenceStatus.FINISHED finished_this_step.append(req) # Remove finished, free their blocks for req in finished_this_step: running.remove(req) block_pool.free(req.num_kv_blocks) # Try to admit waiting requests while waiting and len(running) \u0026lt; max_batch_size: req = waiting[0] if block_pool.can_allocate(blocks_needed): waiting.popleft() running.append(req) Preemption：优雅应对显存压力 Continuous batching 还引入了一个关键机制：preemption（抢占）。当 KV Cache 显存不够为所有 running 请求分配新 block 时，调度器可以：\n暂停一个或多个 running 请求（通常选最新加入的——LIFO 策略，因为它们做的工作最少） 释放被暂停请求的 KV Cache blocks 用腾出的空间继续服务其他请求 稍后恢复被暂停的请求（重新做 prefill 来重建 KV Cache） flowchart TD A[\u0026#34;正常运行\u0026lt;br/\u0026gt;Running: Req 0, 1, 2, 3\u0026lt;br/\u0026gt;KV blocks: 200/256 (78%)\u0026#34;] B[\u0026#34;Req 1 生成长回复, KV 持续增长...\u0026lt;br/\u0026gt;KV blocks: 253/256 (99%) — 即将 OOM!\u0026#34;] C[\u0026#34;触发 Preemption\u0026lt;br/\u0026gt;暂停 Req 3 (LIFO: 最新, 做的工作最少)\u0026lt;br/\u0026gt;释放 Req 3 的 20 个 blocks\u0026#34;] D[\u0026#34;KV blocks: 233/256 (91%)\u0026lt;br/\u0026gt;有空间继续了\u0026#34;] E[\u0026#34;稍后 Req 0/1 结束, 空间释放\u0026lt;br/\u0026gt;→ 恢复 Req 3, 重新 prefill\u0026#34;] F[\u0026#34;替代方案: 不做 preemption → OOM crash → 所有请求失败!\u0026#34;] A --\u0026gt; B --\u0026gt; C --\u0026gt; D --\u0026gt; E B -.-\u0026gt; F style A fill:#d4edda,stroke:#155724 style B fill:#f8d7da,stroke:#721c24 style C fill:#fff3cd,stroke:#856404 style D fill:#d4edda,stroke:#155724 style F fill:#f8d7da,stroke:#721c24 Preemption 的代价是被暂停的请求需要重新做 prefill（重建 KV Cache），增加了该请求的延迟。但比起 OOM 导致所有请求崩溃，这是一个合理的 trade-off。\n配套代码中的 preemption demo 显示，即使在非常受限的显存下（32 blocks），preemption 也能保证所有请求最终完成：\nConfig: 20 requests, 32 KV blocks, max_batch=4 Preemptions triggered: 12 Preemption allows the system to handle load spikes without OOM crashes. All 20 requests completed: True Chunked Prefill 长 Prompt 阻塞解码 Prefill 阶段（处理用户输入的 prompt）和 decode 阶段（逐 token 生成）有截然不同的计算特性：\nPrefill：一次性处理整个 prompt（可能数千个 token），是 Compute-bound 的大 GEMM 操作 Decode：每步只生成 1 个 token，是 Memory-bound 的 GEMV 操作 问题在于：如果一个请求的 prompt 有 8000 个 token，做 prefill 可能需要几百毫秒。在这段时间里，所有正在 decode 的请求都被阻塞了——它们在等 GPU 做完这个长 prefill。\nflowchart TD subgraph NOCHUNK[\u0026#34;无 Chunked Prefill\u0026#34;] GPU1[\u0026#34;GPU: ====== Prefill 8K tokens (200ms) ====== → D D D D\u0026#34;] RA1[\u0026#34;Req A: 等待 200ms... 然后才能 decode\u0026#34;] RB1[\u0026#34;Req B: 等待 200ms... 然后才能 decode\u0026#34;] end PROBLEM[\u0026#34;问题: Req A, B 的 TPOT 从 10ms 暴增到 200ms+\u0026lt;br/\u0026gt;用户体验: 正在打字的对话突然卡住 200ms\u0026#34;] style NOCHUNK fill:#fff3cd,stroke:#856404 style PROBLEM fill:#f8d7da,stroke:#721c24 对于正在和模型对话的用户来说，TPOT 从 10ms 跳到 200ms 会造成明显的\u0026quot;卡顿感\u0026quot;——字在流式输出时突然停了一下。\n分块 Prefill Chunked Prefill 的解决方案直截了当：把长 prompt 的 prefill 分成多个小块（chunk），每个 chunk 处理固定数量的 token（比如 512 个），每个 chunk 之间穿插一轮 decode step。\nflowchart TD subgraph CHUNK[\u0026#34;Chunked Prefill — 8192 tokens, chunk_size=512, 共 16 chunk\u0026#34;] GPU2[\u0026#34;GPU: P1 → D → P2 → D → P3 → D → P4 → D → P5 ...\u0026#34;] LEGEND[\u0026#34;P = prefill chunk (512 tokens, ~12ms) · D = decode step (~3ms)\u0026#34;] end subgraph DECODE[\u0026#34;Decode 请求视角\u0026#34;] RA2[\u0026#34;Req A: D — D — D — D — D (每 ~15ms 一次)\u0026#34;] RB2[\u0026#34;Req B: D — D — D — D — D (每 ~15ms 一次)\u0026#34;] end RESULT2[\u0026#34;TPOT: ~15ms（稳定!）vs 无 chunking 时的 200ms 尖峰\u0026#34;] CHUNK --\u0026gt; DECODE DECODE --\u0026gt; RESULT2 style CHUNK fill:#d4edda,stroke:#155724 style RESULT2 fill:#cce5ff,stroke:#004085 Trade-off：\nPrefill 总时间变长了：因为每个 chunk 之间插入了 decode step，8K token 的 prefill 从 200ms 变成了约 240ms（TTFT 略增） Decode 延迟变稳定了：TPOT 从 200ms 的尖峰降到稳定的 15ms 这是一个 TTFT vs TPOT 的 trade-off：\n指标 无 Chunked Prefill 有 Chunked Prefill TTFT (首 token) 200ms ~240ms (+20%) TPOT (逐 token，最坏情况) 200ms (被阻塞) ~15ms (稳定) TPOT (P99) 极差 可控 用户体验 偶尔严重卡顿 流畅稳定 在大多数在线服务场景下，稳定的 TPOT 比略低的 TTFT 更重要——用户更在意\u0026quot;打字流畅\u0026quot;而非\u0026quot;快 40ms 看到第一个字\u0026quot;。\nChunk 大小的选择 Chunk 大小的选择需要平衡：\n太小（如 64 tokens）：prefill 的 GPU 利用率低（小矩阵乘法效率差），TTFT 显著增加 太大（如 4096 tokens）：接近不分块的情况，decode 仍然会被长时间阻塞 典型值：512-1024 tokens，兼顾 prefill 效率和 decode 稳定性 SGLang 和 vLLM 都支持 chunked prefill，并可根据当前 running batch 的状态动态调整 chunk 大小——如果当前没有 decode 请求，就不需要分块，直接做完整的 prefill 以最小化 TTFT。\n推理系统全景 将以上所有组件整合起来，一个现代 LLM 推理引擎（如 SGLang）的整体架构如下：\nflowchart TD HTTP[\u0026#34;HTTP Server\u0026lt;br/\u0026gt;接收用户请求\u0026#34;] subgraph SCHED[\u0026#34;Scheduler（调度器 — 大脑）\u0026#34;] WQ[\u0026#34;Waiting Queue\u0026#34;] RT[\u0026#34;Radix Tree\u0026lt;br/\u0026gt;（Prefix Cache 索引）\u0026#34;] STEPS[\u0026#34;每个 iteration:\u0026lt;br/\u0026gt;1. 驱逐已完成请求\u0026lt;br/\u0026gt;2. 匹配新请求前缀\u0026lt;br/\u0026gt;3. Cache-aware 排序 + 纳入\u0026lt;br/\u0026gt;4. 必要时 preempt\u0026lt;br/\u0026gt;5. Chunked prefill 调度\u0026#34;] end subgraph BM[\u0026#34;Block Manager（块管理器 — 内存）\u0026#34;] POOL2[\u0026#34;Physical Block Pool: B0 B1 B2 B3 B4 B5 ...\u0026#34;] BMOPS[\u0026#34;按需分配/释放 · CoW · LRU 逐出\u0026#34;] end subgraph ME[\u0026#34;Model Executor（模型执行器 — 计算）\u0026#34;] ATTN[\u0026#34;Attention kernel (FlashAttention / FlashInfer)\u0026#34;] KV[\u0026#34;从 Block Table 读取 KV（非连续物理地址）\u0026#34;] PAR[\u0026#34;TP / PP 并行推理\u0026#34;] end HTTP --\u0026gt; SCHED --\u0026gt; BM --\u0026gt; ME style SCHED fill:#cce5ff,stroke:#004085 style BM fill:#fff3cd,stroke:#856404 style ME fill:#d4edda,stroke:#155724 调度器是大脑——它决定每一步执行哪些请求，是 prefill 还是 decode，是否需要 preempt。Block Manager 是内存管理器——它负责分配、释放、共享 KV Cache blocks。Model Executor 是执行引擎——它拿到调度器的决策和 block table 的映射，执行实际的 attention 计算。\n三者协同工作，构成了一个高效的推理流水线。\nKey Takeaways KV Cache 是推理的显存瓶颈。每个 token 需要存储 $2 \\times n_{\\text{layers}} \\times n_{\\text{heads}} \\times d_{\\text{head}} \\times \\text{dtype_size}$ 的 KV 数据。LLaMA-70B 下每个 token 约 2.5 MB，100 个并发请求、每个 2048 tokens 就需要 500 GB 的 KV Cache——远超单卡显存。能 batch 多少请求、能处理多长序列，完全取决于 KV Cache 的管理效率。\nPagedAttention 用 OS 虚拟内存思想解决显存碎片。固定大小的 block 按需分配，block table 做逻辑-物理映射，CoW 支持高效的序列分叉。将显存浪费从 60-80% 降到 \u0026lt;5%，同样的 GPU 可以 batch 2-4 倍的请求。\nRadixAttention 实现跨请求的 KV Cache 复用。Radix Tree 索引所有已缓存的 token 前缀，新请求自动匹配最长前缀并复用其 KV Cache。结合 LRU 逐出和 cache-aware scheduling，在共享前缀场景下（system prompt、few-shot、multi-turn chat）大幅降低 TTFT。\nContinuous Batching 将调度粒度从 batch 级降到 iteration 级。每步驱逐已完成请求、纳入等待请求，GPU 利用率从 20-50% 提升到 80-95%。配合 preemption 机制，优雅应对显存压力，避免 OOM 崩溃。\nChunked Prefill 解决长 prompt 阻塞 decode 的问题。将 prefill 分成小块，穿插 decode step，用略高的 TTFT 换取稳定的 TPOT。chunk 大小的选择需要根据负载特征权衡。\n这些技术不是互斥的，而是层层叠加的。现代推理引擎（SGLang、vLLM）同时使用 PagedAttention + Prefix Caching + Continuous Batching + Chunked Prefill，每一层解决不同维度的效率问题，叠加起来实现了比朴素实现 2-10 倍的吞吐量提升。\n配套代码 本文配套代码位于 code/03-inference-sglang/：\nkv_cache_from_scratch.py — 从零实现 KV Cache 和 PagedAttention Block Manager。包含 4 个部分：(1) 朴素 KV Cache 的浪费演示；(2) PagedAttention block 分配、释放、CoW 的完整实现；(3) 用 block-managed KV cache 跑一个 mini decode loop；(4) 朴素 vs PagedAttention 的内存效率对比。运行方式：python kv_cache_from_scratch.py\ncontinuous_batching.py — 静态 batching vs continuous batching 的对比模拟。包含 4 个部分：(1) 静态/连续 batching 的吞吐量对比；(2) ASCII 调度时间线可视化；(3) preemption 在显存压力下的表现；(4) 不同 batch size 下的 scaling 分析。运行方式：python continuous_batching.py\n所有代码支持 CPU 模式（无需 GPU），方便在任何环境下学习核心逻辑。\n参考资料 Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP'23 (vLLM) — PagedAttention 的原始论文，将 OS 虚拟内存思想引入 KV Cache 管理 SGLang: Efficient Execution of Structured Language Model Programs — Zheng et al., 2024 (LMSYS) — RadixAttention 和 cache-aware scheduling 的原始论文 Orca: A Distributed Serving System for Transformer-Based Generative Models — Yu et al., OSDI'22 — Continuous Batching（iteration-level scheduling）的开创性工作 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al., NeurIPS'22 — 推理引擎中 attention kernel 的核心优化 vLLM Documentation — docs.vllm.ai — vLLM 官方文档 SGLang Documentation — sgl-project.github.io — SGLang 官方文档 Sarathi-Serve: Chunked-Prefills for Efficient LLM Serving — Agrawal et al., 2024 — Chunked Prefill 的系统化分析 ","permalink":"https://mzf666.github.io/llm-infra/zh/posts/03-inference-sglang/","summary":"深入 PagedAttention 与 RadixAttention，理解现代 LLM 推理引擎的核心设计。","title":"LLM 推理系统架构（以 SGLang 为例）"},{"content":"Motivation 如果你去问做 RLHF 的工程师\u0026quot;RLHF 最难的地方是什么\u0026quot;，答案很可能不是 PPO 算法本身——PPO 在强化学习领域已经是相当成熟的算法。真正的困难在于：RLHF 需要同时驱动四个大模型，在它们之间编排复杂的数据流，并且在训练循环内部嵌入了一个完整的推理系统。\n让我们先算一笔账感受一下。假设你要对一个 7B 模型做 SFT，FP16 下参数占 14 GB，加上 Adam 优化器状态（momentum + variance 各一份）和梯度，总计大约 $14 \\times 4 = 56$ GB——一张 A100-80GB 刚好能装下。\n但如果要做 RLHF，你需要四个模型：\nActor（被训练的 LLM）：14 GB 参数 + 42 GB 优化器/梯度 = 56 GB Critic（价值函数）：14 GB 参数 + 42 GB 优化器/梯度 = 56 GB Reward Model（奖励模型，冻结）：14 GB Reference Model（参考模型，冻结）：14 GB 总计：140 GB，至少需要 2 张 A100-80GB，而这还没算激活值和 KV Cache。\n除了显存，RLHF 训练循环内部还嵌套了一个推理过程——Actor 需要自回归地生成 response。这个生成过程是 memory-bound 的推理问题，但它发生在 compute-bound 的训练循环内部。你需要同时优化这两种截然不同的计算模式。\n本文将从系统视角出发，帮你理解 RLHF 训练的全貌：四模型的角色与交互、PPO 的完整数据流、为什么这是一个系统问题，以及 verl 框架如何用\u0026quot;混合引擎\u0026quot;优雅地解决这些挑战。\n前置知识 GPU 显存模型与分布式通信基础（第 1 篇）——理解显存瓶颈和通信原语 分布式并行策略全景（第 2 篇）——特别是 FSDP 和 Tensor Parallelism LLM 推理系统架构（第 3 篇）——理解自回归生成和 KV Cache 强化学习基本概念（Policy、Reward、Value Function），不要求精通 先看一张全局架构图，后面逐一展开：\nflowchart TD subgraph RLHF[\u0026#34;RLHF 训练系统全景\u0026#34;] A[\u0026#34;**Actor (Policy)**\u0026lt;br/\u0026gt;生成 response\u0026lt;br/\u0026gt;PPO 更新\u0026lt;br/\u0026gt;✓ 可训练\u0026#34;] R[\u0026#34;**Reference (Frozen)**\u0026lt;br/\u0026gt;KL 锚点\u0026lt;br/\u0026gt;防止漂移\u0026lt;br/\u0026gt;✗ 冻结\u0026#34;] RM[\u0026#34;**Reward Model (Frozen)**\u0026lt;br/\u0026gt;打分\u0026lt;br/\u0026gt;标量奖励\u0026lt;br/\u0026gt;✗ 冻结\u0026#34;] C[\u0026#34;**Critic (Value)**\u0026lt;br/\u0026gt;估计价值\u0026lt;br/\u0026gt;计算优势\u0026lt;br/\u0026gt;✓ 可训练\u0026#34;] A --\u0026gt; Flow R --\u0026gt; Flow RM --\u0026gt; Flow C --\u0026gt; Flow Flow[\u0026#34;PPO 数据流：生成 → 打分 → 优势估计 → 更新\u0026#34;] end subgraph Challenges[\u0026#34;系统挑战\u0026#34;] CH1[\u0026#34;4x 显存 vs SFT\u0026#34;] CH2[\u0026#34;推理嵌套在训练中\u0026#34;] CH3[\u0026#34;复杂数据依赖\u0026#34;] CH4[\u0026#34;权重同步\u0026#34;] CH5[\u0026#34;异构计算模式\u0026#34;] CH6[\u0026#34;模型放置策略\u0026#34;] end RLHF --\u0026gt; Challenges style A fill:#d4edda,stroke:#28a745 style C fill:#d4edda,stroke:#28a745 style R fill:#fff3cd,stroke:#ffc107 style RM fill:#fff3cd,stroke:#ffc107 style Flow fill:#cce5ff,stroke:#007bff RLHF 四模型架构 RLHF 之所以是系统问题，根源在于它需要四个模型协同工作。理解每个模型的角色是理解整个系统的前提。\nActor（策略模型） Actor 就是我们要训练的 LLM——它的任务是根据 prompt 生成高质量的 response。在 RLHF 之前，它通常已经经过 SFT（Supervised Fine-Tuning），具备基本的指令遵循能力。\n从架构上看，Actor 就是一个标准的 Causal Language Model：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 class CausalLM(nn.Module): \u0026#34;\u0026#34;\u0026#34; Causal Language Model (GPT-style). 生产环境中就是 LLaMA、Qwen 等 7B-70B+ 的模型。 \u0026#34;\u0026#34;\u0026#34; def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=4, d_ff=512, max_seq_len=256): super().__init__() self.token_embed = nn.Embedding(vocab_size, d_model) self.pos_embed = nn.Embedding(max_seq_len, d_model) self.blocks = nn.ModuleList([ TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers) ]) self.ln_f = nn.LayerNorm(d_model) self.lm_head = nn.Linear(d_model, vocab_size, bias=False) self.lm_head.weight = self.token_embed.weight # Weight tying Actor 在 RLHF 中有两种工作模式，这正是系统设计的难点所在：\n生成模式（推理）：自回归地采样 token，memory-bound，受益于 KV Cache 和 TP 训练模式：基于 PPO 目标函数计算梯度并更新参数，compute-bound，受益于 FSDP Reference Model（参考模型） Reference Model 是 Actor 在 RLHF 训练开始前的一份冻结副本。它的唯一作用是计算 KL 散度惩罚——防止 Actor 在追逐高奖励的过程中偏离原始行为太远。\n1 2 3 4 5 6 # Reference = Actor 的冻结副本（训练开始前复制一次，此后永不更新） reference = CausalLM(vocab_size, d_model, n_heads, n_layers, d_ff, max_seq_len) reference.load_state_dict(actor.state_dict()) for param in reference.parameters(): param.requires_grad = False # 完全冻结 为什么需要 Reference？ 这涉及 RLHF 中一个著名的问题——Reward Hacking。Reward Model 并不完美，它是一个学出来的近似函数。如果没有任何约束，Actor 很容易找到 Reward Model 的\u0026quot;漏洞\u0026quot;：生成的 response 得到高分，但实际上是无意义的、重复的、或者过度冗长的文本。KL 惩罚通过约束 Actor 不能偏离 Reference 太远，有效缓解了这个问题。\n系统开销：Reference 虽然冻结不需要优化器状态，但它仍然需要完整的前向传播来计算 log probabilities。对于 7B 模型，这意味着额外 14 GB 显存和一次完整的前向计算。\nReward Model（奖励模型） Reward Model 是在人类偏好数据上预训练好的——给定 (prompt, response_A, response_B)，它学会为人类偏好的那个 response 打更高的分。在 RLHF 训练期间，它作为\u0026quot;评委\u0026quot;对 Actor 生成的 response 打分。\n架构上，它和 Actor 共享相同的 Transformer backbone，但最后一层的 language modeling head 被替换为一个 value head，输出标量奖励：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 class RewardModel(nn.Module): \u0026#34;\u0026#34;\u0026#34; Reward Model：将 (prompt, response) 映射到标量奖励值。 使用最后一个 token 的 hidden state 作为序列表示。 \u0026#34;\u0026#34;\u0026#34; def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=4, d_ff=512, max_seq_len=256): super().__init__() # ... 与 CausalLM 相同的 Transformer backbone ... # 关键区别：value head 替代 lm_head self.value_head = nn.Linear(d_model, 1, bias=False) def forward(self, input_ids): # ... Transformer 前向传播 ... last_hidden = x[:, -1, :] # 最后一个 token 的 hidden state reward = self.value_head(last_hidden).squeeze(-1) # 标量 return reward # (batch,) Reward Model 在 RLHF 训练中完全冻结，只做推理。但不要小看它的系统开销——每个 PPO iteration 都需要对一整个 batch 的 (prompt + response) 做前向传播。\nCritic（价值模型） Critic 估计每个 token 位置的期望未来奖励 $V(s_t)$。这个价值估计用于计算 GAE（Generalized Advantage Estimation），从而降低策略梯度的方差。\n没有 Critic，PPO 退化为 REINFORCE——虽然理论上也能工作，但方差极高，训练极不稳定。\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 class CriticModel(nn.Module): \u0026#34;\u0026#34;\u0026#34; Critic：估计每个 token 位置的 V(s) 值。 与 Reward Model 不同，它输出 per-token 的标量值，而非整个序列的标量。 \u0026#34;\u0026#34;\u0026#34; def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=4, d_ff=512, max_seq_len=256): super().__init__() # ... 与 CausalLM 相同的 Transformer backbone ... self.value_head = nn.Linear(d_model, 1, bias=False) def forward(self, input_ids): # ... Transformer 前向传播 ... values = self.value_head(x).squeeze(-1) # (batch, seq_len) return values # 每个位置一个标量值 Critic 和 Actor 一样是可训练的——它与 Actor 同步更新，通常用 MSE 损失拟合 GAE 计算出的 returns。在实践中，Critic 常用 Reward Model 的权重来初始化（因为两者的目标相似：估计\u0026quot;response 有多好\u0026quot;）。\n四模型总览 把四个模型放在一起比较：\nflowchart TD subgraph Compare[\u0026#34;RLHF 四模型对比\u0026#34;] direction LR subgraph Trainable[\u0026#34;✓ 可训练\u0026#34;] Actor[\u0026#34;**Actor**\u0026lt;br/\u0026gt;Adam 优化器\u0026lt;br/\u0026gt;前向 + 反向\u0026lt;br/\u0026gt;56 GB (7B FP16)\u0026#34;] Critic[\u0026#34;**Critic**\u0026lt;br/\u0026gt;Adam 优化器\u0026lt;br/\u0026gt;前向 + 反向\u0026lt;br/\u0026gt;56 GB (7B FP16)\u0026#34;] end subgraph Frozen[\u0026#34;✗ 冻结\u0026#34;] Reward[\u0026#34;**Reward Model**\u0026lt;br/\u0026gt;无优化器\u0026lt;br/\u0026gt;仅前向\u0026lt;br/\u0026gt;14 GB (7B FP16)\u0026#34;] Reference[\u0026#34;**Reference**\u0026lt;br/\u0026gt;无优化器\u0026lt;br/\u0026gt;仅前向\u0026lt;br/\u0026gt;14 GB (7B FP16)\u0026#34;] end end Compare --\u0026gt; Total[\u0026#34;**TOTAL: 140 GB**\u0026lt;br/\u0026gt;对比：SFT 只需 1 个模型 + Adam ≈ 56 GB\u0026lt;br/\u0026gt;RLHF 需要 2.5x SFT 的显存\u0026#34;] style Actor fill:#d4edda,stroke:#28a745 style Critic fill:#d4edda,stroke:#28a745 style Reward fill:#fff3cd,stroke:#ffc107 style Reference fill:#fff3cd,stroke:#ffc107 style Total fill:#cce5ff,stroke:#007bff 更大规模模型的显存需求更加惊人：\n模型规模 单模型 (FP16) 四模型权重 含优化器状态 最少 GPU 数 (A100-80GB) 7B 14 GB 56 GB 140 GB 2 13B 26 GB 104 GB 260 GB 4 70B 140 GB 560 GB 1400 GB 18 PPO 数据流 理解了四个模型各自的角色之后，关键问题是：它们之间的数据是如何流动的？ PPO 的每个 iteration 包含两个阶段：Rollout Phase（经验收集）和 Training Phase（梯度更新）。\nPhase 1: Rollout（经验收集） Rollout 阶段的目标是收集一批\u0026quot;经验\u0026quot;——Actor 生成的 response、Reward Model 的评分、以及 Critic 的价值估计。整个过程的核心数据流如下：\nflowchart TD subgraph Rollout[\u0026#34;ROLLOUT PHASE (推理，无梯度)\u0026#34;] S1[\u0026#34;**Step 1: 生成**\u0026lt;br/\u0026gt;Prompts → Actor → 自回归采样 → Responses\u0026#34;] S1 --\u0026gt; S2a \u0026amp; S2b \u0026amp; S3 S2a[\u0026#34;**Step 2a: Reward 打分**\u0026lt;br/\u0026gt;Reward Model → 标量 Rewards\u0026#34;] S2b[\u0026#34;**Step 2b: KL 计算**\u0026lt;br/\u0026gt;Reference Model → log_probs → KL Penalties\u0026#34;] S3[\u0026#34;**Step 3: 价值估计**\u0026lt;br/\u0026gt;Critic → per-token Values\u0026#34;] S2a --\u0026gt; S4 S2b --\u0026gt; S4 S3 --\u0026gt; S4 S4[\u0026#34;**Step 4: 优势计算**\u0026lt;br/\u0026gt;GAE(Rewards, KL, Values) → Advantages\u0026#34;] end style S1 fill:#cce5ff,stroke:#007bff style S2a fill:#fff3cd,stroke:#ffc107 style S2b fill:#fff3cd,stroke:#ffc107 style S3 fill:#fff3cd,stroke:#ffc107 style S4 fill:#d4edda,stroke:#28a745 让我们逐步拆解。\nStep 1: 自回归生成 Actor 拿到一批 prompt，自回归地生成 response。这本质上是一个推理问题：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 @torch.no_grad() def generate_responses(actor, prompt_ids, max_new_tokens, temperature=1.0): \u0026#34;\u0026#34;\u0026#34; Actor 自回归生成 response。 生产环境中会使用 KV Cache、Continuous Batching、Tensor Parallelism。 \u0026#34;\u0026#34;\u0026#34; generated = prompt_ids.clone() for _ in range(max_new_tokens): logits = actor(generated) # 完整前向传播 next_logits = logits[:, -1, :] # 只取最后一个位置 probs = F.softmax(next_logits / temperature, dim=-1) next_token = torch.multinomial(probs, num_samples=1) generated = torch.cat([generated, next_token], dim=1) return generated # (batch, prompt_len + response_len) 系统洞察：这个循环的每一步都需要一次完整的前向传播（没有 KV Cache 的情况下）。生成 $T$ 个 token 就需要 $T$ 次前向传播。这是 RLHF 中计算开销最大的单个环节。在生产系统（如 verl）中，这里会使用 KV Cache + Tensor Parallelism + Continuous Batching 来加速——本质上需要在训练循环内部嵌入一个推理引擎。\nStep 2: Reward 打分与 KL 计算 生成 response 后，三个模型需要分别处理这些 response：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 def compute_rewards_and_kl(full_ids, prompt_len, actor, reference, reward_model, kl_coeff): # 1. Reward Model 打分（冻结，无梯度） with torch.no_grad(): rewards = reward_model(full_ids) # (batch,) 标量奖励 # 2. Actor 计算 log probabilities actor_log_probs = actor.get_log_probs(full_ids) response_actor_lp = actor_log_probs[:, prompt_len - 1:] # 3. Reference 计算 log probabilities（冻结，无梯度） with torch.no_grad(): ref_log_probs = reference.get_log_probs(full_ids) response_ref_lp = ref_log_probs[:, prompt_len - 1:] # 4. KL 惩罚：KL(Actor || Reference) ≈ actor_lp - ref_lp kl_penalties = kl_coeff * (response_actor_lp.detach() - response_ref_lp) return rewards, response_actor_lp, response_ref_lp, kl_penalties KL 惩罚是 RLHF 稳定性的关键。在实践中，KL 散度可以用 per-token 的近似来计算：\n$$D_{KL}(\\pi_\\theta | \\pi_\\text{ref}) \\approx \\sum_t \\left[\\log \\pi_\\theta(a_t | s_t) - \\log \\pi_\\text{ref}(a_t | s_t)\\right]$$\n当 KL 为正值时，说明 Actor 在该 token 上的概率比 Reference 更高——即 Actor 正在偏离原始行为。kl_coeff 控制惩罚力度：\n太大：Actor 几乎学不到东西（被\u0026quot;锁死\u0026quot;在 Reference 附近） 太小：容易出现 Reward Hacking（Actor 找到 Reward Model 的漏洞） 常见取值：0.01 - 0.2，有些系统（如 InstructGPT）会自适应调整 Step 3: GAE 优势估计 有了 rewards、KL penalties 和 Critic 的 value estimates 之后，就可以计算 GAE（Generalized Advantage Estimation）了。优势函数 $A(s_t, a_t)$ 告诉我们：\u0026ldquo;这个 action 比期望水平好多少？\u0026rdquo;\n$$A^{GAE(\\gamma,\\lambda)}t = \\sum{l=0}^{T-t} (\\gamma\\lambda)^l \\delta_{t+l}$$\n其中 TD error $\\delta_t = r_t + \\gamma V(s_{t+1}) - V(s_t)$。\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 @torch.no_grad() def compute_advantages_gae(rewards, kl_penalties, values, gamma=1.0, lam=0.95): \u0026#34;\u0026#34;\u0026#34; GAE 通过 lambda 参数平衡偏差与方差： - lam=1: 高方差，低偏差（Monte Carlo） - lam=0: 低方差，高偏差（1-step TD） - lam=0.95: 常用的折衷选择 \u0026#34;\u0026#34;\u0026#34; B, T = values.shape # 构造 per-token rewards： # - 每个 token 承受 KL 惩罚 # - 只有最后一个 token 获得 Reward Model 的序列级奖励 per_token_rewards = -kl_penalties.clone() per_token_rewards[:, -1] += rewards # 序列奖励分配到最后一个 token # 从后往前计算 GAE advantages = torch.zeros_like(values) last_gae = torch.zeros(B, device=values.device) for t in reversed(range(T)): next_value = values[:, t + 1] if t \u0026lt; T - 1 else torch.zeros(B) delta = per_token_rewards[:, t] + gamma * next_value - values[:, t] last_gae = delta + gamma * lam * last_gae advantages[:, t] = last_gae returns = advantages + values # Critic 的训练目标 return advantages, returns 注意这里的一个设计选择：序列级的 reward 被分配到最后一个 token，而 KL 惩罚是 per-token 的。这是 RLHF 中的标准做法。\nPhase 2: Training（PPO 更新） 收集完经验之后，进入 PPO 更新阶段。PPO 的核心思想是：用同一批经验数据做多次梯度更新，但通过 clipping 防止策略变化过大。\nflowchart TD subgraph Training[\u0026#34;TRAINING PHASE (梯度更新)\u0026#34;] Input[\u0026#34;**输入：Rollout 经验数据**\u0026lt;br/\u0026gt;full_ids, old_log_probs, advantages, returns\u0026#34;] Input --\u0026gt; Loop subgraph Loop[\u0026#34;for epoch in range(ppo_epochs) — 通常 2-4 个 epoch\u0026#34;] direction LR ActorUpdate[\u0026#34;**Actor PPO 更新**\u0026lt;br/\u0026gt;new_lp = Actor(full_ids)\u0026lt;br/\u0026gt;ratio = exp(new - old)\u0026lt;br/\u0026gt;loss = -min(r*A, clip)\u0026lt;br/\u0026gt;backward + step\u0026#34;] CriticUpdate[\u0026#34;**Critic 更新**\u0026lt;br/\u0026gt;values = Critic(full_ids)\u0026lt;br/\u0026gt;loss = MSE(values, returns)\u0026lt;br/\u0026gt;backward + step\u0026#34;] end Note[\u0026#34;Reference 和 Reward Model：本阶段不参与\u0026#34;] end style ActorUpdate fill:#d4edda,stroke:#28a745 style CriticUpdate fill:#d4edda,stroke:#28a745 style Input fill:#cce5ff,stroke:#007bff style Note fill:#fff3cd,stroke:#ffc107 PPO Clipped Surrogate Loss PPO 的核心是 clipped surrogate objective，它是 PPO 相比 vanilla policy gradient 的关键创新：\n$$L^{CLIP}(\\theta) = \\mathbb{E}\\left[\\min\\left(r_t(\\theta)\\hat{A}_t,; \\text{clip}(r_t(\\theta),; 1-\\epsilon,; 1+\\epsilon)\\hat{A}_t\\right)\\right]$$\n其中 policy ratio $r_t(\\theta) = \\frac{\\pi_\\theta(a_t | s_t)}{\\pi_{\\theta_{old}}(a_t | s_t)} = \\exp(\\log\\pi_\\theta - \\log\\pi_{\\theta_{old}})$。\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 def ppo_actor_loss(actor, full_ids, prompt_len, old_log_probs, advantages, clip_eps): \u0026#34;\u0026#34;\u0026#34; PPO clipped surrogate loss。 ratio \u0026gt; 1: action 现在更可能被选中 ratio \u0026lt; 1: action 现在更不可能被选中 clipping 防止 ratio 偏离 [1-eps, 1+eps] 区间 \u0026#34;\u0026#34;\u0026#34; new_log_probs = actor.get_log_probs(full_ids)[:, prompt_len - 1:] # Policy ratio ratio = torch.exp(new_log_probs - old_log_probs.detach()) # 标准化 advantages（训练稳定性的标准做法） adv = (advantages - advantages.mean()) / (advantages.std() + 1e-8) # Clipped surrogate loss surr1 = ratio * adv surr2 = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * adv loss = -torch.min(surr1, surr2).mean() return loss, ratio Clipping 的直觉是：\n当 $\\hat{A}_t \u0026gt; 0$（好的 action）：我们想增大其概率，但 clipping 防止 $r_t$ 超过 $1+\\epsilon$ 当 $\\hat{A}_t \u0026lt; 0$（差的 action）：我们想减小其概率，但 clipping 防止 $r_t$ 低于 $1-\\epsilon$ 这确保了每次更新的\u0026quot;步长\u0026quot;有界，让 PPO 能安全地在同一批数据上做多次更新。\nCritic Loss Critic 用 MSE 损失拟合 GAE 计算出的 returns：\n$$L^{VF} = \\frac{1}{2}\\mathbb{E}\\left[(V_\\phi(s_t) - R_t)^2\\right]$$\n1 2 3 4 5 6 def ppo_critic_loss(critic, full_ids, prompt_len, returns): \u0026#34;\u0026#34;\u0026#34;Critic 价值函数损失：MSE(predicted_value, returns)\u0026#34;\u0026#34;\u0026#34; values = critic(full_ids)[:, prompt_len - 1:-1] min_len = min(values.shape[1], returns.shape[1]) loss = F.mse_loss(values[:, :min_len], returns[:, :min_len].detach()) return loss 完整的 PPO 训练循环 把 Rollout 和 Training 两个阶段拼在一起，就是完整的 PPO 训练循环：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 # 只有 Actor 和 Critic 需要优化器 actor_optim = torch.optim.Adam(actor.parameters(), lr=1e-5) critic_optim = torch.optim.Adam(critic.parameters(), lr=1e-5) for iteration in range(num_iterations): # ===== Phase 1: Rollout (推理模式) ===== with torch.no_grad(): full_ids = generate_responses(actor, prompt_ids, response_len, temperature) rewards = reward_model(full_ids) old_log_probs = actor.get_log_probs(full_ids)[:, prompt_len-1:] ref_log_probs = reference.get_log_probs(full_ids)[:, prompt_len-1:] kl_penalties = kl_coeff * (old_log_probs - ref_log_probs) values = critic(full_ids)[:, prompt_len-1:-1] advantages, returns = compute_advantages_gae( rewards, kl_penalties, values, gamma, lam) # ===== Phase 2: PPO Update (训练模式) ===== for epoch in range(ppo_epochs): # 同一批数据做多次更新 # Actor 更新 actor_optim.zero_grad() a_loss, ratio = ppo_actor_loss(actor, full_ids, prompt_len, old_log_probs, advantages, clip_eps) a_loss.backward() torch.nn.utils.clip_grad_norm_(actor.parameters(), max_norm=1.0) actor_optim.step() # Critic 更新 critic_optim.zero_grad() c_loss = ppo_critic_loss(critic, full_ids, prompt_len, returns) c_loss.backward() torch.nn.utils.clip_grad_norm_(critic.parameters(), max_norm=1.0) critic_optim.step() 注意几个关键细节：\n学习率极低（1e-5）：RLHF 中 Actor 的学习率通常比 SFT 低 10 倍以上，因为我们只想做微小调整 梯度裁剪（max_norm=1.0）：防止梯度爆炸，RLHF 训练中几乎是必须的 多个 PPO epoch：同一批 rollout 数据做 2-4 次更新，clipping 确保不会过度更新 为什么 RLHF 是系统问题 到这里，你应该已经感受到 RLHF 的复杂性了。让我们系统性地总结它带来的四大系统挑战。\n挑战 1：显存压力——4 倍于 SFT 最直观的挑战是显存。我们在前面已经算过：7B 模型的 RLHF 需要 ~140 GB，而 SFT 只需要 ~56 GB。\nflowchart TD subgraph ACTOR[\u0026#34;Actor — 56 GB\u0026#34;] direction TB A1[\u0026#34;参数: 14 GB\u0026#34;] A2[\u0026#34;梯度: 14 GB\u0026#34;] A3[\u0026#34;Adam momentum: 14 GB\u0026#34;] A4[\u0026#34;Adam variance: 14 GB\u0026#34;] end subgraph CRITIC[\u0026#34;Critic — 56 GB\u0026#34;] direction TB C1[\u0026#34;参数: 14 GB\u0026#34;] C2[\u0026#34;梯度: 14 GB\u0026#34;] C3[\u0026#34;Adam momentum: 14 GB\u0026#34;] C4[\u0026#34;Adam variance: 14 GB\u0026#34;] end subgraph REWARD[\u0026#34;Reward Model — 14 GB\u0026#34;] R1[\u0026#34;参数: 14 GB (冻结，仅推理)\u0026#34;] end subgraph REF[\u0026#34;Reference — 14 GB\u0026#34;] RF1[\u0026#34;参数: 14 GB (冻结，仅推理)\u0026#34;] end ACTOR \u0026amp; CRITIC \u0026amp; REWARD \u0026amp; REF --\u0026gt; TOTAL[\u0026#34;**总计: 140 GB**\\n(SFT: ~56 GB)\\n+ 激活值 + KV Cache + 通信 buffer\u0026#34;] style ACTOR fill:#d4edda,stroke:#28a745 style CRITIC fill:#d4edda,stroke:#28a745 style REWARD fill:#fff3cd,stroke:#ffc107 style REF fill:#fff3cd,stroke:#ffc107 style TOTAL fill:#cce5ff,stroke:#007bff 在实际生产中，还需要考虑激活值（可以用 activation checkpointing 压缩）和通信 buffer。70B 模型的 RLHF 需要 1400 GB 显存，即使 18 张 A100-80GB 也仅够放下模型参数和优化器。\n挑战 2：计算异构——推理嵌套在训练中 RLHF 最独特的系统挑战是：训练循环内部嵌入了一个完整的推理过程。\n在 SFT 中，整个训练循环都是 compute-bound 的：前向传播 → 反向传播 → 参数更新，计算模式统一。但 RLHF 的 Rollout Phase 需要自回归生成——这是一个典型的 memory-bound 推理任务。\nflowchart LR subgraph SFT[\u0026#34;SFT 训练循环 (全部 compute-bound)\u0026#34;] direction LR SF[\u0026#34;前向\u0026#34;] --\u0026gt; SB[\u0026#34;反向\u0026#34;] --\u0026gt; SU[\u0026#34;更新\u0026#34;] SStrategy[\u0026#34;最佳策略: FSDP / DDP\u0026#34;] end subgraph RLHF[\u0026#34;RLHF 训练循环 (异构计算)\u0026#34;] direction LR RG[\u0026#34;自回归生成\u0026lt;br/\u0026gt;⚡ memory-bound\u0026lt;br/\u0026gt;最佳: TP + KV Cache\u0026#34;] --\u0026gt; RS[\u0026#34;打分\u0026#34;] --\u0026gt; RGAE[\u0026#34;GAE\u0026#34;] --\u0026gt; RPPO[\u0026#34;PPO 更新\u0026lt;br/\u0026gt;⚡ compute-bound\u0026lt;br/\u0026gt;最佳: FSDP\u0026#34;] end style RG fill:#fff3cd,stroke:#ffc107 style RPPO fill:#d4edda,stroke:#28a745 这意味着你需要在同一组 GPU 上切换两种截然不同的并行策略：\n阶段 计算特性 最佳并行策略 瓶颈 自回归生成 Memory-bound Tensor Parallelism HBM 带宽 Reward 打分 Compute-bound Data Parallelism 算力 PPO 更新 Compute-bound FSDP (ZeRO-3) 算力 + 显存 挑战 3：数据流复杂——四模型的严格依赖顺序 RLHF 的数据流不是简单的\u0026quot;输入 → 输出\u0026quot;，而是四个模型之间的有向无环图（DAG）。每个 PPO iteration 的执行顺序是严格确定的：\nflowchart LR P[\u0026#34;Prompts\u0026#34;] --\u0026gt; Actor Actor --\u0026gt; resp[\u0026#34;responses\u0026#34;] resp --\u0026gt; RM[\u0026#34;Reward Model\u0026#34;] resp --\u0026gt; Ref[\u0026#34;Reference\u0026#34;] resp --\u0026gt; Crit[\u0026#34;Critic\u0026#34;] RM --\u0026gt; rewards Ref --\u0026gt; KL[\u0026#34;KL penalties\u0026#34;] Crit --\u0026gt; values rewards --\u0026gt; GAE KL --\u0026gt; GAE values --\u0026gt; GAE GAE --\u0026gt; advantages advantages --\u0026gt; ActorPPO[\u0026#34;Actor PPO update\u0026lt;br/\u0026gt;← old_log_probs\u0026#34;] advantages --\u0026gt; CriticMSE[\u0026#34;Critic MSE update\u0026lt;br/\u0026gt;← returns\u0026#34;] style Actor fill:#d4edda,stroke:#28a745 style GAE fill:#cce5ff,stroke:#007bff style ActorPPO fill:#d4edda,stroke:#28a745 style CriticMSE fill:#d4edda,stroke:#28a745 style RM fill:#fff3cd,stroke:#ffc107 style Ref fill:#fff3cd,stroke:#ffc107 style Crit fill:#fff3cd,stroke:#ffc107 这个 DAG 的关键约束：\n生成必须先完成：Reward、Reference、Critic 都依赖 Actor 生成的 response GAE 依赖三方输入：rewards + KL penalties + values 必须全部就绪才能计算 PPO 更新依赖 GAE：advantages 和 returns 是 Actor 和 Critic 更新的输入 Actor 更新后需要权重同步：下一次生成要用更新后的权重 在分布式场景下，如果四个模型放在不同的 GPU 组上，这些依赖关系就变成了跨设备的数据传输，调度复杂度大幅上升。\n挑战 4：权重同步 每个 PPO iteration 结束后，Actor 的权重被更新了。但下一个 iteration 的生成阶段需要使用更新后的权重进行推理。如果训练和推理使用不同的并行策略（比如训练用 FSDP，推理用 TP），那就需要在两种权重格式之间进行权重重整（weight resharding）：\nflowchart TD A[\u0026#34;PPO 更新结束\u0026lt;br/\u0026gt;(FSDP 格式: 每个 GPU 持有 1/N 参数分片)\u0026#34;] A --\u0026gt; B[\u0026#34;权重重整 (Resharding)\u0026#34;] B --\u0026gt; C[\u0026#34;下一轮生成开始\u0026lt;br/\u0026gt;(TP 格式: 每个 GPU 持有所有层的一个切片)\u0026#34;] style A fill:#d4edda,stroke:#28a745 style B fill:#fff3cd,stroke:#ffc107 style C fill:#cce5ff,stroke:#007bff 这个重整过程需要 all-gather 通信来收集所有分片，然后按 TP 的切分方式重新分发。对于大模型来说，这是一笔不小的通信开销。\n计算开销对比 把 RLHF 和 SFT 的计算开销放在一起比较：\n操作 SFT RLHF 额外倍数 Actor 前向传播 1x 2x 2x Actor 反向传播 1x 1x 1x 自回归生成 0 $N$x $+N$x Reference 前向传播 0 1x +1x Reward Model 前向传播 0 1x +1x Critic 前向传播 0 2x +2x Critic 反向传播 0 1x +1x GAE 计算 0 1x +1x 总计（近似） ~2x ~10-16x 5-8x 其中 $N$ 是生成的 token 数。如果生成 256 个 token，仅生成阶段就需要 256 次前向传播（无 KV Cache 时）。RLHF 单步训练的计算量大约是 SFT 的 5-8 倍。\nverl 架构深入分析 面对上述挑战，业界提出了多种解决方案。verl（Volcano Engine Reinforcement Learning）是字节跳动开源的 RLHF 训练框架，它通过**混合引擎（Hybrid Engine）和共置策略（Colocated Strategy）**优雅地解决了这些问题。\n核心设计理念 verl 的核心观察是：RLHF 的两个阶段需要不同的并行策略，但传统方案要么用 Separated 策略（浪费 GPU），要么用 Colocated 策略（显存不够）。verl 的解法是：同一组 GPU 上动态切换并行模式。\nflowchart LR subgraph HybridEngine[\u0026#34;verl 混合引擎 — 同一组 GPU 动态切换并行策略\u0026#34;] subgraph Gen[\u0026#34;GENERATION 阶段\u0026#34;] G1[\u0026#34;Tensor Parallel\u0026lt;br/\u0026gt;(低延迟推理)\u0026#34;] G2[\u0026#34;每个 GPU 持有\u0026lt;br/\u0026gt;所有层的一个切片\u0026#34;] G3[\u0026#34;all-reduce 通信\u0026#34;] end subgraph Train[\u0026#34;TRAINING 阶段\u0026#34;] T1[\u0026#34;FSDP (ZeRO-3)\u0026lt;br/\u0026gt;(高效训练)\u0026#34;] T2[\u0026#34;每个 GPU 持有\u0026lt;br/\u0026gt;1/N 的参数分片\u0026#34;] T3[\u0026#34;all-gather +\u0026lt;br/\u0026gt;reduce-scatter\u0026#34;] end Gen \u0026lt;-- \u0026#34;weight\u0026lt;br/\u0026gt;resharding\u0026#34; --\u0026gt; Train end Adv[\u0026#34;**关键优势**\u0026lt;br/\u0026gt;生成用 TP: 低延迟\u0026lt;br/\u0026gt;训练用 FSDP: 显存高效\u0026lt;br/\u0026gt;无 GPU 空闲浪费\u0026#34;] HybridEngine --\u0026gt; Adv style Gen fill:#cce5ff,stroke:#007bff style Train fill:#d4edda,stroke:#28a745 style Adv fill:#fff3cd,stroke:#ffc107 共置策略 vs 分离策略 在分布式 RLHF 中，四个模型的放置策略是一个核心决策。verl 支持两种策略：\n策略 1: Colocated（共置） 所有四个模型放在同一组 GPU 上。\nflowchart TD subgraph Colocated[\u0026#34;Colocated — GPU 0-7: 所有模型共置\u0026#34;] Models[\u0026#34;Actor(FSDP) + Critic(FSDP) + Reward(TP) + Reference(TP)\u0026#34;] Pro1[\u0026#34;(+) 无跨组数据传输\u0026#34;] Pro2[\u0026#34;(+) 调度简单：按阶段顺序执行\u0026#34;] Pro3[\u0026#34;(+) GPU 利用率高\u0026#34;] Con1[\u0026#34;(-) 显存压力大：4 模型共享\u0026#34;] Con2[\u0026#34;(-) 无法独立优化并行策略\u0026#34;] Con3[\u0026#34;(-) 需要精细的显存管理\u0026#34;] end style Models fill:#cce5ff,stroke:#007bff style Pro1 fill:#d4edda,stroke:#28a745 style Pro2 fill:#d4edda,stroke:#28a745 style Pro3 fill:#d4edda,stroke:#28a745 style Con1 fill:#fff3cd,stroke:#ffc107 style Con2 fill:#fff3cd,stroke:#ffc107 style Con3 fill:#fff3cd,stroke:#ffc107 策略 2: Separated（分离） 每个模型放在独立的 GPU 组上。\nflowchart TD subgraph Separated[\u0026#34;Separated — 每个模型独立 GPU 组\u0026#34;] direction LR subgraph AG[\u0026#34;Actor GPUs 0-3\u0026#34;] A1[\u0026#34;Actor (FSDP + TP)\u0026#34;] end subgraph CG[\u0026#34;Critic GPUs 4-5\u0026#34;] C1[\u0026#34;Critic (FSDP)\u0026#34;] end subgraph RG[\u0026#34;Reward GPU 6\u0026#34;] R1[\u0026#34;Reward (TP)\u0026#34;] end subgraph RefG[\u0026#34;Ref GPU 7\u0026#34;] Ref1[\u0026#34;Reference (TP)\u0026#34;] end end Pros[\u0026#34;(+) 每个模型有充足显存\u0026lt;br/\u0026gt;(+) 可独立选择并行策略\u0026lt;br/\u0026gt;(+) 模型间解耦\u0026#34;] Cons[\u0026#34;(-) 必须跨组传输数据\u0026lt;br/\u0026gt;(-) GPU 空闲时间长\u0026lt;br/\u0026gt;(-) 资源分配不灵活\u0026#34;] Separated --\u0026gt; Pros Separated --\u0026gt; Cons style Pros fill:#d4edda,stroke:#28a745 style Cons fill:#fff3cd,stroke:#ffc107 verl 选择了 Colocated 策略，因为它避免了分离策略中最大的问题：GPU 空闲和数据传输开销。为了解决显存压力，verl 使用 FSDP 来分片参数，并在不同阶段动态切换模型的工作模式。\n权重更新与 Resharding verl 混合引擎中最关键的操作是权重重整（resharding）——在 FSDP 分片格式和 TP 切片格式之间转换。\nflowchart TD subgraph FSDP_Format[\u0026#34;FSDP 格式 (训练后)\u0026#34;] direction LR G0s[\u0026#34;GPU 0\u0026lt;br/\u0026gt;shard 0\u0026#34;] G1s[\u0026#34;GPU 1\u0026lt;br/\u0026gt;shard 1\u0026#34;] G2s[\u0026#34;GPU 2\u0026lt;br/\u0026gt;shard 2\u0026#34;] G3s[\u0026#34;GPU 3\u0026lt;br/\u0026gt;shard 3\u0026#34;] end FSDP_Format --\u0026gt; |\u0026#34;all-gather\u0026#34;| Full[\u0026#34;Full Model\u0026#34;] Full --\u0026gt; |\u0026#34;split by dimension\u0026#34;| TP_Format subgraph TP_Format[\u0026#34;TP 格式 (生成时)\u0026#34;] direction LR G0t[\u0026#34;GPU 0\u0026lt;br/\u0026gt;col 0\u0026lt;br/\u0026gt;所有层\u0026#34;] G1t[\u0026#34;GPU 1\u0026lt;br/\u0026gt;col 1\u0026lt;br/\u0026gt;所有层\u0026#34;] G2t[\u0026#34;GPU 2\u0026lt;br/\u0026gt;col 2\u0026lt;br/\u0026gt;所有层\u0026#34;] G3t[\u0026#34;GPU 3\u0026lt;br/\u0026gt;col 3\u0026lt;br/\u0026gt;所有层\u0026#34;] end style FSDP_Format fill:#d4edda,stroke:#28a745 style TP_Format fill:#cce5ff,stroke:#007bff style Full fill:#fff3cd,stroke:#ffc107 这个过程的通信开销是 $O(P)$，其中 $P$ 是模型参数量。对于 7B 模型大约是 14 GB 的数据传输。但考虑到生成阶段会运行数百步（每步都需要通信），这个一次性开销是完全可以接受的。\nverl 的训练流程 把上面所有部分串起来，verl 的一次 PPO iteration 流程如下：\nflowchart TD subgraph Iter[\u0026#34;verl PPO Iteration\u0026#34;] subgraph Phase1[\u0026#34;1. Rollout Phase\u0026#34;] R1[\u0026#34;a) FSDP → TP resharding\u0026lt;br/\u0026gt;(Actor 权重重整)\u0026#34;] R2[\u0026#34;b) Actor (TP) 自回归生成 responses\u0026#34;] R3[\u0026#34;c) Reward Model (TP) 打分\u0026#34;] R4[\u0026#34;d) Reference (TP) 计算 log probs\u0026#34;] R5[\u0026#34;e) Critic (FSDP) 估计 values\u0026#34;] R6[\u0026#34;f) 计算 KL penalties + GAE advantages\u0026#34;] R1 --\u0026gt; R2 --\u0026gt; R3 --\u0026gt; R4 --\u0026gt; R5 --\u0026gt; R6 end subgraph Phase2[\u0026#34;2. Training Phase\u0026#34;] T1[\u0026#34;a) TP → FSDP resharding\u0026lt;br/\u0026gt;(Actor 权重重整回 FSDP)\u0026#34;] T2[\u0026#34;b) Actor (FSDP) PPO 更新 x ppo_epochs\u0026#34;] T3[\u0026#34;c) Critic (FSDP) Value 更新 x ppo_epochs\u0026#34;] T1 --\u0026gt; T2 --\u0026gt; T3 end Phase1 --\u0026gt; Phase2 Phase2 --\u0026gt; Repeat[\u0026#34;3. 重复\u0026#34;] end style Phase1 fill:#cce5ff,stroke:#007bff style Phase2 fill:#d4edda,stroke:#28a745 分布式 RLHF 策略 当模型规模增大到需要数十甚至上百张 GPU 时，RLHF 的通信模式变得非常复杂。让我们系统地梳理各个阶段的通信需求。\n各阶段通信模式 阶段 通信操作 特性 生成 (TP) All-reduce（每层之后） Latency-bound 生成 (PP) Point-to-point（流水线） Bubble overhead Actor 训练 (FSDP) All-gather + reduce-scatter Bandwidth-bound Critic 训练 (FSDP) All-gather + reduce-scatter Bandwidth-bound Reward 打分 Broadcast prompts + responses 一次性开销 Reference log probs Broadcast prompts + responses 一次性开销 权重同步 (Actor) All-gather / broadcast 每个 iteration 一次 几个关键观察：\n生成阶段是 latency-bound 的：自回归生成的每一步都需要一次 all-reduce（TP 的情况），$T$ 个 token 就需要 $T$ 次。这就是为什么生成阶段更适合用 TP 而非 FSDP——TP 的 all-reduce 量小（只有 hidden dimension），而 FSDP 的 all-gather 量大（整个参数分片）。\n训练阶段是 bandwidth-bound 的：FSDP 的 all-gather 和 reduce-scatter 传输的是完整的参数和梯度分片，数据量大但通信次数少（每层一次 all-gather + 一次 reduce-scatter）。\n权重同步是固定开销：每个 PPO iteration 只需做一次，开销与模型参数量成正比。\n生产级配置示例 以 70B 模型在 64 张 A100-80GB 上的 RLHF 训练为例：\nflowchart TD subgraph Config[\u0026#34;70B RLHF 生产配置 (共置策略)\u0026#34;] HW[\u0026#34;**硬件** 8 节点 x 8 GPU = 64 A100-80GB\u0026lt;br/\u0026gt;互联: 节点内 NVLink, 节点间 RDMA\u0026#34;] subgraph Models[\u0026#34;模型并行策略\u0026#34;] AR[\u0026#34;**Actor + Reference**\u0026lt;br/\u0026gt;训练: FSDP across 64 GPUs\u0026lt;br/\u0026gt;生成: TP=8 + PP=8\u0026#34;] CR[\u0026#34;**Critic**\u0026lt;br/\u0026gt;训练: FSDP across 64 GPUs\u0026#34;] RW[\u0026#34;**Reward Model**\u0026lt;br/\u0026gt;推理: TP=8 within node\u0026#34;] end subgraph Mem[\u0026#34;显存预算 (per GPU) — 总计 ~73 GB\u0026#34;] M1[\u0026#34;Actor FSDP shard: ~2.2 GB\u0026#34;] M2[\u0026#34;Actor optimizer: ~4.4 GB\u0026#34;] M3[\u0026#34;Critic shard + optimizer: ~6.6 GB\u0026#34;] M4[\u0026#34;Reward (TP=8): ~17.5 GB\u0026#34;] M5[\u0026#34;Reference (TP=8): ~17.5 GB\u0026#34;] M6[\u0026#34;激活值 + KV Cache: ~20 GB\u0026#34;] M7[\u0026#34;通信 buffer: ~5 GB\u0026#34;] end HW --\u0026gt; Models --\u0026gt; Mem end style HW fill:#cce5ff,stroke:#007bff style AR fill:#d4edda,stroke:#28a745 style CR fill:#d4edda,stroke:#28a745 style RW fill:#fff3cd,stroke:#ffc107 在实际部署中，还需要考虑：\nActivation Checkpointing：用计算换显存，对 Actor 和 Critic 的训练阶段尤为重要 Mixed Precision：BF16 训练 + FP32 优化器状态 Gradient Accumulation：增大有效 batch size，减少通信频率 Prompt 排序：按长度排序 prompt，减少 padding 浪费 RLHF 的前沿方向 值得一提的是，近年来 RLHF 系统还在快速演进：\nGRPO（Group Relative Policy Optimization）：DeepSeek 提出的方法，去掉了 Critic 模型，用组内 response 的相对排名来估计优势。这直接减少了 25-40% 的显存需求和相应的计算开销。\nOnline DPO：结合了 DPO（Direct Preference Optimization）和在线生成，在某些场景下可以替代 PPO。\n异步 PPO：将生成和训练异步化，提高 GPU 利用率。Actor 在训练的同时，用旧版本的权重进行下一批生成。\nKey Takeaways RLHF 需要四个模型：Actor（生成 response）、Critic（估计价值）、Reward Model（打分）、Reference（KL 锚点）。加上优化器状态，总显存需求约为 SFT 的 2.5 倍。\nPPO 的数据流是严格有序的：生成 → 打分 → KL 计算 → GAE 优势估计 → PPO 更新。四个模型之间形成 DAG 依赖，任何一步都不能跳过或并行化。\nRLHF 的核心系统挑战是推理嵌套在训练中：生成阶段是 memory-bound 的推理问题（受益于 TP + KV Cache），训练阶段是 compute-bound 的优化问题（受益于 FSDP）。同一组 GPU 需要在两种模式间切换。\nPPO 通过两个机制稳定训练：\nKL 惩罚：$D_{KL}(\\pi_\\theta | \\pi_\\text{ref})$ 防止 Reward Hacking Clipping：$\\text{clip}(r_t, 1-\\epsilon, 1+\\epsilon)$ 限制每步更新幅度 verl 的核心创新是\u0026quot;混合引擎\u0026quot;：在同一组 GPU 上动态切换 FSDP（训练）和 TP（推理）。通过共置策略避免 GPU 空闲，通过权重重整（resharding）在两种并行格式之间切换。\n规模效应：70B 模型的 RLHF 至少需要 1400 GB 显存（18 张 A100-80GB），实际部署通常使用 64-128 张 GPU 以获得合理的训练吞吐量。\n配套代码 本文配套代码位于 code/04-rlhf-system/：\nminimal_rlhf.py — 从零实现完整的 RLHF 训练循环。包含四个模型的定义、PPO 数据生成管线、PPO 训练步骤、以及系统挑战的可视化。使用小模型 (d_model=256) 在 CPU 上运行，但架构和数据流与生产系统完全一致。 运行方式：\n1 2 cd code/04-rlhf-system python minimal_rlhf.py 代码分为四个部分，对应本文的四个核心主题：\nPart 1: Four-Model Architecture — 创建四个模型，展示参数量和显存估算 Part 2: PPO Data Generation — 完整的 rollout 管线：生成 → 打分 → KL → GAE Part 3: PPO Training Step — 端到端的 PPO 训练循环，包含 Actor 和 Critic 更新 Part 4: System Challenges — 数据流图、计算开销对比、分布式策略分析 参考资料 Ouyang et al., 2022. Training language models to follow instructions with human feedback — InstructGPT，RLHF 的开创性工作\nSchulman et al., 2017. Proximal Policy Optimization Algorithms — PPO 算法原论文\nSheng et al., 2024. HybridFlow: A Flexible and Efficient RLHF Framework — verl 的技术论文\nZheng et al., 2023. Secrets of RLHF in Large Language Models — RLHF 训练的实践经验总结\nRafailov et al., 2023. Direct Preference Optimization: Your Language Model is Secretly a Reward Model — DPO，RLHF 的替代方案\nDeepSeek-AI, 2024. DeepSeek-R1 — GRPO 方法，去掉 Critic 的简化版 RLHF\nverl GitHub Repository. https://github.com/volcengine/verl — 字节跳动开源的 RLHF 训练框架\n","permalink":"https://mzf666.github.io/llm-infra/zh/posts/04-rlhf-system/","summary":"从 RLHF 四模型架构到 verl 系统实现，理解为什么 RLHF 本质上是一个系统问题。","title":"RLHF 系统设计入门"},{"content":"Motivation 假设你拿到了一个 DeepSeek-style 的 30B-A3B MoE 模型：128 个 expert，每个 token 激活 top-6，sigmoid routing。怎么训？\n先算一笔账。128 个 expert，每个 expert 是一个 SwiGLU FFN（hidden_size=1856），单个 expert 参数量约 $3 \\times 1856 \\times 2688 \\approx 15M$（gate_proj + up_proj + down_proj）。128 个 expert 共 $128 \\times 15M \\approx 1.9B$ 参数——光 expert 就占了总参数量 30B 的大头。加上 attention、embedding、shared expert 等非 expert 参数，BF16 下模型权重约 60 GB。\n这只是权重。训练时还需要：\n优化器状态：Adam 的 momentum + variance，FP32 下又是 $2 \\times 60 = 120$ GB 梯度：与权重同规模，~60 GB 激活值：随 expert 数线性增长，因为每个 token 要经过 6 个 expert 的前向计算 总计轻松突破 300 GB——4 张 A100-80GB 都不够。而且这还是\u0026quot;小\u0026quot;模型，DeepSeek V3 有 671B 参数、256 个 expert。\n显存只是第一道坎。MoE 训练面临三大工程挑战：\n显存墙：128 个 expert 的参数量远超 dense 模型。即使 active 参数只有 3B，total 参数是 30B——你需要为所有 expert 分配显存，即使大部分 expert 在每个 batch 中并不被大多数 token 激活。\n通信墙：Expert Parallel 的核心操作是 All-to-All dispatch——每个 GPU 需要把自己持有的 token 发送到对应 expert 所在的 GPU，计算完再收回来。通信量与 expert 数和 token 数成正比。当 128 个 expert 分布在 64 张 GPU 上时，All-to-All 的通信量可以轻松达到每步数 GB。\n负载不均衡：token routing 天然不均匀。某些 expert 因为学到了更通用的特征而被频繁选中（\u0026ldquo;热门 expert\u0026rdquo;），而另一些几乎无人问津。这导致部分 GPU 过载、部分 GPU 空转——系统吞吐量由最慢的 GPU 决定。\n这三个问题不是孤立的——显存约束限制了你的并行策略选择，通信开销受并行拓扑影响，负载均衡策略又会改变通信模式。MoE 训练本质上是一个联合优化问题：需要在并行策略、通信调度、计算融合、精度管理之间找到全局最优解。\nMegatron-Core + Megatron-Bridge 是这套系统方案的工业级代表。Megatron-Core 提供了 TP/PP/DP/EP/CP 五维并行的底层引擎和 MoE 加速原语（Grouped GEMM、Flex Dispatcher、通信重叠）；Megatron-Bridge 则在此之上架起了通往 HuggingFace 生态的桥梁——让你能直接加载 HF 格式的预训练权重，用 Megatron 训练，再导出回 HF 格式部署。\n本文将从系统视角出发，帮你理解：Megatron 如何组织五维并行、Bridge 如何实现权重双向转换、MoE 训练中的四层加速策略，以及如何用这套体系端到端地训练一个 30B MoE 模型。\n前置知识 GPU 显存模型与分布式通信基础（第 1 篇）——理解 All-to-All 等通信原语和 NVLink/RDMA 拓扑 分布式并行策略全景（第 2 篇）——特别是 EP、TP、FSDP 和混合并行 了解 MoE（Mixture-of-Experts）的基本概念：Router、Top-K、Expert FFN 先看一张全局架构图，后面逐一展开：\nflowchart LR subgraph HF[\u0026#34;HuggingFace 生态\u0026#34;] HF_CK[\u0026#34;HF Pretrained Checkpoint\u0026lt;br/\u0026gt;(DeepSeek, Llama, ...)\u0026#34;] end subgraph MCore[\u0026#34;Megatron 训练引擎 (Megatron-Core)\u0026#34;] PS[\u0026#34;parallel_state\u0026lt;br/\u0026gt;TP/PP/DP/EP/CP 五维网格\u0026#34;] MOE[\u0026#34;transformer.moe\u0026lt;br/\u0026gt;Router + Dispatcher\u0026lt;br/\u0026gt;Grouped GEMM\u0026#34;] ACC[\u0026#34;加速原语\u0026lt;br/\u0026gt;DeepEP / HybridEP\u0026lt;br/\u0026gt;FP8 / MXFP8 / FP4\u0026lt;br/\u0026gt;TP Comm Overlap\u0026#34;] end HF_CK -- \u0026#34;AutoBridge\u0026lt;br/\u0026gt;权重映射+切分\u0026#34; --\u0026gt; PS PS --- MOE MOE --- ACC MCore -- \u0026#34;反向转换\u0026lt;br/\u0026gt;训练完成后导出回 HF 格式\u0026lt;br/\u0026gt;用于推理部署\u0026#34; --\u0026gt; HF_CK style HF fill:#fff3cd,stroke:#856404 style MCore fill:#cce5ff,stroke:#004085 Megatron-Core 架构速览 在深入 Bridge 和 MoE 加速之前，先建立对 Megatron-Core 的整体认知。Megatron-Core 是 NVIDIA 的分布式训练引擎，为 LLM 训练提供了工业级的并行基础设施。\n五维并行网格 Megatron-Core 的并行策略组织为五维网格：TP（Tensor Parallel）、PP（Pipeline Parallel）、DP（Data Parallel）、EP（Expert Parallel）、CP（Context Parallel）。每张 GPU 在五个维度上各属于一个并行组。\nflowchart TD subgraph Grid[\u0026#34;Megatron-Core 五维并行网格\u0026lt;br/\u0026gt;假设 64 张 GPU，配置：TP=4, PP=2, DP=4, EP=2, CP=1\u0026#34;] TP[\u0026#34;\u0026lt;b\u0026gt;TP (4)\u0026lt;/b\u0026gt;\u0026lt;br/\u0026gt;张量切分\u0026lt;br/\u0026gt;all-reduce\u0026lt;br/\u0026gt;切分权重矩阵的行/列\u0026#34;] PP[\u0026#34;\u0026lt;b\u0026gt;PP (2)\u0026lt;/b\u0026gt;\u0026lt;br/\u0026gt;流水线分层\u0026lt;br/\u0026gt;send/recv\u0026lt;br/\u0026gt;切分 Transformer 层\u0026#34;] DP[\u0026#34;\u0026lt;b\u0026gt;DP (4)\u0026lt;/b\u0026gt;\u0026lt;br/\u0026gt;数据并行\u0026lt;br/\u0026gt;all-reduce\u0026lt;br/\u0026gt;切分梯度\u0026#34;] EP[\u0026#34;\u0026lt;b\u0026gt;EP (2)\u0026lt;/b\u0026gt;\u0026lt;br/\u0026gt;专家分配\u0026lt;br/\u0026gt;all-to-all\u0026lt;br/\u0026gt;切分 MoE expert\u0026#34;] CP[\u0026#34;\u0026lt;b\u0026gt;CP (1)\u0026lt;/b\u0026gt;\u0026lt;br/\u0026gt;上下文并行\u0026lt;br/\u0026gt;ring attention\u0026lt;br/\u0026gt;切分序列维度\u0026#34;] end TOTAL[\u0026#34;总 GPU 数 = TP × PP × DP × EP × CP = 4 × 2 × 4 × 2 × 1 = 64\u0026lt;br/\u0026gt;\u0026lt;i\u0026gt;注意：EP 和 DP 共享同一组 GPU，EP × DP = 数据并行维度的总 GPU 数\u0026lt;/i\u0026gt;\u0026#34;] Grid --\u0026gt; TOTAL style Grid fill:#cce5ff,stroke:#004085 style TOTAL fill:#fff3cd,stroke:#856404 parallel_state.initialize_model_parallel_group() 是并行组初始化的核心。它将所有 GPU 按五个维度划分成互不相交的通信组：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, expert_model_parallel_size: int = 1, context_parallel_size: int = 1, ... ): \u0026#34;\u0026#34;\u0026#34; 将 world_size 个 GPU 划分为五维并行组。 每个 rank 属于每个维度的一个组，用于后续通信。 \u0026#34;\u0026#34;\u0026#34; world_size = torch.distributed.get_world_size() data_parallel_size = world_size // ( tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size ) # EP 从 DP 维度中\u0026#34;借\u0026#34;GPU： # 实际 DP = data_parallel_size // expert_model_parallel_size # EP 组内的 GPU 持有不同的 expert，DP 组内的 GPU 持有相同的数据分片 关键洞察：EP 和 DP 是\u0026quot;竞争\u0026quot;关系——EP 越大，实际的 DP 越小。如果 64 张 GPU 配置 TP=4, PP=2, EP=8，那么 DP = 64 / (4 × 2) / 8 = 1——没有数据并行！这意味着 batch size 受限于单个数据并行组。生产中的 DeepSeek V3 用 TP=2, PP=16, EP=64，通过极大的 PP 来释放 EP 所需的 GPU 数量。\n核心模块架构 graph TD subgraph GPT[\u0026#34;megatron.core.models.gpt\u0026#34;] GPTModel[\u0026#34;GPTModel\u0026#34;] EMB[\u0026#34;embedding\u0026lt;br/\u0026gt;(VocabParallelEmbedding)\u0026lt;br/\u0026gt;TP 切分词表\u0026#34;] DEC[\u0026#34;decoder (TransformerBlock)\u0026#34;] OUT[\u0026#34;output_layer\u0026lt;br/\u0026gt;(ColumnParallelLinear)\u0026#34;] GPTModel --\u0026gt; EMB GPTModel --\u0026gt; DEC GPTModel --\u0026gt; OUT TL[\u0026#34;layers[] (TransformerLayer)\u0026#34;] DEC --\u0026gt; TL SA[\u0026#34;self_attention\u0026#34;] MLP[\u0026#34;mlp\u0026#34;] TL --\u0026gt; SA TL --\u0026gt; MLP QKV[\u0026#34;linear_qkv\u0026lt;br/\u0026gt;fused QKV, TP column parallel\u0026#34;] PROJ[\u0026#34;linear_proj\u0026lt;br/\u0026gt;output proj, TP row parallel\u0026#34;] SA --\u0026gt; QKV SA --\u0026gt; PROJ DENSE[\u0026#34;[Dense] linear_fc1 / linear_fc2\u0026#34;] MOE_SUB[\u0026#34;[MoE] router + experts\u0026#34;] MLP --\u0026gt; DENSE MLP --\u0026gt; MOE_SUB ROUTER[\u0026#34;router\u0026lt;br/\u0026gt;Top-K / Sigmoid 路由\u0026#34;] DISP[\u0026#34;token_dispatcher\u0026lt;br/\u0026gt;All-to-All 调度\u0026#34;] EXPERTS[\u0026#34;experts\u0026lt;br/\u0026gt;GroupedMLP / SequentialMLP\u0026#34;] SHARED[\u0026#34;shared_experts\u0026lt;br/\u0026gt;共享专家 (可选)\u0026#34;] MOE_SUB --\u0026gt; ROUTER MOE_SUB --\u0026gt; DISP MOE_SUB --\u0026gt; EXPERTS MOE_SUB --\u0026gt; SHARED end subgraph DIST[\u0026#34;megatron.core.distributed\u0026#34;] DDP[\u0026#34;DistributedDataParallel\u0026lt;br/\u0026gt;gradient all-reduce\u0026#34;] DOPT[\u0026#34;DistributedOptimizer\u0026lt;br/\u0026gt;ZeRO-style 优化器状态切分\u0026#34;] end subgraph TMOE[\u0026#34;megatron.core.transformer.moe\u0026#34;] MOEL[\u0026#34;MoELayer — MoE 层封装\u0026#34;] TKR[\u0026#34;TopKRouter / SigmoidRouter\u0026#34;] A2A[\u0026#34;AllToAllTokenDispatcher\u0026#34;] FLEX[\u0026#34;FlexTokenDispatcher\u0026lt;br/\u0026gt;DeepEP / HybridEP 后端\u0026#34;] GMLP[\u0026#34;GroupedMLP\u0026lt;br/\u0026gt;多 expert 融合计算\u0026#34;] end style GPT fill:#cce5ff,stroke:#004085 style DIST fill:#d4edda,stroke:#155724 style TMOE fill:#fff3cd,stroke:#856404 Megatron-Core 通过 get_gpt_decoder_block_spec() 定义 decoder 的层规格（layer spec）。对于 MoE 模型，MoE 层和 dense 层可以交替出现——这是 Nemotron-3-Nano 等 hybrid 架构的基础。\nTransformerConfig 控制所有配置，包括 MoE 相关的关键参数：\n1 2 3 4 5 6 7 8 9 10 TransformerConfig( num_moe_experts=128, # expert 总数 moe_router_topk=6, # 每个 token 激活的 expert 数 moe_router_score_function=\u0026#34;sigmoid\u0026#34;, # 路由评分函数 moe_grouped_gemm=True, # 是否用 Grouped GEMM 融合计算 moe_token_dispatcher_type=\u0026#34;flex\u0026#34;, # 调度器类型 moe_shared_expert_overlap=True, # 共享专家与路由专家并行 expert_model_parallel_size=8, # EP 大小 ... ) Megatron-Bridge 桥接机制 Megatron-Core 的训练性能是一流的，但它有一个实际门槛：你不能直接拿 HuggingFace 格式的预训练权重开始训练。Megatron 有自己的权重命名和存储格式——fused QKV、column/row parallel 切分、Grouped GEMM 的 expert 权重布局——与 HF 格式完全不同。\nMegatron-Bridge 解决的就是这个问题：它在 HF 生态和 Megatron 训练引擎之间架起了一座双向桥梁。\n数据流全景 flowchart TD A[\u0026#34;① 加载 HF 权重\u0026lt;br/\u0026gt;HuggingFace Hub\u0026lt;br/\u0026gt;\u0026lt;i\u0026gt;deepseek-ai/DeepSeek-V3\u0026lt;/i\u0026gt;\u0026#34;] B[\u0026#34;② AutoBridge 自动识别\u0026lt;br/\u0026gt;检测 HF config → 匹配 DeepSeekV3Bridge\u0026#34;] C[\u0026#34;③ 配置翻译 (provider_bridge)\u0026lt;br/\u0026gt;HF config → TransformerConfig\u0026lt;br/\u0026gt;num_attention_heads → num_attention_heads\u0026lt;br/\u0026gt;num_key_value_heads → num_query_groups\u0026lt;br/\u0026gt;intermediate_size → ffn_hidden_size\u0026lt;br/\u0026gt;n_routed_experts → num_moe_experts\u0026#34;] D[\u0026#34;④ 权重映射 (mapping_registry)\u0026lt;br/\u0026gt;Q, K, V 分离 → fused QKV\u0026lt;br/\u0026gt;gate_proj + up_proj → fused fc1\u0026lt;br/\u0026gt;per-expert 权重 → GroupedMLP 布局\u0026#34;] E[\u0026#34;⑤ 分布式切分 (scatter_to_tp_ranks)\u0026lt;br/\u0026gt;根据 TP/PP/EP 配置，将权重切分到各 GPU\u0026#34;] F[\u0026#34;⑥ Megatron GPTModel 就绪\u0026lt;br/\u0026gt;开始训练！\u0026#34;] G[\u0026#34;⑦ 训练完成后：反向转换\u0026lt;br/\u0026gt;Megatron → gather_from_tp_ranks → unfuse → HF 格式\u0026lt;br/\u0026gt;导出到 HF Hub，用于推理部署\u0026#34;] A --\u0026gt; B --\u0026gt; C --\u0026gt; D --\u0026gt; E --\u0026gt; F F --\u0026gt; G G -.-\u0026gt; A style A fill:#fff3cd,stroke:#856404 style F fill:#d4edda,stroke:#155724 style G fill:#cce5ff,stroke:#004085 AutoBridge：一行代码加载任意 HF 模型 AutoBridge 是 Bridge 的用户入口。它通过注册机制自动匹配 HF 模型类型：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 from megatron.bridge import AutoBridge # 一行代码：加载 HF 预训练模型，自动检测架构 bridge = AutoBridge.from_hf_pretrained( \u0026#34;deepseek-ai/DeepSeek-V3\u0026#34;, trust_remote_code=True ) # 转换为 Megatron provider，配置并行策略 provider = bridge.to_megatron_provider() provider.tensor_model_parallel_size = 2 provider.pipeline_model_parallel_size = 16 provider.expert_model_parallel_size = 64 provider.finalize() # 执行权重转换 + 分布式切分 # 获取分布式模型，开始训练 model = provider.provide_distributed_model(wrap_with_ddp=False) 自动检测机制：每个模型 bridge 通过装饰器注册自己支持的 HF 模型类：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 @MegatronModelBridge.register_bridge( source=\u0026#34;DeepseekV3ForCausalLM\u0026#34;, # HF 模型类名 target=GPTModel, # Megatron 目标模型 provider=MLAModelProvider, # 模型提供器 model_type=\u0026#34;deepseek_v3\u0026#34;, ) class DeepSeekV3Bridge(MegatronModelBridge): CONFIG_MAPPING = [ (\u0026#34;num_attention_heads\u0026#34;, \u0026#34;num_attention_heads\u0026#34;), (\u0026#34;num_key_value_heads\u0026#34;, \u0026#34;num_query_groups\u0026#34;), (\u0026#34;n_routed_experts\u0026#34;, \u0026#34;num_moe_experts\u0026#34;), (\u0026#34;num_experts_per_tok\u0026#34;, \u0026#34;moe_router_topk\u0026#34;), # ... 更多字段映射 ] 调用 AutoBridge.from_hf_pretrained() 时，它读取 HF config 中的 architectures 字段（如 \u0026quot;DeepseekV3ForCausalLM\u0026quot;），在注册表中查找对应的 bridge class，然后执行配置翻译和权重映射。\n目前已注册的模型覆盖了主流 LLM 家族：Llama 2/3、DeepSeek V2/V3、Qwen 2/2.5/3、Gemma、Mistral、Mamba，以及 NVIDIA 自家的 Nemotron 系列。\n权重映射策略 权重映射是 Bridge 的核心技术挑战。HF 和 Megatron 在以下方面存在本质差异：\nflowchart LR subgraph HF[\u0026#34;HuggingFace 格式\u0026#34;] HF_QKV[\u0026#34;QKV 投影: 分离的\u0026lt;br/\u0026gt;q_proj, k_proj, v_proj\u0026#34;] HF_MLP[\u0026#34;MLP 上投影: 分离的\u0026lt;br/\u0026gt;gate_proj, up_proj\u0026#34;] HF_EXP[\u0026#34;Expert 存储:\u0026lt;br/\u0026gt;per-expert 独立参数\u0026#34;] HF_TP[\u0026#34;TP 切分: 无\u0026lt;br/\u0026gt;(单卡完整权重)\u0026#34;] HF_PP[\u0026#34;PP 切分: 无\u0026lt;br/\u0026gt;(所有层在一起)\u0026#34;] HF_EP[\u0026#34;EP 切分: 无\u0026lt;br/\u0026gt;(所有 expert 在一起)\u0026#34;] end subgraph MG[\u0026#34;Megatron-Core 格式\u0026#34;] MG_QKV[\u0026#34;QKV 投影:\u0026lt;br/\u0026gt;fused linear_qkv\u0026#34;] MG_MLP[\u0026#34;MLP 上投影:\u0026lt;br/\u0026gt;fused linear_fc1\u0026#34;] MG_EXP[\u0026#34;Expert 存储:\u0026lt;br/\u0026gt;GroupedMLP 融合\u0026#34;] MG_TP[\u0026#34;TP 切分:\u0026lt;br/\u0026gt;按行/列切分到多卡\u0026#34;] MG_PP[\u0026#34;PP 切分:\u0026lt;br/\u0026gt;按层分配到不同 stage\u0026#34;] MG_EP[\u0026#34;EP 切分:\u0026lt;br/\u0026gt;expert 分配到不同卡\u0026#34;] end HF_QKV -.-\u0026gt; MG_QKV HF_MLP -.-\u0026gt; MG_MLP HF_EXP -.-\u0026gt; MG_EXP HF_TP -.-\u0026gt; MG_TP HF_PP -.-\u0026gt; MG_PP HF_EP -.-\u0026gt; MG_EP style HF fill:#fff3cd,stroke:#856404 style MG fill:#cce5ff,stroke:#004085 Bridge 通过一组 MegatronParamMapping 子类处理这些差异：\nQKVMapping — 最复杂的映射之一。HF 存储分离的 Q、K、V 权重，Megatron 将它们融合为一个 linear_qkv 张量。融合时还要考虑 GQA（Grouped Query Attention）：Q 的 head 数可能是 K/V 的 N 倍，交错排列确保 TP 切分后每个 GPU 拿到完整的 Q-K-V head 组。\n1 2 3 4 5 6 7 8 9 10 11 # HF: 分离的 Q, K, V # q_proj.weight: (num_heads * head_dim, hidden_size) # k_proj.weight: (num_kv_heads * head_dim, hidden_size) # v_proj.weight: (num_kv_heads * head_dim, hidden_size) # Megatron: 融合的 QKV，交错排列以支持 TP # linear_qkv.weight: ((num_heads + 2*num_kv_heads) * head_dim, hidden_size) # 排列方式：[Q_group0, K0, V0, Q_group1, K1, V1, ...] # TP 切分：沿第一维（output dim）均匀切分 # GPU 0 拿到前 1/TP 的 head 组，GPU 1 拿到下一组 ... GatedMLPMapping — 处理 SwiGLU 等 gated activation 的 MLP。HF 分开存储 gate_proj 和 up_proj，Megatron 将它们拼接为 linear_fc1：\n1 2 3 4 5 6 # HF: gate_proj.weight (ffn_hidden, hidden) + up_proj.weight (ffn_hidden, hidden) # Megatron: linear_fc1.weight (2 * ffn_hidden, hidden) ← 拼接 # 对于 MoE expert，每个 expert 的 gate+up 独立融合 # GroupedMLP 模式下，所有 expert 的 fc1 存储为 weight0, weight1, ... weightN # 支持 CUTLASS GroupedGEMM 一次性计算所有 expert 反向转换同样重要：训练完成后，megatron_to_hf 将 fused 权重拆分回 HF 格式。Bridge 的 gather_from_tp_ranks 先从各 TP rank 收集切片，然后 unfuse 回分离的权重。这使得 Megatron 训练的模型可以无缝用 HF 生态（vLLM、SGLang 等）部署。\nBridge 的价值 Bridge 消除了 Megatron 的生态锁定。在此之前，要用 Megatron 训练一个 HF 模型，需要手动编写转换脚本——不同模型架构的 QKV 排列方式、MoE expert 布局、attention 变体（MHA/GQA/MLA）各不相同，每个模型都需要定制转换逻辑。Bridge 将这些转换标准化为可复用的 mapping 策略，新增模型支持只需注册新的 bridge class 和 mapping。\nMoE 加速深度剖析 这是本文的核心章节。MoE 训练的加速可以分为四个层次：并行策略、通信调度、计算融合、模型实例。我们逐层展开。\nExpert Parallel 基础 第 2 篇已经介绍了 EP 的基本概念：将 expert 分配到不同 GPU，通过 All-to-All dispatch/combine 实现 token 路由。这里聚焦 EP 在实际大规模训练中的工程细节。\nEP 的工作流分为三步：\nflowchart TD subgraph S1[\u0026#34;Step 1: Router 决策\u0026#34;] INPUT[\u0026#34;Input tokens: t0, t1, t2, ..., t7\u0026#34;] ROUTER[\u0026#34;Router(x) → TopK scores\u0026lt;br/\u0026gt;t0 → Expert 3,7,12,45,67,99 (top-6 sigmoid)\u0026lt;br/\u0026gt;t1 → Expert 1,5,23,44,88,101\u0026lt;br/\u0026gt;...\u0026#34;] INPUT --\u0026gt; ROUTER end subgraph S2[\u0026#34;Step 2: All-to-All Dispatch (token 重分布)\u0026#34;] BEFORE[\u0026#34;\u0026lt;b\u0026gt;Before:\u0026lt;/b\u0026gt; GPU 按数据分片持有 token\u0026lt;br/\u0026gt;GPU 0: [t0, t1] \u0026amp;nbsp; GPU 1: [t2, t3] \u0026amp;nbsp; ...\u0026#34;] A2A1[\u0026#34;═══ All-to-All ═══\u0026#34;] AFTER[\u0026#34;\u0026lt;b\u0026gt;After:\u0026lt;/b\u0026gt; GPU 按 expert 分组持有 token\u0026lt;br/\u0026gt;GPU 0 (E0-E15): 收到所有发给 E0-E15 的 token\u0026lt;br/\u0026gt;GPU 1 (E16-E31): 收到所有发给 E16-E31 的 token\u0026#34;] BEFORE --\u0026gt; A2A1 --\u0026gt; AFTER end subgraph S3[\u0026#34;Step 3: Expert Compute + All-to-All Combine\u0026#34;] COMPUTE[\u0026#34;各 GPU 对收到的 token 运行自己的 expert FFN\u0026lt;br/\u0026gt;GPU 0: GroupedGEMM(E0-E15, tokens)\u0026#34;] A2A2[\u0026#34;═══ All-to-All ═══\u0026#34;] RESULT[\u0026#34;token 结果送回原始 GPU\u0026lt;br/\u0026gt;恢复数据分片布局\u0026#34;] COMPUTE --\u0026gt; A2A2 --\u0026gt; RESULT end S1 --\u0026gt; S2 --\u0026gt; S3 style S1 fill:#fff3cd,stroke:#856404 style S2 fill:#cce5ff,stroke:#004085 style S3 fill:#d4edda,stroke:#155724 与其他并行维度的组合是 MoE 训练的关键设计决策。以 DeepSeek V3（671B，256 expert）为例：\n并行维度 大小 原因 TP 2 MLA attention 的 KV 维度较小，TP=2 即可切分 PP 16 61 层 Transformer，16 stage 流水线，每 stage ~4 层 EP 64 256 expert / 64 = 4 expert per GPU，显存可控 DP 1 GPU 总数有限，EP 和 PP 占满后无多余 GPU 做数据并行 CP 1 训练序列长度 8K，不需要上下文并行 总 GPU 数：$2 \\times 16 \\times 64 \\times 1 \\times 1 = 2048$ 张。注意 DP=1 意味着每个 micro-batch 只在一组 GPU 上处理，全局 batch size 通过 gradient accumulation 实现。\nDeepEP 与 Flex Dispatcher 标准的 All-to-All 通信有一个问题：它是同步的。所有 GPU 必须同时参与 dispatch 和 combine，期间 GPU 的计算单元闲置。对于 128 expert、EP=8 的配置，每个 MoE 层需要两次 All-to-All，而一个 30B 模型可能有 40+ 个 MoE 层——通信开销累积起来非常可观。\nDeepEP（Deep Expert Parallelism）和 Flex Dispatcher 是 Megatron-Core 对 All-to-All 的优化方案。\nflowchart TD subgraph STD[\u0026#34;标准 All-to-All (AllToAllTokenDispatcher)\u0026#34;] direction LR S_DISP[\u0026#34;dispatch\u0026#34;] --\u0026gt; S_WAIT1[\u0026#34;等待通信完成\u0026lt;br/\u0026gt;(GPU 计算闲置)\u0026#34;] --\u0026gt; S_COMP[\u0026#34;expert compute\u0026#34;] S_COMP --\u0026gt; S_COMB[\u0026#34;combine\u0026#34;] --\u0026gt; S_WAIT2[\u0026#34;等待通信完成\u0026#34;] end subgraph FLEX[\u0026#34;Flex Dispatcher (FlexTokenDispatcher + DeepEP)\u0026#34;] direction LR F_DISP[\u0026#34;dispatch 开始\u0026lt;br/\u0026gt;(异步通信)\u0026#34;] --\u0026gt; F_COMP[\u0026#34;expert compute 开始\u0026lt;br/\u0026gt;(收到一部分就开始算)\u0026#34;] F_DISP --\u0026gt; F_DONE[\u0026#34;dispatch 完成\u0026#34;] F_COMP --\u0026gt; F_ALL[\u0026#34;全部 expert compute 完成\u0026#34;] F_ALL --\u0026gt; F_COMB[\u0026#34;combine (异步)\u0026#34;] end KEY[\u0026#34;关键：通信和计算可以部分重叠\u0026#34;] STD ~~~ FLEX FLEX ~~~ KEY style STD fill:#fff3cd,stroke:#856404 style FLEX fill:#d4edda,stroke:#155724 style KEY fill:#cce5ff,stroke:#004085 Megatron-Bridge 通过配置选择 dispatcher 后端：\n1 2 3 4 5 6 7 8 9 10 11 # 标准 All-to-All（默认） cfg.model.moe_token_dispatcher_type = \u0026#34;alltoall\u0026#34; # DeepEP 后端（Ampere/Hopper/Blackwell） cfg.model.moe_token_dispatcher_type = \u0026#34;flex\u0026#34; cfg.model.moe_flex_dispatcher_backend = \u0026#34;deepep\u0026#34; # HybridEP 后端（GB200 NVL72 最优，也支持其他 GPU） cfg.model.moe_token_dispatcher_type = \u0026#34;flex\u0026#34; cfg.model.moe_flex_dispatcher_backend = \u0026#34;hybridep\u0026#34; cfg.model.moe_hybridep_num_sms = 16 # 分配给通信的 SM 数量 DeepEP 的核心思想：利用 RDMA 和 NVLink 的异步传输能力，将 All-to-All 拆分为多个小批次，每个批次的 token 到达后立即开始 expert 计算，无需等待全部 token 到齐。这在 expert 计算量不均（负载不均衡）时尤其有效——热门 expert 的 GPU 可以尽早开始计算，不用等冷门 expert 的 GPU。\nHybridEP 更进一步，专门为 NVIDIA GB200 NVL72 拓扑优化——NVL72 将 72 个 GPU 通过 NVSwitch 全互联，提供 1.8 TB/s 的对分带宽。HybridEP 根据 GPU 间是 NVLink 直连还是需要跨 switch，自动选择最优的通信路径。\n重要约束：使用 Flex Dispatcher 时，moe_shared_expert_overlap 必须设为 False。因为 Flex Dispatcher 本身已经在做通信-计算重叠，再叠加 shared expert overlap 会导致 CUDA stream 竞争。\n计算融合 通信优化解决的是\u0026quot;数据搬运\u0026quot;问题，计算融合解决的是\u0026quot;计算效率\u0026quot;问题。MoE 层的计算有三个维度的融合机会。\nGrouped GEMM 这是 MoE 计算加速最重要的一项优化。\n问题：标准 MoE 实现中，N 个 expert 各自独立做矩阵乘法。如果 EP=8、每 GPU 持有 16 个 expert，每个 MoE 层需要 16 次独立的 GEMM 调用。每次 GEMM 都有 kernel launch 开销和 GPU 利用率损失（因为单个 expert 处理的 token 数可能很少）。\n解决方案：将所有 expert 的矩阵乘法合并为一次 Grouped GEMM 调用。CUTLASS 的 GroupedGEMM kernel 接受一组形状可能不同的矩阵乘法，在 GPU 上统一调度执行。\nflowchart TD subgraph SEQ[\u0026#34;SequentialMLP (每个 expert 独立)\u0026#34;] direction LR E0[\u0026#34;GEMM E0\u0026lt;br/\u0026gt;20 tok\u0026#34;] E1[\u0026#34;GEMM E1\u0026lt;br/\u0026gt;5 tok\u0026#34;] E2[\u0026#34;GEMM E2\u0026lt;br/\u0026gt;35 tok\u0026#34;] EN[\u0026#34;...\u0026lt;br/\u0026gt;GEMM E15\u0026lt;br/\u0026gt;8 tok\u0026#34;] E0 --\u0026gt; E1 --\u0026gt; E2 --\u0026gt; EN end SEQ_NOTE[\u0026#34;16 次 kernel launch，GPU 利用率低\u0026#34;] subgraph GRP[\u0026#34;GroupedMLP (融合)\u0026#34;] GGEMM[\u0026#34;GroupedGEMM\u0026lt;br/\u0026gt;[E0:20tok, E1:5tok, E2:35tok, ...]\u0026lt;br/\u0026gt;1 次 kernel launch，GPU 利用率高\u0026#34;] end subgraph STORE[\u0026#34;权重存储对比\u0026#34;] direction LR S1[\u0026#34;Sequential:\u0026lt;br/\u0026gt;expert_0.fc1.weight\u0026lt;br/\u0026gt;expert_1.fc1.weight, ...\u0026#34;] S2[\u0026#34;Grouped:\u0026lt;br/\u0026gt;linear_fc1.weight0\u0026lt;br/\u0026gt;linear_fc1.weight1, ...\u0026lt;br/\u0026gt;(连续存储，便于 CUTLASS 批量访问)\u0026#34;] end SEQ --- SEQ_NOTE SEQ_NOTE ~~~ GRP GRP ~~~ STORE style SEQ fill:#fff3cd,stroke:#856404 style GRP fill:#d4edda,stroke:#155724 style STORE fill:#cce5ff,stroke:#004085 在 Megatron-Core 中启用 Grouped GEMM 只需一个配置项：\n1 cfg.model.moe_grouped_gemm = True # 默认 False 开启后，MoE 层使用 GroupedMLP 而非 SequentialMLP，expert 权重以融合格式存储。Bridge 的 GatedMLPMapping 自动处理 HF 格式到 Grouped 格式的转换。\nPermute Fusion 在 All-to-All dispatch 之后，token 需要按 expert 分组重排（permute），然后送入 expert 计算。标准实现中，permute 是一次独立的内存拷贝操作。\nPermute Fusion 将 token 重排与后续的 GEMM 融合为一个 kernel——在读取 token hidden states 时\u0026quot;顺便\u0026quot;完成重排，避免额外的内存读写。\n1 cfg.model.moe_permute_fusion = True # 融合 permute 与 GEMM 这个优化看起来简单，但对于 hidden_size=2688、batch 内 token 数以万计的场景，减少一次全量 hidden states 的内存拷贝可以节省数 GB/s 的带宽。\nShared Expert Overlap DeepSeek-style MoE 有两种 expert：所有 token 都要经过的 shared expert，和经过 router 选择的 routed expert。标准实现中两者串行执行。\nShared Expert Overlap 让 shared expert 和 routed expert 在不同 CUDA stream 上并行计算：\nflowchart TD subgraph SERIAL[\u0026#34;串行（默认）\u0026#34;] direction LR R_DISP[\u0026#34;Routed Expert\u0026lt;br/\u0026gt;Dispatch\u0026#34;] --\u0026gt; R_GEMM[\u0026#34;Routed\u0026lt;br/\u0026gt;GEMM\u0026#34;] --\u0026gt; R_COMB[\u0026#34;Combine\u0026#34;] --\u0026gt; S_GEMM1[\u0026#34;Shared\u0026lt;br/\u0026gt;GEMM\u0026#34;] end SERIAL_T[\u0026#34;总时间 = T_dispatch + T_routed + T_combine + T_shared\u0026#34;] subgraph OVERLAP[\u0026#34;并行（overlap=True）\u0026#34;] direction LR STREAM0[\u0026#34;Stream 0:\u0026lt;br/\u0026gt;Routed Dispatch → Routed GEMM → Combine\u0026#34;] STREAM1[\u0026#34;Stream 1:\u0026lt;br/\u0026gt;Shared GEMM (与 Stream 0 并行)\u0026#34;] end OVERLAP_T[\u0026#34;总时间 = max(T_dispatch + T_routed + T_combine, T_shared)\u0026lt;br/\u0026gt;节省：T_shared 被完全隐藏（通常 shared expert 较小）\u0026#34;] SERIAL --- SERIAL_T SERIAL_T ~~~ OVERLAP OVERLAP --- OVERLAP_T style SERIAL fill:#fff3cd,stroke:#856404 style OVERLAP fill:#d4edda,stroke:#155724 style OVERLAP_T fill:#cce5ff,stroke:#004085 1 2 cfg.model.moe_shared_expert_overlap = True # 注意：与 Flex Dispatcher 不兼容，二者只能选其一 其他融合优化 Megatron-Core 还提供了一系列通用计算融合：\n1 2 3 4 cfg.model.cross_entropy_loss_fusion = True # 融合 cross-entropy loss 计算 cfg.model.gradient_accumulation_fusion = True # 梯度累积融合 cfg.model.bias_activation_fusion = True # bias + activation 融合 cfg.model.bias_dropout_fusion = True # bias + dropout 融合 这些融合减少了 kernel launch 次数和中间结果的显存占用，对 MoE 训练的影响虽然不如 Grouped GEMM 显著，但累积效果不可忽视。\nNemotron-3-Nano (30B-A3B) 实例 NVIDIA 的 Nemotron-3-Nano 是一个很好的 MoE 工程实例——它不仅是 MoE，还是 Hybrid Mamba + Transformer + MoE 架构。\n模型架构 参数 值 总参数量 30B 活跃参数量 3B (per token) 层数 52 Hidden size 2688 Attention heads 32 (GQA, 2 KV groups) Expert 数 128 Expert FFN hidden 1856 Shared expert hidden 3712 (2x expert) Router top-k 6 路由函数 Sigmoid + expert bias 激活函数 Squared ReLU 最大序列长度 262,144 Hybrid 层模式：52 层中混合了三种层类型：\nflowchart LR subgraph Pattern[\u0026#34;Nemotron-3-Nano 52 层 Hybrid 模式\u0026#34;] direction TB ROW1[\u0026#34;M E M E M * E M E M E M * E M E M E M * E M E M E M *\u0026#34;] ROW2[\u0026#34;E M E M E M E M * E M E M E M E M E\u0026#34;] subgraph Legend[\u0026#34;图例\u0026#34;] direction LR LM[\u0026#34;M = Mamba (SSM)\\n线性复杂度\u0026#34;] LE[\u0026#34;E = MoE (专家层)\\n稀疏激活\u0026#34;] LA[\u0026#34;* = Attention\\n标准注意力\u0026#34;] end NOTE[\u0026#34;52 层中：Attention 仅 ~7 层\\n大部分是 Mamba 和 MoE 交替\u0026#34;] end style LM fill:#d4edda,stroke:#155724 style LE fill:#fff3cd,stroke:#856404 style LA fill:#cce5ff,stroke:#004085 style NOTE fill:#f8f9fa,stroke:#6c757d 这个设计的动机是：Attention 的计算复杂度是 $O(n^2)$，Mamba 是 $O(n)$。对于长序列（262K tokens），用少量 Attention 层捕获全局依赖，用 Mamba 层处理局部上下文，比纯 Transformer 在推理效率上有巨大优势——同时通过 MoE 保持了大参数量带来的模型容量。\n并行配置 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 # 并行策略 cfg.model.tensor_model_parallel_size = 4 # TP=4 cfg.model.pipeline_model_parallel_size = 1 # 无流水线并行 cfg.model.expert_model_parallel_size = 8 # EP=8 cfg.model.sequence_parallel = True # SP 开启 # MoE 加速 cfg.model.moe_grouped_gemm = True # Grouped GEMM cfg.model.moe_permute_fusion = True # Permute 融合 cfg.model.moe_router_score_function = \u0026#34;sigmoid\u0026#34; # Sigmoid 路由 cfg.model.moe_router_enable_expert_bias = True # Expert bias 修正 cfg.model.moe_aux_loss_coeff = 0.0001 # 负载均衡 loss 系数 # 计算融合 cfg.model.cross_entropy_loss_fusion = True cfg.model.gradient_accumulation_fusion = True cfg.model.bias_activation_fusion = True cfg.model.use_fused_weighted_squared_relu = True # Nemotron 专用融合 128 个 expert，EP=8，每张 GPU 持有 16 个 expert。TP=4 意味着非 expert 参数（attention、embedding 等）在 4 张 GPU 间切分。假设用 32 张 GPU（4 节点 × 8 GPU），则 DP = 32 / (4 × 1) / 8 = 1——单数据并行，全局 batch size 通过 gradient accumulation 扩大。\n与 DeepSeek V3 的对比 维度 Nemotron-3-Nano (30B-A3B) DeepSeek V3 (671B) 架构 Hybrid Mamba + Transformer + MoE 纯 Transformer + MoE Expert 数 128 256 Top-K 6 8 Shared expert 有（hidden=3712） 有（hidden=2×FFN） Attention MLA (Multi-Latent Attention) via GQA MLA (原生) TP 4 2 PP 1 16 EP 8 64 总 GPU ~32 ~2048 特殊层 Mamba SSM Multi-Token Prediction 路由 Sigmoid + bias Sigmoid + bias 两者的 MoE 路由策略几乎相同（sigmoid scoring + expert bias 修正），但并行策略差异巨大——DeepSeek V3 的 EP=64 和 PP=16 是其 256 expert、61 层深度的必然选择。Nemotron-3-Nano 因为模型较小（52 层、128 expert）且有 Mamba 层替代部分 Attention，可以用更简单的 TP=4 + EP=8 配置。\n通信隐藏与混合精度 上一节聚焦 MoE 层内部的加速。但 Transformer 模型不只有 MoE 层——attention、embedding、LayerNorm 等 dense 组件同样需要 TP 通信，而整个训练过程的精度策略直接影响显存和吞吐量。\nTP 通信重叠 Tensor Parallelism 的每一层都需要通信：column parallel 层在前向传播后做 all-reduce（或 reduce-scatter），row parallel 层在前向传播前做 all-gather。对于一个 52 层的模型，每层 2 次 TP 通信（attention + MLP），每步就是 100+ 次通信操作。\n**通信重叠（Comm Overlap）**的核心思想：不要等通信完成再开始下一步计算，而是让通信和计算在不同的硬件单元上同时进行——GEMM 用 Tensor Core，通信用 NVLink/PCIe 的 DMA 引擎。\nMegatron-Core 提供三种重叠策略：\nflowchart TD subgraph PIPE[\u0026#34;1. Pipeline Overlap (流水线重叠)\u0026#34;] direction LR C0[\u0026#34;Chunk 0: GEMM\u0026#34;] --\u0026gt; CC0[\u0026#34;comm\u0026#34;] C1[\u0026#34;Chunk 1: GEMM\u0026#34;] --\u0026gt; CC1[\u0026#34;comm\u0026#34;] C2[\u0026#34;Chunk 2: GEMM\u0026#34;] --\u0026gt; CC2[\u0026#34;comm\u0026#34;] end PIPE_NOTE[\u0026#34;通信和下一个 chunk 的 GEMM 在不同 SM 上并行\u0026lt;br/\u0026gt;适用于：前向传播的 fprop\u0026#34;] subgraph RING[\u0026#34;2. Ring-Exchange Overlap (环形交换重叠)\u0026#34;] direction LR RS0[\u0026#34;Step 0: recv chunk\u0026lt;br/\u0026gt;+ GEMM on local chunk\u0026#34;] RS1[\u0026#34;Step 1: recv chunk\u0026lt;br/\u0026gt;+ GEMM on received chunk\u0026#34;] RS0 --\u0026gt; RS1 end RING_NOTE[\u0026#34;每步接收一个 chunk，同时计算上一步收到的 chunk\u0026lt;br/\u0026gt;适用于：前向传播，特别是 FP8 场景\u0026#34;] subgraph BULK[\u0026#34;3. Bulk Overlap (批量重叠)\u0026#34;] direction LR B_GEMM[\u0026#34;GEMM (大部分 SM)\u0026#34;] B_COMM[\u0026#34;Comm (少量 SM)\u0026#34;] end BULK_NOTE[\u0026#34;不拆分 GEMM，在专用 SM 上启动整个通信操作\u0026lt;br/\u0026gt;适用于：反向传播的 dgrad / wgrad\u0026#34;] PIPE --- PIPE_NOTE PIPE_NOTE ~~~ RING RING --- RING_NOTE RING_NOTE ~~~ BULK BULK --- BULK_NOTE style PIPE fill:#cce5ff,stroke:#004085 style RING fill:#d4edda,stroke:#155724 style BULK fill:#fff3cd,stroke:#856404 实际配置模式：Megatron-Bridge 提供了针对不同硬件的预配置 profile。以 H100 + TP=4 为例：\n操作 阶段 BF16 策略 FP8 策略 qkv_fprop 前向 Pipeline Ring-Exchange proj_fprop 前向 Pipeline Ring-Exchange fc1_fprop 前向 Pipeline Ring-Exchange fc2_fprop 前向 Ring-Exchange Ring-Exchange qkv_dgrad 反向 Bulk Bulk proj_dgrad 反向 Bulk Bulk qkv_wgrad 反向 Bulk Bulk 规律：前向传播倾向于 Pipeline 或 Ring-Exchange（GEMM 和通信量相当），反向传播倾向于 Bulk（GEMM 更大，通信可以被完全隐藏在 GEMM 之后）。FP8 场景下，GEMM 速度翻倍但通信量不变，因此通信成为更大的瓶颈，Ring-Exchange 的细粒度重叠更有优势。\nFP8 训练 FP8（8-bit floating point）是 Hopper 及更新架构的核心加速能力。相比 BF16，FP8 的 Tensor Core 算力翻倍，显存带宽需求减半。但 FP8 的动态范围很小（E4M3 格式只有 ~480），需要精心管理缩放因子。\nMegatron-Core 支持四种 FP8 策略：\n策略 说明 适用硬件 tensorwise 每个 tensor 一个缩放因子 Hopper / Blackwell delayed 基于历史 amax 延迟计算缩放因子 Hopper / Blackwell blockwise 按 block 粒度缩放 Hopper mxfp8 Microscaling FP8，更细粒度的缩放 Blackwell 1 2 3 4 5 6 7 8 9 10 11 from megatron.bridge.training import MixedPrecisionConfig cfg.mixed_precision = MixedPrecisionConfig( fp8=\u0026#34;e4m3\u0026#34;, # 启用 FP8 E4M3 格式 fp8_recipe=\u0026#34;tensorwise\u0026#34;, # 缩放策略 fp8_param=True, # 参数也存为 FP8 fp8_param_gather=True, # all-gather 时用 FP8 通信 fp8_wgrad=True, # 权重梯度用 FP8 first_last_layers_bf16=True, # 首尾层保持 BF16（稳定性） moe_router_padding_for_fp8=True, # MoE router 输出对齐 FP8 ) 对 MoE 的影响：FP8 对 MoE 的 Grouped GEMM 特别有益。128 个 expert 的矩阵乘法在 BF16 下已经是 memory-bound（每个 expert 处理的 token 少），FP8 将访存量减半，直接提升 roofline 利用率。fp8_param_gather=True 还减少了 FSDP all-gather 的通信量——参数用 FP8 传输，到目标 GPU 后再 upcast。\nActivation Recomputation MoE 模型的激活值显存随 expert 数线性增长。每个 token 经过 top-6 个 expert，每个 expert 产生独立的中间激活，128 expert 模型的激活值显存可以是 dense 模型的数倍。\nMegatron-Core 的 activation recomputation 有两种模式：\n1 2 3 4 5 6 7 8 9 10 # Selective recomputation（推荐）：只重计算特定算子 cfg.model.recompute_granularity = \u0026#34;selective\u0026#34; # 重计算 attention 的 softmax 和 dropout，保留 GEMM 结果 # MoE 场景下：重计算 router softmax、dispatch permute # Full recomputation：整层重计算 cfg.model.recompute_granularity = \u0026#34;full\u0026#34; cfg.model.recompute_method = \u0026#34;uniform\u0026#34; cfg.model.recompute_num_layers = 1 # 每 1 层做一次完整重计算，显存节省最大但计算增加 ~33% DeepSeek V3 的 32 节点配置就使用了 full recomputation——在 2048 张 GPU 上，显存是更紧张的约束（EP=64 意味着每 GPU 仍持有 4 个 expert 的完整参数和优化器状态），而额外的计算开销可以被大量 GPU 分摊。\n端到端实战 前面介绍了架构、Bridge、加速原语和精度策略。把它们组合起来，端到端训练一个 MoE 模型是什么样的？以 Nemotron-3-Nano (30B-A3B) 为例。\nRecipe 配置解读 Megatron-Bridge 通过 Python recipe 文件定义完整的训练配置。以下是 Nemotron-3-Nano 预训练 recipe 的关键参数逐项注释：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 from megatron.bridge import AutoBridge from megatron.bridge.training import TrainingConfig, MixedPrecisionConfig def nemotron_3_nano_pretrain_config(): cfg = TrainingConfig() # ─────────── 模型加载 ─────────── cfg.model = AutoBridge.from_hf_pretrained( \u0026#34;nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16\u0026#34; ).to_megatron_provider(load_weights=True) # ─────────── 并行策略 ─────────── cfg.model.tensor_model_parallel_size = 4 # TP=4, 节点内 NVLink cfg.model.pipeline_model_parallel_size = 1 # 无 PP（模型够小） cfg.model.expert_model_parallel_size = 8 # 128 experts / 8 = 16 per GPU cfg.model.sequence_parallel = True # 与 TP 配合的序列并行 # ─────────── MoE 加速 ─────────── cfg.model.moe_grouped_gemm = True # Grouped GEMM（必开） cfg.model.moe_permute_fusion = True # Permute 融合 cfg.model.moe_token_dispatcher_type = \u0026#34;alltoall\u0026#34; # 标准 All-to-All cfg.model.moe_shared_expert_overlap = True # Shared expert 并行 # ─────────── MoE 路由 ─────────── cfg.model.moe_router_score_function = \u0026#34;sigmoid\u0026#34; cfg.model.moe_router_enable_expert_bias = True # DeepSeek-style bias cfg.model.moe_aux_loss_coeff = 0.0001 # 负载均衡 loss # ─────────── 计算融合 ─────────── cfg.model.cross_entropy_loss_fusion = True cfg.model.gradient_accumulation_fusion = True cfg.model.bias_activation_fusion = True cfg.model.use_fused_weighted_squared_relu = True # ─────────── 训练超参 ─────────── cfg.training.max_steps = 100000 cfg.training.global_batch_size = 256 cfg.training.micro_batch_size = 1 cfg.training.seq_length = 8192 # ─────────── 优化器 ─────────── cfg.optimizer.lr = 1e-4 cfg.optimizer.min_lr = 1e-5 cfg.optimizer.weight_decay = 0.1 cfg.optimizer.adam_beta1 = 0.9 cfg.optimizer.adam_beta2 = 0.95 # ─────────── 混合精度 ─────────── cfg.mixed_precision = MixedPrecisionConfig() # 默认 BF16 return cfg 几个关键决策的解读：\n为什么 PP=1？ Nemotron-3-Nano 只有 30B 参数，TP=4 + EP=8 已经能让每张 GPU 的显存需求可控。PP 会引入流水线 bubble，对吞吐量有损，小模型没必要。\n为什么 micro_batch_size=1？ MoE 的 All-to-All 通信量与 micro-batch 内的 token 数成正比。MBS=1 意味着每次前向/反向只处理 8192 个 token，限制了 All-to-All 的峰值通信量。全局 batch size 通过 gradient accumulation 扩大到 256。\n为什么 moe_aux_loss_coeff=0.0001？ 负载均衡 loss 鼓励 token 均匀分布到各 expert，但系数过大会干扰主任务 loss。0.0001 是经验值——足够防止 expert 坍缩（所有 token 只选少数 expert），又不至于损害模型质量。\n启动命令 1 2 3 4 5 6 7 8 9 10 11 12 13 # 4 节点 × 8 GPU = 32 GPU # 使用 torchrun 启动，Megatron-Bridge 自动初始化并行组 torchrun --nproc_per_node=8 --nnodes=4 \\ --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \\ -m megatron.bridge.training.train \\ --recipe nemotron_3_nano_pretrain_config # 环境要求： # - PyTorch \u0026gt;= 2.4 with CUDA 12.x # - Megatron-Core (via 3rdparty/Megatron-LM submodule) # - Transformer Engine (TE) \u0026gt;= 1.7 # - NCCL \u0026gt;= 2.18 # - 节点间 InfiniBand / RoCE RDMA 互联 SFT 场景差异 SFT（Supervised Fine-Tuning）与预训练在配置上有几个关键差异：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 # 并行策略简化：SFT 通常数据量较小，可以降低并行度 cfg.model.tensor_model_parallel_size = 1 # TP=1（不切分权重） cfg.model.expert_model_parallel_size = 8 # EP 保持（expert 数不变） # 学习率降低 cfg.optimizer.lr = 2e-5 # 预训练的 1/5 cfg.optimizer.min_lr = 2e-6 # 序列长度可能更长（SFT 数据含完整对话） cfg.training.seq_length = 16384 # PEFT：LoRA/DoRA 只微调部分参数 cfg.peft.method = \u0026#34;lora\u0026#34; cfg.peft.lora_rank = 16 cfg.peft.lora_alpha = 32 cfg.peft.target_modules = [ \u0026#34;linear_qkv\u0026#34;, # Attention QKV \u0026#34;linear_proj\u0026#34;, # Attention output \u0026#34;linear_fc1\u0026#34;, # MLP (含 expert) \u0026#34;linear_fc2\u0026#34;, # MLP down \u0026#34;in_proj\u0026#34;, # Mamba in projection \u0026#34;out_proj\u0026#34;, # Mamba out projection ] 注意：LoRA 的 target_modules 覆盖了 Transformer 和 Mamba 两种层类型的关键投影——这是 Hybrid 架构特有的需求。只微调 Transformer 部分的 LoRA 会丢失 Mamba 层的适配能力。\n常见调参陷阱 1. EP 与 DP 的平衡：EP 越大，每 GPU 持有的 expert 越少，显存越宽裕。但 EP 会\u0026quot;吃掉\u0026quot;DP 维度——如果 EP 过大导致 DP=1，全局 batch size 只能通过 gradient accumulation 实现，训练吞吐量可能下降（因为 GA 的 micro-step 之间没有计算-通信重叠）。\n2. Grouped GEMM 必须开启：不开 moe_grouped_gemm，128 个 expert 意味着每 GPU 16 次独立 GEMM，kernel launch overhead 会吃掉大量 GPU 时间。这是 MoE 训练最容易忽略的性能杀手。\n3. Router score function 要与预训练一致：如果预训练用的是 sigmoid routing，SFT 也必须用 sigmoid。切换为 softmax 会导致 token 分布突变，训练不稳定。\n4. FP8 + MoE 的陷阱：FP8 训练时，MoE router 的输出需要对齐到 FP8 kernel 的要求。忘记设置 moe_router_padding_for_fp8=True 可能导致 CUDA kernel crash 或静默精度问题。\n5. Shared expert overlap 与 Flex Dispatcher 互斥：不能同时开启 moe_shared_expert_overlap=True 和 moe_token_dispatcher_type=\u0026quot;flex\u0026quot;。前者用 CUDA multi-stream 并行 shared/routed expert，后者也使用 multi-stream 做异步 dispatch，两者会产生 stream 竞争。\n如何进一步加速？ MoE 训练的性能优化是一个不断演进的领域。以下是当前最有潜力的几个方向。\n通信瓶颈分析 在做任何优化之前，先定位瓶颈。Megatron-Core 集成了 PyTorch 的 profiling 工具：\n1 2 3 4 5 6 7 8 9 10 11 12 # 使用 PyTorch Profiler 抓取 trace with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], schedule=torch.profiler.schedule(wait=5, warmup=2, active=3), on_trace_ready=torch.profiler.tensorboard_trace_handler(\u0026#39;./log/trace\u0026#39;), ) as prof: for step, batch in enumerate(dataloader): train_step(model, batch) prof.step() 在 TensorBoard 中查看 trace，重点关注：\nAll-to-All 占比：如果 All-to-All dispatch + combine 占单步时间超过 30%，说明 EP 通信是瓶颈，可以考虑 Flex Dispatcher 或增大 EP（减少每次 All-to-All 的通信量） GEMM 利用率：如果 expert GEMM 的 GPU 利用率低于 50%，说明每个 expert 处理的 token 太少，应该增大 micro-batch size 或减小 EP 通信-计算重叠率：如果 TP comm 和 GEMM 没有重叠（在 trace 中表现为通信和计算交替出现），需要开启 comm overlap Kernel Fusion 前沿 FlashAttention-3：针对 Hopper 架构的 warp-specialization 优化，进一步提升 attention 计算的 FLOPs 利用率。对 Nemotron-3-Nano 的 7 个 attention 层虽然影响有限，但对纯 Transformer MoE（如 DeepSeek V3 的 61 层 attention）效果显著。\nMLA 专用 Kernel：DeepSeek V3 使用的 Multi-Latent Attention 将 KV 投影到低维空间，减少 KV Cache 显存。针对 MLA 的专用 CUDA kernel 可以将 latent 压缩/解压与 attention 计算融合，避免额外的内存读写。\n通信-计算极致重叠 当前的 Shared Expert Overlap 和 Flex Dispatcher 各自做局部的通信-计算重叠，但它们之间是互斥的。未来的方向是全局调度：\ngantt title 理想的 MoE 层执行流水线 dateFormat X axisFormat %s section Pipeline Router :r, 0, 3 Dispatch (异步) :d, 3, 10 Shared Exp (与 dispatch 并行) :s, 10, 18 Routed Exp (dispatch 完成后) :re, 18, 28 Combine (异步) :c, 28, 35 Next Layer :nl, 35, 39 EP dispatch 与上一层的 dense 计算流水线化，shared expert 与 dispatch 并行，combine 与下一层的 router 重叠——将空闲 bubble 压缩到极致。\n新硬件适配 GB200 NVL72：72 个 GPU 通过第五代 NVSwitch 全互联，对分带宽 1.8 TB/s。这个拓扑对 MoE 的 All-to-All 是革命性的——NVL72 内的 All-to-All 带宽比当前 8-GPU NVLink 域提升 ~9 倍。HybridEP 已经为这个拓扑做了优化，区分域内（NVSwitch 直连）和域间（跨机架）的通信路径。\nMXFP8 / FP4：Blackwell 架构引入了 Microscaling FP8（更细粒度的缩放因子）和 FP4（4-bit 浮点）。FP4 可以将 expert 权重的显存进一步减半，让单 GPU 能容纳更多 expert——这可能改变 EP 的最优配置，因为 EP 可以更小（更多 expert per GPU），从而释放更多 GPU 给 DP。\n系统级优化 异步 Checkpoint：大规模训练中，checkpoint 保存可以花费数分钟。异步 checkpoint 在后台线程完成存储 I/O，训练不中断。对于 MoE 模型尤其重要——128 expert 的全量 checkpoint 可能达数十 GB，同步保存会严重拖慢训练。\nFault Tolerance：2048 张 GPU 的集群中，单卡故障是常态而非异常。弹性训练框架需要支持：检测故障 GPU → 重新分配并行组 → 从最近 checkpoint 恢复 → 继续训练。Megatron-Core 目前通过与 NVIDIA NeMo 的集成支持部分故障恢复能力。\nKey Takeaways MoE 训练的三大挑战：显存墙（128 expert 参数量巨大）、通信墙（All-to-All dispatch 开销与 expert 数成正比）、负载不均衡（token routing 天然不均匀）。这三者相互耦合，需要联合优化。\nMegatron-Core 的五维并行：TP/PP/DP/EP/CP 构成完整的并行空间。EP 和 DP 竞争同一组 GPU——EP 越大，DP 越小。DeepSeek V3 用 TP=2, PP=16, EP=64 训练 671B 模型，需要 2048 张 GPU。\nMegatron-Bridge 消除了生态锁定：AutoBridge 自动检测 HF 模型类型，QKVMapping 和 GatedMLPMapping 处理权重格式差异，训练完成后可无缝导出回 HF 格式用于推理部署。\nMoE 加速的四个层次：\n并行策略：EP 分配 expert，与 TP/PP 组合 通信调度：Flex Dispatcher（DeepEP/HybridEP）实现异步 All-to-All 计算融合：Grouped GEMM 将 N 个 expert 合并为一次 kernel 调用 精度优化：FP8 将 GEMM 算力翻倍、通信量减半 Nemotron-3-Nano 是 Hybrid 架构的代表：52 层中混合 Mamba（SSM）、MoE、Attention 三种层类型，用少量 Attention 捕获全局依赖，Mamba 处理局部上下文，MoE 提供大参数容量。\n实战中的关键配置：moe_grouped_gemm=True 是必开项；moe_shared_expert_overlap 与 moe_flex_dispatcher_backend 互斥；FP8 训练需要 moe_router_padding_for_fp8；SFT 的 router score function 必须与预训练一致。\n推荐源码阅读路径 如果你想深入 Megatron-Core 和 Bridge 的源码，推荐以下阅读顺序：\nmegatron/core/parallel_state.py — 理解五维并行组的初始化逻辑 megatron/core/transformer/moe/moe_layer.py — MoE 层的完整工作流 megatron/core/transformer/moe/token_dispatcher.py — All-to-All dispatch/combine 实现 megatron/core/transformer/moe/grouped_mlp.py — Grouped GEMM 的权重布局和计算 megatron/bridge/models/conversion/auto_bridge.py — AutoBridge 自动检测机制 megatron/bridge/models/conversion/param_mapping.py — 权重映射基类和 QKV/GatedMLP 映射 megatron/bridge/training/comm_overlap.py — TP 通信重叠的三种策略配置 参考资料 Shoeybi et al., 2019. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Megatron 的开创性工作，定义了 TP 列切分/行切分范式\nNVIDIA, 2025. Megatron-Bridge GitHub Repository — Megatron-Bridge 开源仓库\nDeepSeek-AI, 2024. DeepSeek-V3 Technical Report — 671B MoE 模型的训练细节，包括 EP=64 的并行配置\nFedus et al., 2022. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity — MoE 在 Transformer 中的经典应用\nDai et al., 2024. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models — Fine-grained expert + shared expert 的设计\nNVIDIA, 2025. Nemotron-3-Nano Technical Blog — Hybrid Mamba + Transformer + MoE 架构解析\nGu \u0026amp; Dao, 2023. Mamba: Linear-Time Sequence Modeling with Selective State Spaces — Mamba SSM 架构，Nemotron-3-Nano 的层类型之一\n","permalink":"https://mzf666.github.io/llm-infra/zh/posts/05-megatron-bridge/","summary":"从 Megatron-Core 架构到 Megatron-Bridge 桥接机制，深入剖析 MoE 大模型训练中的并行策略、通信优化与计算融合。","title":"Megatron-Bridge 深度剖析：MoE 大模型训练加速的工程之道"}]