DeepSeek V4 Infra Deep Dive 4.4: PyTorch 训练显存管理

来源：DeepSeekV4 Infra Deep Dive · 第四节 4.4: PyTorch 训练显存管理：从 Caching Allocator 到 Stream 同步

系列：可灵 AI infra 训练团队出品，全文系列见原文末链接。本文是该系列的 4.4 一篇，专门讲训练显存怎么分、怎么 OOM、怎么省。

笔记以”按生命周期分四类”为主线展开。比按”哪一项占得多”分类更能看清 OOM 卡在哪一环：从初始化到 step 结束，四类显存分别在什么时刻驻留、用什么手段能挪走，挪到哪里去。

01 allocated 与 reserved 的语义

这一段是后面所有判断的前提。

• allocated = caching allocator 实际切给存活 tensor 的字节数
• reserved = allocator 从驱动一次性申请到的 segment 总量；驱动一次分一大块连续显存，叫一个 segment

两者永远满足 reserved ≥ allocated，差出来的是 allocator 持有的、暂时既不还驱动也没分出去的空闲 block，也就是碎片。

reserved = allocated + 碎片 + 瞬态峰值被抬高后留下的「台阶」

「台阶」是 OOM 判断里最容易踩的坑：caching allocator 的 free 只把 block 退回池里复用、并不向驱动退还，整段 segment 全部空闲才有资格归还。瞬态峰值那几块大 block 常常和长寿命 block（激活缓存、通信预留）位于同一 segment，始终无法同时空出来，于是 reserved 只升难降——它记录的不是此刻用了多少，而是训练至今某一瞬间最多需要过多少连续显存。

「allocated 不高却 OOM」通常就是这道 reserved 台阶造成的。

02 四类训练显存：按生命周期划分

训练显存的增量主要来自反向传播。推理时单卡 HBM 上基本只有 weights，加一份随生成长度增长的 KV cache；训练多出的十几倍几乎都来自反向——chain rule 要求前向中间结果和梯度副本一直驻留到反向用完，于是 activation、梯度、optimizer state、重算中间值才同时占用卡上 HBM。下面四类是同一个反向需求在不同生命周期上的分布，去掉反向（即推理）它们全部消失。

按生命周期划分，关注每一类显存什么时候分配、什么时候释放、有没有可能跨 step 还被保留。四类按占用从高到低排列：当前最高的那类降下来后，下一类就成为主导项。

下面逐类拆开。

02.1 恒定占用：weights / gradients / optimizer states

生命周期跨 step 不变：从初始化那一刻分配出来，持续保留到训练结束。每一步前向 / 反向 / optimizer 都要读它们，没有任何释放窗口。

体现在 reserved 曲线上：训练启动后几个 step 内，reserved 分几次跳到一个稳定值，此后整段训练再不变。

必须提前一次性分配。DDP 下模型刚 build 出来就占满。FSDP / ZeRO 是特例，按 rank 分片到多卡。除了沿 DP 副本维度分片的 FSDP / ZeRO，TP（按 head / hidden 划分）、PP（按层划分）、EP（按 expert 划分）也都在削减这条恒定基线，只是各占一个不同的维度。

02.2 Activation：前向累积、反向释放

生命周期由 chain rule 决定：前向算到第 i 层时产生的中间结果，要保留到反向传播回第 i 层、读取完成后才释放。中间隔着「forward 跑完剩下层 + backward 倒回这一层」几十毫秒到秒级的窗口——跨 kernel 却不跨 step。

allocated 曲线上是随 forward 单调累积、随 backward 单调消费的单峰曲线——前向到达峰值时全部 activation 都保留在 HBM 上，backward 每过一层释放一份。这段累积-释放曲线是 Activation Checkpointing 的优化目标。

每个 micro-batch 临时申请：每一步开始前 activation 都不存在；前向算到哪层，那层中间结果就被 allocator 划出 block；分配路径仍留在 caching allocator 内部。

上一个 micro-batch 释放的同档 block 留在池里复用。warmup 之后几乎所有 activation 分配都在 allocator 内部命中已有 block，驱动那条较慢的路径只在分桶不匹配（请求大小不落进任一现成 block 的尺寸区间）或 segment 已用尽时才走。

