nemo-mbridge-perf-activation-recompute

激活重计算

稳定版文档：@docs/training/activation-recomputation.md 卡片：@skills/nemo-mbridge-perf-activation-recompute/card.yaml

激活重计算（激活检查点）以在反向传播期间增加前向计算为代价，减少保留激活所需的内存。合适的检查点边界取决于模型架构、注意力后端、并行方式，以及实际决定每个 rank 峰值的张量。

快速决策指南

1. 确认内存压力来自真实分配，而非分配器碎片。比较每个 rank 上的 max_memory_allocated() 和 max_memory_reserved()。
2. 如果工作负载能够运行，保留一个明确的不重计算对照组。在选择性粒度下，recompute_modules=[] 是有效且有用的对照。
3. 根据架构和观测到的峰值选择第一个边界：
   • 标准注意力： core_attn 通常是首选。如果未融合的注意力会物化分数/概率张量，它的效果最显著。使用 Transformer Engine 融合注意力或 Flash Attention 时，请将它与 [] 对比，因为这些后端本身就会重新物化注意力内部状态。
   • 多潜在注意力（MLA）： 如果扩展后的 Q/K/V 投影占主导，先尝试 mla_up_proj。只有在注意力核心状态仍然重要时，才添加 core_attn。
   • 分组 MoE： 如果专家中间激活占主导，先尝试 moe_act；如果归一化输出占用明显，再添加 layernorm。使用整个 moe 重计算前，应先考虑它会重放额外的专家计算和通信。
   • 稠密 FFN： mlp 可以节省整个稠密 MLP 激活区域的内存，但其计算开销通常高于较窄的输出丢弃边界。
4. 每次只更改一个标签。记录每个 rank 的已分配/已预留峰值，以及稳态步时或吞吐量；不要根据单个配置推断全局模块排名。
5. 只有在有针对性的选择性边界无法让工作负载运行时，才使用整层重计算。整层重计算的内存影响最大，重放开销也最高。
6. 将 CUDA graphs、FP8、上下文并行通信和重叠功能视为兼容性约束，而非事后才考虑的因素。

当 PCIe/NVLink 传输开销比重放计算更合适时，Megatron Core 的 cpu_offloading=True 是一种替代方案。它不能与激活重计算同时使用，并且与大于 1 的流水线并行度不兼容。

启用

选择性重计算

cfg.model.recompute_granularity = "selective"
cfg.model.recompute_modules = ["core_attn"]  # Common standard-attention candidate, not a universal default.

对于 MLA、MoE、稠密 MLP 或 GDN 工作负载，请根据下方决策表替换或扩展该列表。

整层重计算

cfg.model.recompute_granularity = "full"
cfg.model.recompute_method = "uniform"
cfg.model.recompute_num_layers = 1

• uniform：按 recompute_num_layers 指定的固定组对 Transformer 层设置检查点。
• block：在每个流水线阶段为前 recompute_num_layers 层设置检查点，并按虚拟流水线感知的方式进行分配。

选择性模块决策表

当前固定版本的 Megatron Core 接受以下标签。开发分支可能会添加特定于模型的标签，因此请针对确切的目标修订版本进行验证，而不要将某个分支的列表照搬到其他分支。

模块	检查点边界	适用测试场景	主要成本或注意事项
core_attn	核心注意力	标准注意力，尤其是保留注意力中间结果的未融合后端	重放注意力。使用 TE 融合注意力／Flash Attention 时，额外节省可能很少；上下文并行可能导致注意力通信被重放。
mla_up_proj	MLA Q/KV 上投影及 RoPE 区域	保留扩展后 Q/K/V 张量的 MLA 模型	重放 MLA 扩展路径。它是与 core_attn 不同、可能叠加的边界。
layernorm	输入和预 MLP 归一化输出	归一化输出对峰值内存占用贡献显著，通常与 MoE 或 MLA 边界同时出现	通常范围较窄，但节省效果取决于隐藏维度、序列长度以及处于活动状态的图路径。
moe_act	分组专家 FC1 与 FC2 之间的激活输出	分组 MoE 的专家中间激活占主导	范围较窄的输出丢弃检查点。它不会重放分发、FC1 或 FC2，但受到 FP8 延迟缩放限制。
mlp	整个稠密 MLP	用尽较窄边界后，稠密层仍占主导	重放整个稠密 MLP。对 MLP 为 MoE 的层无效。
moe	整个 MoE 前向过程	必须丢弃较大的 MoE 区域才能使工作负载正常运行	重放路由、分发／合并通信、专家计算和共享专家计算。与专家并行重叠不兼容。
shared_experts	未重叠的共享专家 MLP	共享专家构成独立且显著的峰值	重放共享专家 MLP；与共享专家重叠不兼容。外层 moe 已移除其原始前向过程保存的结果，但嵌套仍可能改变反向重放期间的瞬时峰值。
gdn_norm_out	GDN 门控归一化输出	GDN／混合模型会保留此输出	重放归一化及其 HP 到 CP 的全互联路径。

