| 名称 | nemo-mbridge-perf-activation-recompute |
| 描述 | >- Validate and use selective and full activation recompute in Megatron Bridge to reduce GPU memory usage at the cost of extra compute. Use for activation memory OOMs or regressions involving recompute_granularity, recompute_num_layers, recompute_modules, recompute_method, selective recompute, full recompute, or activation checkpointing. |
| 开源协议 | Apache-2.0 |
激活重计算
Stable docs: @docs/training/activation-recomputation.md Card: @skills/nemo-mbridge-perf-activation-recompute/card.yaml
<!-- Guidance refreshed: 2026-08-12. -->
激活重计算（activation checkpointing）用额外的正向计算换取更低的保留激活内存。有效的检查点边界取决于模型架构、注意力后端、并行度和实际驱动每个rank峰值的张量。
快速决策指南
- 确认压力是真实分配，而不是分配器碎片。在每个rank上比较
max_memory_allocated()和max_memory_reserved()。 - 当工作负载可以适应时保留一个显式的无重计算对照组。在选择性粒度下，
recompute_modules=[]是有效且有用的对照。 - 从架构和观察到的峰值中选择第一个边界：
- 标准注意力：
core_attn是常见的第一候选。当未融合注意力物化score/probability张量时最强。在Transformer Engine融合或Flash Attention下，由于这些后端已经重新物化注意力内部，应与[]比较。 - 多头潜在注意力（MLA）： 当扩展的Q/K/V投影占主导时，从
mla_up_proj开始。仅当注意力核心状态仍然重要时才添加core_attn。 - 分组MoE： 当专家中间激活占主导时从
moe_act开始；当norm输出重要时添加layernorm。只有在计入额外专家计算和通信重放后，才使用整个moe重计算。 - 稠密FFN：
mlp可以保存整个稠密MLP激活区域，但通常比窄输出丢弃边界计算成本更高。
- 标准注意力：
- 一次只更改一个标签。记录每个rank的allocated/reserved峰值以及稳定状态的步时间或吞吐；不要从一个配置推断全局模块排名。
- 仅当目标选择性边界无法让工作负载适应时才使用整层重计算。全量重计算内存影响最广，重放成本最大。
- 将CUDA图、FP8、上下文并行通信和重叠特性视为兼容性约束，而不是事后考虑。
Megatron Core的 cpu_offloading=True 是替代方案，当PCIe/NVLink传输开销优于重放计算时使用。它不能与激活重计算组合，且与大于1的流水线并行不兼容。
启用
选择性重计算
cfg.model.recompute_granularity = "selective"
cfg.model.recompute_modules = ["core_attn"] # 常见标准注意力候选，不是通用默认值。
使用下面的决策表为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区域 | MLA模型保留扩展的Q/K/V张量 | 重放MLA扩展路径。这是与core_attn不同且可能可叠加的边界。 |
layernorm |
输入和pre-MLP归一化输出 | Norm输出显著贡献峰值，通常与MoE或MLA边界同时 | 通常较窄，但节省取决于隐藏大小、序列长度和活动图路径。 |
moe_act |
分组专家FC1和FC2之间的激活输出 | 分组MoE专家中间激活占主导 | 窄输出丢弃检查点。它不重放dispatch、FC1或FC2，但具有FP8延迟缩放限制。 |
mlp |
整个稠密MLP | 在较窄边界耗尽后，稠密层占主导 | 重放完整稠密MLP。对MLP为MoE的层没有影响。 |
moe |
整个MoE前向 | 必须丢弃MoE区域才能让工作负载适应 | 重放路由、dispatch/combine通信、专家和共享专家工作。与专家并行重叠不兼容。 |
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或layernorm加moe_act； - 更高压力的MoE配方有时使用更宽的组合如
moe加layernorm。
这些是候选模式，不是排序保证。峰值归属和匹配测量决定最终列表。
测量契约
对每个候选项，捕获：
- 精确的Bridge和Megatron Core修订版；
- 模型、序列长度、micro/global批大小、精度、注意力后端和并行度；
- 精确的
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步。该模型混合了1个稠密层和26个MLA+MoE层。每行只更改 recompute_modules；所有20个loss均有限，零跳过或NaN迭代。
峰值分配内存是迭代2后报告的最大后优化器值。时间和吞吐是迭代11-20的平均值。
| 选择性模块 | 峰值分配 (GB) | 步时间 (ms) | TFLOP/s/GPU | 分配 vs [] |
时间 vs [] |
|---|---|---|---|---|---|
[] |
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 一样多的分配内存。如果 mla_up_proj 额外约39 MB的减少很重要，则它是下一个候选。将 mlp 加到 mla_up_proj 或 layernorm 加到 moe_act 没有改善观测到的峰值且使步变慢。显式 core_attn 在融合注意力下增加成本而没有实际内存收益。
最大保留内存保持约40 GB左右，并未单调下降。那是分配器缓存，不是反证：本研究中的边界选择基于分配内存和成功的端到端步骤。
H100匹配证据：Nemotron 3 Nano
同一个2026-08-12研究使用本地16-H100 BF16性能配方，针对52层混合Mamba/融合注意力MoE模型。匹配短运行配置使用序列长度8192，MBS=1，GBS=16，TP=1，PP=1，CP=1，EP=8，DP=16，expert-DP=2，HybridEP，分组GEMM，TE CUDA图用于注意力和Mamba，模拟数据，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。 |
两个成功行都有有限loss和零跳过或NaN迭代。对于该确切容量受限配方，整个moe重计算是通过测试的最小边界。添加layernorm以多3.07%步时间换取额外1.014 GB（1.57%）rank-0峰值降低，因此当需要该余量时，配方的更宽组合是合理的。窄的moe_act产生了实际激活缓解，但未使整个训练步骤可行。
探索性的原生8-H100布局在FP32优化器状态初始化期间即使序列长度4096也失败。那是优化器容量问题，不是选择性边界吞吐基线；这里不使用这些运行的时序比较。
跨模型结论
这些测量不定义唯一排名。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图
- 仅当被检查点的模块完全位于所选图作用域内部或外部时，选择性重计算才有效。跨越图边界的检查点边界无效。
- 捕获/预热可能绕过检查点包装器，因此最终验证图作用域和重放路径，而不是假设急切行为。
- 全量重计算搭配CUDA图需要在固定Megatron Core中使用
cuda_graph_impl="full_iteration"。否则应禁用CUDA图；作用域/局部图捕获不是本场景全迭代捕获的替代。
历史测量：背景而非模块排名
Bridge PR #3107的历史H100测量使用Llama 3 70B SFT，在32个H100 80GB GPU上，FP8 current scaling，序列长度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。 |
MLA在core_attn后仍OOM |
扩展的Q/K/V投影张量而非注意力核心张量占主导 | 测试mla_up_proj；只有匹配证据支持时才加core_attn。 |
| MoE峰值仍然高 | 专家中间或norm输出占主导 | 测试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.mdskills/nemo-mbridge-perf-memory-tuning/SKILL.mdskills/nemo-mbridge-perf-cuda-graphs/SKILL.mdskills/nemo-mbridge-perf-cpu-offloading/SKILL.md- Megatron Core激活重计算指南：https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/index.html