02.3 单个 kernel 内的瞬态峰值

进入 kernel 时申请、退出 kernel 立刻释放，存在时间只有毫秒甚至微秒级。MoE 的 expert FFN intermediate、FlashAttention backward 的 work tensor、cuBLAS GEMM 的 workspace 都属于这一类。

在 MemTrace 上表现为一道 kernel 刻度宽的尖峰，按时间平均下来几乎看不出来，瞬时却常常达到单层 GiB 级。因为存在时间太短、不进入前向保存的激活，手算 saved activation 时容易漏掉。

常驻还是临时：看上去是临时的，kernel 跑完就标记为空闲。但 caching allocator 的 free 只把 block 退回池里复用、并不向驱动退还，整段 segment 全部空闲才有资格归还；瞬态峰值那几块大 block 常常和长寿命 block 位于同一 segment，始终无法同时空出来，于是 reserved 只升难降。

allocated 曲线规律地起落，reserved 却是台阶式抬升。

这类瞬时峰值的根本解决办法是 fused kernel，把多个 kernel 的中间结果保留在 SRAM、不写回 HBM，单次数据量不会产生几 GiB 的瞬态 buffer。

02.4 通信 buffer：集合通信的中转缓冲

训练一开始就预分配，此后一直占用。NCCL 按上限一次性预留每条通道的发送 / 接收缓冲、ring 上的中转 buffer、reduce 算法的工作区。

尺寸由四个因素共同决定：message size、rank 数、NCCL 选定的算法（ring / tree / NVLS）、bucket 聚合。

若是等训练主循环跑起来之后才临时申请，cudaMalloc 会阻塞发起它的 CPU 线程，并触发 device-wide 同步：整卡所有 stream 已入队的工作要先排空，卡在集合通信的关键路径上。中后期 reserved 已接近上限，驱动这一刻可能分配不出连续的几百 MiB，等到通信发不出去才发现显存不够，训练 hang 在集合通信上。预分配在启动时一次性解决这些不确定性，跑起来之后通信 buffer 占用恒定不变，不再与主计算争用驱动。

02.5 通信 stream：overlap 与显存池分裂的取舍

同一条 stream 上的 kernel 严格按提交顺序串行执行，通信和计算共用同一条 stream 时，allreduce 一发就阻塞后面的 GEMM，重叠执行无从谈起。

为了 overlap，NCCL 的集合通信通常跑在独立的通信 stream 上：集合通信 kernel 独立排队执行，主计算 stream 同时还能发 GEMM。这不是强制要求——集合通信留在 default stream 一样能跑通，只是退化成串行，但 buffer 进入主池反倒不分裂、碎片更少。只有真正存在数据依赖的关键点才用 wait_event 显式同步，比如 backward 算到 dW 那一刻等待上一步 allreduce 的 event。

副作用是显存池跟着分裂。caching allocator 按「一个池只服务一条 stream」的规则过滤空闲 block（隔离用于保证正确性：A stream 的 kernel 可能还没读完这块 block，B stream 拿去写会发生数据竞争）。主计算 stream 的池占满之后，通信 stream 名下整段 segment 即便全部空闲也无法被主计算 stream 复用，OOM 由此产生。

反过来，通信 stream 上要是没多少计算 / 通信可重叠，或显存已经紧张到池分裂的额外开销划不来，就不要分出去——让通信留在主 stream，用一点串行换回不分裂的显存池更划算。

03 四类合计：DeepSeek V3 建模

四类合计对 V4 同样成立。下表 per-GPU 数字取自 NVIDIA Megatron Core MoE 报告对 V3 配置的建模（报告设 256 GPU、PP4×VPP4×EP64，对标 H100 80 GB 容量）：

表中通信 buffer 只计 EP all-to-all + DP allreduce 预留；再加上 PP send/recv（点对点，非集合通信）预留，整体 ≈ 3 GB。

第一行给出恒定基线 68.5 GB，只剩 11.5 GB 供动态显存使用。第二行 Activation 131 GB 是这点余量的十倍有余——recompute、FP8 storage、offload 都是要把它压缩到 11.5 GB 以内；压缩不下去，就只能继续划分恒定基线，腾出更多余量。