例如，DeepSeek V4 配置只有在使用其所需的 Megatron Core 开发分支时才能使用特定于模型的 mhc 标签。它不是适用于固定版本的通用标签，因此未列入上表。

因此，常见性能配置通常分为几种模式，而不是遵循一套通用列表：

• 标准 Transformer 配方通常使用 core_attn；
• MLA 配方通常使用 mla_up_proj，有时也会搭配 mlp；
• 分组 MoE 配方通常使用 moe_act，或搭配 moe_act 使用 layernorm；
• 内存压力更大的 MoE 配方有时会使用更宽泛的组合，例如 moe 加 layernorm。

这些是候选模式，并不保证适用顺序。最终列表应由峰值归因和匹配测量结果决定。

测量约定

对于每个候选项，请记录：

• 准确的 Bridge 和 Megatron Core 修订版本；
• 模型、序列长度、微批次／全局批次大小、精度、注意力后端和并行方式；
• 准确的 recompute_granularity、模块列表、方法和层数；
• 每个 rank 的 max_memory_allocated() 和 max_memory_reserved()；
• 预热后的稳态步时或吞吐量；
• 与任务相适应的简短收敛性或数值健全性检查。

使用匹配的无重计算对照，并一次只更改一种重计算选择。不同作业、后端或并行布局的峰值内存不能作为模块排名基准。

不要仅仅因为候选配置比对照配置多运行了几步，就判定其成功。请运行到优化器状态初始化完成，并继续执行多个稳态步骤：选择性重计算可能会将内存瓶颈从前向过程转移到梯度同步或优化器阶段，却仍无法使工作负载正常运行。

匹配的 H100 证据：Moonlight 16B

一项于 2026-08-12 进行的短程研究使用了确切的 Bridge 修订版本 600d069b824dd5ce50367a311a5a3244478faf22 和 Megatron Core 修订版本 24bad8e677d22625d86ef2a54c9506b6e4992c93。Moonlight 16B BF16 预训练配方在 8 张 H100 80GB GPU 上运行，序列长度为 4096，MBS=1、 GBS=4、TP=2、PP=1、CP=1、EP=8，使用模拟数据并运行 20 步。该模型由一个 稠密层和 26 个 MLA+MoE 层组成。每一行仅更改了 recompute_modules；20 次迭代的损失值均为有限值，跳过或出现 NaN 的 迭代次数均为零。

峰值已分配内存是第 2 次迭代后报告的优化器之后数值中的最大值。时间和吞吐量为第 11--20 次迭代的平均值。

选择性模块	峰值已分配内存（GB）	步时（ms）	TFLOP/s/GPU	相对于 [] 的已分配内存	相对于 [] 的时间
[]	36.618	457.18	77.50	对照	对照
core_attn	36.614	474.80	74.44	-0.01%	+3.85%
mla_up_proj	35.902	480.89	73.72	-1.96%	+5.19%
mla_up_proj、mlp	35.917	496.73	71.50	-1.91%	+8.65%
moe_act	35.941	466.26	75.50	-1.85%	+1.99%
layernorm、moe_act	35.949	506.53	70.27	-1.83%	+10.79%

对于这一确切工作负载，moe_act 是最佳的首选边界：它以较低的重放开销，回收了几乎与 mla_up_proj 一样多的已分配内存。 如果额外减少约 39 MB 内存很重要，mla_up_proj 是下一个候选项。将 mlp 加入 mla_up_proj，或将 layernorm 加入 moe_act，都没有降低观测到的峰值，反而使步时变长。在融合注意力下，显式指定 core_attn 会增加开销，却未带来显著的内存收益。

最大保留内存保持在约 40 GB 附近，并未单调下降。这是分配器缓存导致的，并非相反证据：本研究依据已分配内存和端到端步骤成功完成情况来选择边界。

匹配的 H100 证据：Nemotron 3 Nano

同一项于 2026-08-12 进行的研究使用了适用于这款 52 层混合 Mamba/融合注意力 MoE 模型的原生 16 张 H100 BF16 性能配方。匹配的短程配置使用序列长度 8192、MBS=1、GBS=16、TP=1、PP=1、CP=1、 EP=8、DP=16、expert-DP=2、HybridEP、分组 GEMM、用于注意力和 Mamba 的 TE CUDA 图、 模拟数据并运行 12 步。每一行仅更改了 recompute_modules。