未优化前合计 199.5 GB ≈ 2.5× H100。目标是把 per-GPU 总占用压回 80 GB；131 GB 的 activation 是首要压缩目标。

03.1 真实峰值不是各类最大值的简单相加

四类峰并不在同一时刻。判断某套配置会不会超出容量，要逐时刻看占用：

• 恒定占用与通信 buffer 是两条水平线
• activation 是一个单峰
• 计算峰值是一个尖峰

某一刻的占用 = 那一刻各类的瞬时值相加；真峰值 = 这个逐时刻之和在整个 step 里取到的最大值。峰值时刻 \(t^*\) 会随配置变化：

• 不开 recompute 时，activation 单峰停在 forward 末端，\(t^*\) 位于那里。
• 开 full recompute 后，forward 末端那段保存的激活被削平，但 backward 早期出现一个新窗口：层间保存的激活还没卸掉几层、当前层 recompute 重算出的 forward intermediate、当前层新分配的 grad buffer 三类在同一窗口叠加，\(t^*\) 从 forward 末端移到 backward 早期。

这只是把它推迟到反向、并非删除。峰值是否下降取决于移入窗口的高度：低于被削平那段就降低（V3 full recompute 即把峰值从 199.5 GB 压到 ~85 GB），要是比削平那段还高，反而会更早 OOM。

04 MemTrace 与 OOM 诊断

04.1 怎么打 snapshot

torch.cuda.memory._record_memory_history(max_entries=100000)
# 跑几个 step
torch.cuda.memory._dump_snapshot(path)

得到的 pickle 文件里，每一次 alloc / free 都带着时间戳、字节数、所在 stream 和 Python stack trace。

把 allocated 和 reserved 按时间画出来，是 MemTrace 时间线；取 allocated 最高那一刻各类的瞬时值、按 tensor 归类排列，是峰值占用分解。

04.2 案例：peak ≈ 93.5 GB（部分 recompute 后）

下面这张 MemTrace 时间线，浅色区域是 allocated，深色线是 reserved；峰值落在 backward 早期 recompute 窗口，OPT 段另有约 10 GB 的临时 m / v 副本尖峰。

下面这张是同一时刻峰值截面按 tensor 来源分解，合计约 93.5 GB，高出 H100 80 GB 上限约 13.5 GB。

同一时刻的占用分解里，六项是：

这六类与开头四类并非一一对应：Activation 在峰值处拆成层间保存的激活和 recompute intermediate（计算峰值瞬态归入后者的重算窗口），另起 grad buffer 与 allocator overhead；恒定占用与通信 buffer 在占用分解里保持不变。

04.3 OOM 的两种典型情况

情况 A：reserved 与 allocated 都接近上限。 空间已经用尽，瓶颈在峰值时刻各类的同时占用量。看峰值时刻哪一类（恒定基线、层间激活缓存、瞬态、还是 allocator 空隙）的量级最大，针对最大的那类优化。开 recompute 才有意义，它压缩的层间激活缓存正是触及上限的主因。

情况 B：reserved 已达上限，allocated 却低出好几 GB。 可用字节并没有真的被填满，实际是 reserved 被虚高抬起。抬高 reserved 的三种成因都不靠多用显存：

• 瞬态峰值被抬进 reserved 后不再回落形成的那道台阶
• shape 变化造成的碎片
• 多 stream 分池后一条 stream 的空闲填不进另一条的请求

再开 recompute 去压缩 allocated 没有收益——压缩的那一维根本不是瓶颈。

OOM 判断要同时看峰值占用分解（时间维：某一刻同时被持有多少）和 allocator 碎片状态（空间维：池能不能凑出这一刻要的连续大块）。低于上限的峰值遇到碎片化的池照样 OOM；等于上限的峰值若 shape 固定、池无碎片，反而能分配成功。

04.4 失效条件

一次 snapshot 只是某个 step、某张 rank 的快照。MoE 路由逐 step 波动，负载最重的那张 rank 也不固定在同一张卡上——这一回没碰到峰值，不能推导出下一个 step、另一张卡也安全。要做出判断，得在几个高负载 step、几张可疑 rank 上各采样分析一次，看差值是稳定接近上限还是偶发接近上限。

04.5 各分量能否静态估计

各分量字节量可以静态估到 ± 数 GB；真正无法静态预测的只有两样：

• 峰值落点 \(t^*\)（取决于运行时 stream 交错、各 kernel 时长、MoE 路由逐 step 波动）
• allocator 碎片余量

都只能事后从 snapshot 里分析出来。

05 优化手段：四类各自的取舍

05.1 恒定占用：分片、降精度、offload

V3 配置下，恒定占用这三块不分片的话，单 rank 要装：

\[ m_{const} = \underbrace{2N}_{\text{BF16 weights}} + \underbrace{2N}_{\text{BF16 gradients}} + \underbrace{(4+4)N}_{\text{Adam FP32 m/v}} \approx 218 \text{ GB} \]

若加 FP32 master weight（4 bytes/param），完整是 16 bytes/param ≈ 291 GB。前面表里的 68.5 GB 是 Megatron 报告口径下按 V3 并行度切分后的 per-GPU 值。

分片

四种切法各对应一个维度，互不争用同一维度，可以相乘：

\[ \text{per-GPU 恒定占用} = \frac{\text{全量}}{\prod \text{各相关并行维}} \]

降精度

• FP8 weight：把 forward / backward 实际计算、并参与 DP all-gather 的那份工作权重从 BF16 降到 FP8，单 weights 一项再省 2×；FP32 master weight 仍保持高精度。改变的是 forward / backward 的计算精度与 all-gather 的通信精度，optimizer 更新仍在 FP32 计算。
• FP8 optimizer state：把 Adam 的 m / v 也降到 FP8。

失效条件：master weight、m / v 这类「极小增量反复累加进一个大基数」的量，仅靠 scaling 还不够。每步要加的更新值比基数小几个数量级，量化后落在基数尾数能分辨的台阶之下，加法被整个舍去，即 swamping（大基数把小增量淹没）。表现为参数实际不更新、loss 曲线无报错却不再下降。所以这类累加量必须保留一份高精度累加器。

外移

• Optimizer state offload → CPU：通过 PCIe Gen5（≈ 64 GB/s 单向）移到 CPU 内存
• Weights → NVMe：再往后退一档到 NVMe（PCIe + filesystem，≈ 7 GB/s 单向）

从 NVMe 把 weights 拉回 GPU 计算一层比 HBM 直接读取慢两个数量级，MFU 会大幅下降，所以只在「单卡资源严重受限、必须容纳大模型」的场景才用。

降精度只在划分完之后单项仍占几 GB 以上时才划算。FSDP 划分到 256 rank 后 per-GPU optimizer state 已经只有 0.57 GB，再降到 FP8 也只节省约 0.4 GB，不值得为这点字节数增加复杂度。

05.2 Activation 缓存的压缩手段

Activation 缓存是峰值的大头，V3 配置下未压缩前每卡就有 131 GB（约 1.6× 一张 H100）。

Recompute（时间换空间）

Offload（搬到 CPU / NVMe）

• Activation → CPU offload：走 PCIe Gen5 ≈ 64 GB/s 单向
• Activation → NVMe offload：再慢 10× 以上，预训练中几乎不用

关键：必须把 PCIe DMA 和 forward / backward 计算 overlap：移下去要赶在下一层 forward 还没需要缓存之前完成，移回来要在反向使用缓存之前预取完。两个方向都得对准时序，任一边没与计算重叠 MFU 就下降——recompute 只在反向单点插入一段重算，offload 要双向与计算重叠，调度余量小得多。

序列并行（长上下文）

按 sequence 维切到多 rank 上。序列越长、S 相关的那几项越主导，把 S 划分到多 rank 是唯一随长度同步缩小数据量的方式。

FP8 storage

把前向保存下来的中间值每元素的字节数从 16 bit 降到 8 bit。activation 是一次性写、一次性读的量，配 scaling 和反量化即可。

Fused kernel（中间量留 SRAM）