选择性模块	结果	Rank-0 测得峰值	失败或稳态证据
[]	第 1 次迭代后 OOM	第 1 次迭代后为 66.297 GB	第 2 次迭代的 MoE 路由器内存分配失败；高负载 rank 已分配约 72.9 GiB。
core_attn	第 1 次迭代中 OOM	不可比较	分组专家线性层内存分配失败；显式重计算注意力仍无法让融合注意力工作负载运行。
moe_act	第 1 次迭代后 OOM	第 1 次迭代后为 62.103 GB	在匹配检查点下，比对照低 4.194 GB（6.33%），但第 2 次迭代的输出投影仍需要 2 GiB。
layernorm、moe_act	第 1 次迭代中 OOM	不可比较	输出投影仍需要 2 GiB；CUDA 图私有内存池影响显著。
moe	完成 12 步	第 2 次迭代后为 64.653 GB	第 7--12 次迭代的步时为 657.42 ms，吞吐量为 277.72 TFLOP/s/GPU。
moe、layernorm	完成 12 步	第 2 次迭代后为 63.639 GB	第 7--12 次迭代的步时为 677.62 ms，吞吐量为 270.62 TFLOP/s/GPU。

两种成功配置的损失值均为有限值，跳过或出现 NaN 的迭代次数均为零。 对于这一确切的容量受限配方，整层 moe 重计算是测试过的最小可行边界。加入 layernorm 又回收了 1.014 GB（1.57%）的 rank-0 峰值内存，但步时高出 3.07%；因此，当需要这部分余量时，采用该配方中更广泛的组合是合理的。范围较窄的 moe_act 确实减少了激活内存，但仍无法让整个训练步骤顺利运行。

一次探索性的原生 8 张 H100 布局即使在序列长度为 4096 时，也在 FP32 优化器状态初始化期间失败。这属于优化器容量问题，并非用于比较选择性边界吞吐量的基准；本文不采用这些运行的计时结果。

跨模型结论

这些测量无法得出统一的排名。Moonlight 在空控制组下表现良好，并倾向于使用范围较窄的 moe_act；Nemotron 则需要对整个 moe 进行广泛重计算；历史上的稠密 Llama 证据表明，重放整个 mlp 成本高昂，而且缺少空控制组。因此，正确的首选应是架构和峰值所指向的最窄边界；只有当较窄的选项无法通过完整步骤时，才尝试更广泛的重放。

兼容性与验证

配置语义

• recompute_granularity="selective" 使用 recompute_modules；空列表可作为显式控制组。
• recompute_granularity="full" 使用 recompute_method 和 recompute_num_layers；选择性标签不适用。
• 完全重计算会取代选择性模块选项，而非与之组合。
• 未知标签无法通过 Megatron Core 验证。开发分支中的标签可能不同，因此应以对应精确版本的 TransformerConfig 验证器为准。

注意力后端与上下文并行

• TE 融合注意力和 Flash Attention 已使用内部重计算。显式设置 core_attn 仍可能改变保留的输入/输出，但必须在匹配的 [] 对比中证明其价值。
• 在上下文并行下，注意力检查点可能会重放通信以及计算。测量记录中应包含 CP 大小和拓扑。

MoE 限制

• 整个 moe 的重计算与专家并行重叠不兼容，因为反向重放会重复已重叠的路由/通信区域。
• shared_experts 重计算与共享专家重叠不兼容。
• moe_act 适用于分组 GEMM 专家；如果只需丢弃专家激活，这是范围较窄的选择。
• mlp 针对稠密 MLP，在 MoE 层中不起作用；混合稠密/MoE 模型仍可在其稠密层上受益。

FP8 限制

• moe_act 和 layernorm 重计算不支持 FP8 延迟缩放，且需要兼容的 Transformer Engine 版本。
• 已吸收的 MLA 路径还有额外的 FP8/FP4 限制。选择 mla_up_proj 前，请验证确切的模型/提供方路径。

CUDA 图

• 只有当检查点模块完全位于所选图范围内或完全位于其外时，选择性重计算才有效。跨越图边界的检查点边界无效。
• 捕获/预热可能会绕过检查点包装器，因此应验证最终图范围和重放路径，不要假设急切执行时的行为会延续。
• 在固定版本的 Megatron Core 中，使用 CUDA 图进行完全重计算需要设置 cuda_graph_impl="full_iteration"。否则应禁用 CUDA 图；在此处，范围式/局部图捕获不能替代完整迭代捕获。

历史测量：背景信息，而非模块排名