SRAM 和 HBM 的容量、带宽都差一个数量级，中间量每在 HBM 上往返一趟就消耗一份带宽。fuse 把数据量沿某个维度划分成小到能整个放进 SRAM 的 tile，逐 tile 算完只把最终结果写回 HBM，中间 tensor 始终留在片上、既不占 HBM 也不付往返带宽。融合的边界由 SRAM 预算限定。

判据不在中间量是否瞬时，而看 consumer 要的每个输出 tile，是不是只依赖产出它的那一个 CTA（一个 thread block）：

• 依赖都在本 CTA 就留片上：逐元素 pointwise，或归约轴整条留在 CTA 内 streaming
• 要跨 CTA 读取就必须物化——在 producer 已切散到多 CTA 的那根轴上做 reduction / transpose / shuffle

归约本身不必然导致物化：只要归约轴留在 CTA 内 streaming（FA 的 KV、norm 的 \(D\)，running state 小到进寄存器），归约也能片上；物化只在归约轴被 tile 到多个 CTA 时才出现。

代表实例：

• FlashAttention 系：把 QK^T、softmax 输出 P、dS、dP 全部留在 SRAM tile 内。HBM 上只剩 O ∈ [B,S,D] 和 LSE ∈ [B,H,S]。参考参数 (\(B=4, H=32, S=8192, D=4096\), BF16) 下，原本几块约 40 GiB，FA 后只剩 0.254 GiB
• fused LayerNorm / RMSNorm + residual：归约轴 \(D\) 整条留在 CTA 内
• fused MoE token permute：路由后的 token 重排融进 grouped-GEMM 的数据准备路径（Megatron 开关 moe_permute_fusion）
• fused cross-entropy（长 vocab）：\(V \ge 128k\) 时把 logits 和 NLL loss 合并到一个 kernel 内，按 row 流式计算

05.3 计算峰值瞬态的压缩

瞬态是 kernel 执行中临时占用的 HBM，规模 = f(B, 每卡 token 数, dtype 位宽)。recompute 和 offload 对瞬态均无效：瞬态是 kernel 内 μs 级的瞬时量，recompute 针对的是跨 kernel 存下来给反向用的 saved activation，削减它反而新开一个重算窗口、又产生一份瞬态；offload 走 PCIe，搬运延迟比 μs 慢几个数量级，根本来不及在 kernel 算完前把数据量移出。

走 fused kernel / tiling 这条路：不去缩小数据量的尺寸，而是让中间量根本不写进 HBM。

fused MoE expert FFN 是个典型反例：把 expert 内部的 up_proj + gate_proj + 激活（SwiGLU）+ down_proj 几步在 expert 内 fuse，能不能留片上取决于 kernel 怎么 tiling，而不取决于「融没融」。标准实现（Megatron-Core grouped GEMM、V3 同理）下这份 [C_local, 2n] 实际会落 HBM：FC1 把输出按 N 维切散到多个 CTA，FC2 又要沿同一维（FC1 的输出维 F）跨 CTA 归约。这正是「跨 CTA 读取就得物化」那种情形，只能过 HBM/L2。即便把 FC1 + 激活 + FC2 强融进一个 megakernel 也没用——一个 kernel 不等于一个 CTA，各 stage 间仍要在 global memory 上会合。

要真留片上，得换成每个 CTA 独占 token 行块、自己做完整 \(F\) 归约的 FA 式 tiling，代价是降低 GEMM 并行度 / 吞吐。V3 训练路径同理走 grouped GEMM、这份激活照样落 HBM：按 V3 配置 (\(B=1, S=4096\)) 静态估算，\([C_local, 2n]\) 单层约 0.3 GiB / 卡，不是单层峰值主导项；主峰是随路由数 K 放大的进出口张量（\([C_local, D]\) 量级、随 K 累加）。

cuBLAS workspace shrink：CUBLAS_WORKSPACE_CONFIG=:4096:8 固定 workspace 上限（多 stream 共享 handle 时换取确定性），可省下几十个 handle 累积的几 GiB；代价是某些 SplitK / 多阶段 GEMM 算法用不了，吞吐可能略降。

05.4 通信 buffer：减少 staging 份数与尺寸