Bridge PR #3107 中的历史 H100 测量采用 Llama 3 70B SFT，在 32 张 H100 80GB GPU 上运行，使用 FP8 当前缩放、序列长度 4096、微批次大小 1、全局批次大小 32、TP=4、PP=4、VPP=5 和 DP=2：

配置	TFLOP/s/GPU	峰值内存
该次运行中的 core_attn 基线	~704	58.8 GB（rank 0 上 OOM）
mlp	593.6	55.6 GB
mlp + core_attn	586.8	55.6 GB
core_attn + layernorm	~702	59.6 GB（rank 0 上 OOM）
PR 背景中记录的黄金吞吐量	709.93	非配对内存测量

这项证据的局限：

• 未包含匹配的无重计算行；
• 黄金行并非配对的单模块对比；
• 测量仅涵盖一种稠密 Llama 工作负载，不包括 MLA 或 MoE；
• 该表仅能说明局部的内存/吞吐量权衡，不得用于对所有重计算标签进行排名。

代码锚点

• 选择性标签验证和跨特性检查：3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
• 检查点实现：3rdparty/Megatron-LM/megatron/core/tensor_parallel/random.py
• 标准注意力检查点边界：3rdparty/Megatron-LM/megatron/core/transformer/attention.py
• MLA 上投影边界：3rdparty/Megatron-LM/megatron/core/transformer/multi_latent_attention.py
• Layernorm、密集 MLP 和外层 MoE 的放置：3rdparty/Megatron-LM/megatron/core/transformer/transformer_layer.py
• 分组专家激活边界：3rdparty/Megatron-LM/megatron/core/transformer/moe/experts.py
• 共享专家和完整 MoE 路径：3rdparty/Megatron-LM/megatron/core/transformer/moe/moe_layer.py
• GDN 归一化边界：3rdparty/Megatron-LM/megatron/core/ssm/gated_delta_net/gdn.py

故障诊断

症状	可能原因	后续操作
core_attn 几乎没有或完全没有降低峰值	融合注意力/Flash 注意力已对开销较大的内部计算进行重计算，或者峰值出现在其他位置	与 [] 对比，定位峰值，然后测试架构专属边界，例如 mla_up_proj 或 moe_act。
使用 core_attn 后 MLA 仍然 OOM	扩展后的 Q/K/V 投影张量，而非注意力核心张量，占据了主要内存	测试 mla_up_proj；只有在对比证据支持时才添加 core_attn。
MoE 峰值仍然很高	专家中间结果或归一化输出占据主要内存	先测试 moe_act，再测试 layernorm；仅在内存压力更广泛时才使用完整 moe。
专家重叠验证失败	完整 moe 或 shared_experts 重计算与重叠执行冲突	保留重叠执行并使用兼容的内部边界，或者禁用重叠执行并重新测量整个配置。
所选标签没有可测量的效果	对应模块在测量的层中不存在或未激活，或者图捕获绕过了包装器	检查 provider/层组合和最终图范围；例如，mlp 对纯 MoE 层无效。
完整重计算加 CUDA 图触发断言	图实现不是完整迭代模式	设置 cuda_graph_impl="full_iteration" 或禁用 CUDA 图。
保留内存很高，但已分配内存保持稳定	分配器碎片化或缓存	在添加重计算前尝试设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True。
启用重计算后，OOM 转移到其他 rank	流水线/虚拟流水线的层分布改变了瓶颈	对比每个 rank 的峰值，并针对实际热点阶段调整完整块/均匀放置或选择性边界。
候选配置运行得更远，但仍然 OOM	重计算将峰值转移到了梯度同步或优化器状态初始化阶段	记录故障阶段的变化，作为诊断证据；但在优化器初始化和多个稳定运行步骤完成前，不要判定为通过。

已知限制

• 模块列表无法跨模型系列、注意力后端、并行布局或 Megatron Core 版本通用。
• 边界相互重叠或嵌套时，内存节省是非线性的；简单相加不可靠。
• 完整重计算会改变 RNG 执行路径；使用 dropout 的工作负载需要进行数值/收敛检查。
• 激活重计算无法解决参数、优化器状态或分配器碎片化造成的内存压力。
• 正确的结果是满足每个 rank 的内存目标所需的最小实测重放开销，而不是最长的模块列表。

延伸阅读

• docs/performance-guide.md
• skills/nemo-mbridge-perf-memory-tuning/SKILL.md
• skills/nemo-mbridge-perf-cuda-graphs/SKILL.md
• skills/nemo-mbridge-perf-cpu-offloading/SKILL.md
• Megatron Core 激活重计算指南：https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/index.html