NCCL 在每个 GPU 上为每条通信 stream 注册一段 staging buffer，尺寸按最坏情况一次预分配。V3 配置下每卡每方向 EP all-to-all staging 约百 MiB 级（某配置下估算约 230 MiB）。

调度（bucket 聚合）

把同一层多个参数的 gradient 合并成一次 collective。bucket 小则小消息多、staging 总量大、带宽利用率低；bucket 大则单次 staging 大、还要等更多 grad 准备好才发，反向关键路径等待时间长。

配置（NCCL 环境变量）

• NCCL_BUFFSIZE × channel 数 = 每个 rank 预留的总 staging（Simple 默认 4 MiB/channel）
• NCCL_NTHREADS：每 channel 线程数，决定 staging 复用速度
• NCCL_ALGO / NCCL_PROTO：选 ring / tree / collnet / nvls 算法与协议（LL 有效带宽 50% / LL128 95% / Simple 最高延迟也最高）

算法 / 拓扑

复用（对称内存直接写远端）

• NVSHMEM：rank A 可直接 nvshmem_put 把数据写到 rank B 显存上，省去本地 + 远端 staging
• DeepEP：早期 / legacy 版本用 NVSHMEM 省去 EP all-to-all 两侧的中转拷贝（某配置下每卡每方向 ≈ 230 MiB）

05.5 对照表

「优先级」一列表示当前配置下的收益/代价排序：先动富余资源，后动紧缺资源。它会随配置变：通信已是反向关键路径瓶颈时再做恒定占用的拆分反而不划算。

05.6 各手段的失效条件

• FSDP / ZeRO：sharding 组扩大到 all-gather 慢于反向计算时就不再划算——通信从「藏在计算后面」变成关键路径瓶颈
• FP8 storage：失效在反复累加那一类量（swamping），只有一次性读写的激活才换得安全
• Offload：搬运时间藏不进计算就反过来拖慢，PCIe 来回的延迟进入关键路径
• fused kernel：标准 grouped-GEMM 实现下反向会失效——反向要把它整张 fan-out 给多个梯度 GEMM，forward 那种「逐 tile 流过、从不整张落地」的办法用不上
• NVSHMEM：用单边写替换掉随跳数增长的 staging，但同步和完成语义要自己管

06 CP/SP/PP/TP 各分片哪一段显存

06.1 想降低恒定 68.5 GB：用 FSDP / ZeRO

每加一份 DP 维就 \(/N\)。分到极致的代价是每步都要 all-gather + reduce-scatter，再往下分省下的显存会被新增的通信开销抵消——这时要退回到 HSDP（hybrid sharding，组内分片、组间用 DDP）或限制 sharding 组大小。

PP（按层划分）、EP（按 expert 划分）同样把恒定占用降到这条基线，但它们对激活缓存帮助有限。

06.2 层间激活缓存靠什么分片

设每 step 的 Global Token Size = \(G\)，切 sequence 后每个 rank 持有 \(G/N\) 个 token，attention 数据量总和没变。

固定 Global Batch Size 时，划分序列会让 Local Batch Size 上涨，省下的又被原样涨回：总卡数固定，新拿去划分 sequence 的卡只能从 DP 维分出，DP 度一降，\(global\_batch = DP \times local\_batch\) 里每个 DP rank 摊到的 local batch 就成倍上涨。这种抵消对所有 activation 项都成立：

\[ \text{CP 沿 S 切每项} \times \frac{1}{N}, \quad \text{batch 涨回} \times N \;\Rightarrow\; \text{单卡峰值净变化趋近零} \]

抵消假设涨上来的 local batch 由并发 micro-batch 承担——峰值激活只由同时在飞的 micro-batch 数决定。若改用梯度累积让多出来的 batch 串行跑，峰值激活并不回弹，CP 切出的 \(S/N\) 不被抵消、仍净省。

例外条件：当训练目标是训更长 context、要容纳单卡放不下的 sequence 时，CP 加卡用来容纳更长 \(S\)、不从 DP 分卡，local batch 不回弹——此时相对「不切 sequence 直接训练同样长度」，划分 sequence 才真正压低单卡 activation。

06.3 反向当前层的临时量与 kernel 数据量：TP 或 fused kernel

峰值窗口里的动态量分两种：

• Chain rule 决定的反向临时量：重算出的 forward intermediate、与它同形状的中间量梯度（MLP/MoE 这档宽 \(2n\)），以及 dY、dX、dQ、dK、dV 这类宽 D 的边界梯度——这类不能删，只能用 TP 减少每卡持有的 head / hidden 切片（attention / MLP 数据量同样划分 \(/t\)），代价是每层 4 次 all-reduce、对 NVLink 带宽敏感
• 真正 μs 级的 kernel workspace：靠 fused kernel 或 workspace 配置缩小

06.4 PP 对激活缓存帮助有限

PP 大小 \(p\) 在激活缓存公式里被直接约掉：1F1B 第一阶段同时在飞 \(m_{inflight} = p\) 个 micro-batch（\(m \ge p\) 时）。原因：1F1B 启动时先逐级跑完 \(p\) 个 micro-batch 的前向、把 \(p\) 级流水线填满，才轮到第一个 micro-batch 的反向，同时驻留 \(p\) 份激活缓存，正好与 PP 度 \(p\) 约掉。

Korthikanti et al. 2022 (arXiv:2205.05198 §4.2.3) 给出的最优 schedule，bubble 最小化的代价是内存压力集中在 pipeline 第一个 stage。

Schedule multiplier 讨论只在不开 recompute 时才显著：那时激活缓存是整层中间结果，V3 未压缩前每卡 131 GB，\(1\times\) 与 \(1.25\times\) 能差出三十多 GB。开启 Full Recompute 后，每层只 save 一个 layer-input \(2BSD\)，激活缓存缩到 3.5 GiB（V3 量级），multiplier 差异不到 1 GiB。

06.5 分片维度对照

07 FA 之后 HBM 上的残留

今天主流训练框架默认用 FlashAttention，标准 attention 那个 \([B, H, S, S]\) 不再 materialize 到 HBM。

峰值主要来自 FA 之后还留在 HBM 上的：MoE 层的 dispatch / permute / expert FFN 中间值，以及 Full Recompute 重算窗口里的中间值。

FA 那段的参考参数 (\(B=4, H=32, S=8192, D=4096\), BF16) 是单 EP rank、单层、单 micro-batch 的尺度——真实训练里还要再乘上 PP、EP 和 micro-batch 数。

未启用 fused expert kernel 的配置下，MoE FFN intermediate \([C_local, 2n]\) 会成为主导项之一。启用 fused kernel（把 SwiGLU 融进 FC1 epilogue）后，\(2n\) 宽的 pre-activation \(H\) 在 forward 不再单独落 HBM——逐 tile 在 epilogue 里算完即被消费。但这条收益对反向容易失效：反向算 expert FFN 梯度时，forward intermediate \([C_local, 2n]\) 必须在场——无论是 save-for-backward 存下来的，还是 full recompute 现算出来的——它本来就得写出去。此刻还要同时生成它同形状的梯度 \(dH\)（同样 \([C_local, 2n]\)）。forward 不必把它整张落地（随算随被下一段消费），backward 却要整张物化在 HBM 上 fan-out 给多个梯度 GEMM——失效的根因是这道刚性物化边界，不是「tile 超过 SRAM 容量」的问题。

08 Full Recompute 之后峰值仍达 85 GB：差值来自哪几项

上面 MemTrace 占用分解那 93.5 GB，是部分 recompute（层间激活缓存还残留 12 GB）时截下来的快照；这里改为分析 full recompute 把层间激活缓存降到下界后，峰值还剩多少、能不能回到 80 GB 以内。三个配置档的峰值依次递减：

• 未优化 199.5 GB
• 部分 recompute 削到 93.5 GB
• full recompute 继续降低

按 per-GPU 视角逐项算（EP 路由波动下最重 rank 只会更吃紧），V3 配置（\(PP=16, L=61, S=4096, B=1, D=7168\), BF16）下六类分量：

合计：

\[ 68.5 + 3 + 3.76 + 2 + 3 + 5 \approx 85 \text{ GB} \]

这一组分量可代入估算：

\[ \text{恒定基线} \approx \frac{218 \text{ GB}}{\prod \text{FSDP} \times \text{TP} \times \text{PP} \times \text{EP}} \]

\[ \text{层间激活缓存} \approx 2BSD \cdot \frac{L}{p} \cdot \text{multiplier} \cdot \text{recompute\_factor} \]

其中 multiplier 是 PP 调度系数（1F1B ≈ 1×、VPP ≈ 1.23×），recompute_factor 是该粒度下每层的保留比例（full recompute 取 1，不开则是整层中间结果的倍数）。

若峰值超出容量 5 GB，同一时刻占用最大的是恒定那 68.5 GB，该动的就是它：进一步分片（FSDP / TP / PP / EP）或压精度（FP8 master weight / optimizer state）。层间激活缓存只剩 3.5 GiB，再加大重算收益已经很小。

09 Hopper → Blackwell：PP 依赖随硬件扩容下降

之前讨论的前提是单卡 HBM 容量与 NVLink 域大小固定不变；这两项参数一旦变了，最先受影响的就是 PP。

PP 主要解决两件事：

1. 单卡 HBM 容纳不下整个模型，只能按层把恒定占用分摊到多卡
2. 模型必须跨节点划分时，PP 的跨 stage 通信量比 TP/EP 小得多，能用较低带宽的连接承载——代价是流水线 bubble

单卡 HBM 从 H100 80 GB → H200 141 GB → B200（HGX B200 常见可用 180 GB HBM3e，物理口径 192 GB），一张卡能容纳的层数随之上升，容量约束首先放宽。更关键的是 NVLink 域的扩张：Hopper 一域 8 卡，到 NVL72 一域 72 卡。原来必须跨节点、只能交给 PP 走低带宽连接的那部分 TP/EP，现在能整块放进同一个高带宽 NVLink 域里，低带宽连接的约束也随之消失。这两条约束都放宽后，PP 那份 bubble 代价就不再划算，对 PP 的依赖会逐步下降。

10 一些个人化整理

把全文拆成三句话总结：

1. OOM 看的是「峰值时刻各类同时占用」，不是某一类的最大值；reserved 容易被瞬态峰值抬成台阶只升难降。
2. 降显存最划算的顺序：先用 FSDP / ZeRO 把恒定占用摊薄；再用 recompute / fused kernel 砍 activation 峰值；最后才考虑 FP8 / offload 这些有数值或带宽代价的手段。
3. 2026 年的硬件趋势让 PP 的必要性下降：HBM 增大 + NVLink 域扩张，把跨节点划分的硬约束松开了；未来大概率是 TP / EP / CP 在 NVLink 域内整块摊开、PP 留给极少数真正跨域的场景。

跟 [[UltraScale Playbook 笔记]] 互为补充：那本讲「怎么选并行度」，这篇讲「选完之后峰值落在哪、会不会 OOM」。两份笔记里 FSDP / ZeRO / TP 的描述口径一致，差异主要在本文把 MoE + Recompute 的反向窗口单独拎出来画了峰值时刻分解。

参考

• 原文：DeepSeekV4 Infra Deep Dive · 第四节 4.4: PyTorch 训练显存管理
• 同系列：
  • 第一节: infra 的设计原则
  • 第二节: Hybrid Attention - 问题演进视角
  • 第三节：在缓存命中率就是一切的世界里，PagedAttention 失效后会怎样？
  • 4.1: Deepseek V4 为什么不能再按 module 做 checkpoint？
  • 4.2: Activation Checkpointing 的四档粒度
  • 4.3: Tensor-level AC 的三个工程问题：joint graph、最小子图、storage reuse
• 配套阅读：
  • [[UltraScale Playbook 笔记]]
  • Korthikanti et al. 2022 (arXiv:2205.05198) — Reducing Activation Recomputation in Large Transformer Models
  • Liu et al. 2023 (arXiv:2310.01889) — Ring Attention with Blockwise Transformers
  • Jacobs et al. 2023 (arXiv:2309.14509) — DeepSpeed-Ulysses
  • Fang et al. 2024 (arXiv:2405.07719) — USP: A Unified Sequence Parallelism Approach