nemo-mbridge-perf-activation-recompute

激活重计算

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

答案检查清单

对于 OOM 或 CUDA 图相关问题，请按以下确切顺序说明：

1. 首先尝试 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True；许多临界故障是由分配器碎片化导致的，而非激活容量不足。
2. 优先使用选择性重计算，而不是完整层重计算： recompute_granularity="selective" 和 recompute_modules=["core_attn"]。
3. 如果仍处于临界状态，可选地添加 "layernorm"；只有在最后才使用 "mlp"，因为它在宽型稠密 FFN 上会带来较大的计算开销。
4. 只有在选择性重计算仍无法满足显存需求时，才使用完整层重计算，并始终指出必需字段：recompute_granularity="full"、recompute_method 和 recompute_num_layers。
5. 如果启用了 FP8 或限定于 TE 的 CUDA 图，请指出断言风险：完整层重计算与 attn、mlp 和 moe_router 等 TE 范围不兼容。有效的解决方法包括选择性重计算、cuda_graph_impl="none"，或使用 cuda_graph_impl="local" 并设置 cuda_graph_scope="full_iteration"。

功能说明

激活重计算通过在前向传播期间丢弃中间激活，并在反向传播期间重新计算这些激活，以 GPU 计算换取显存。Megatron Bridge 支持两种粒度：

粒度	指定内容	重计算内容	显存节省	计算开销
selective	recompute_modules 列表（例如 core_attn、mlp）	每层中的特定子模块	中等（取决于模块）	低到高
full	recompute_num_layers + recompute_method	整个 Transformer 层（N 层）	最大	最高

注意：MCore 将其称为“selective”（子模块级）和“full”（层级）。这里的“Full”表示重计算完整层，而非整个模型——仍可通过 recompute_num_layers 选择层数。

快速决策

1. 首先设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True——大多数临界 OOM 都是由内存碎片化而非容量不足导致的。此方法零成本即可解决问题。请参阅 @skills/nemo-mbridge-perf-memory-tuning/SKILL.md。
2. 从 recompute_granularity=selective、recompute_modules=[core_attn] 开始（配方中通常已将其设为默认值）。
3. 将 layernorm 添加到重计算模块中——计算开销几乎为零，但节省的显存也微乎其微。仅在极端临界情况下有帮助。
4. 将 mlp 作为最后手段——可节省约 3 GB 显存，但在大型稠密模型（Llama3 70B）上会降低约 16% 的 GPU 利用率。
5. 只有在选择性重计算仍无法满足显存需求时，才使用 recompute_granularity=full。

CPU 卸载（cpu_offloading=True）是另一种方案，可完全避免重计算开销，但与 PP > 1 不兼容。

启用

选择性重计算

cfg.model.recompute_granularity = "selective"
cfg.model.recompute_modules = ["core_attn"]  # add "layernorm", "mlp", or other valid modules as needed

完整层重计算

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

可用的 recompute_modules

模块	重计算内容	计算开销	显存节省
core_attn	注意力 softmax/dropout/QKV 点积	低（Flash Attention 已在内部进行重计算）	中等
layernorm	层归一化	可忽略不计（约 0%）	可忽略不计
mlp	完整 FFN 块	高（Llama3 70B、hidden=28672 时约为 16%）	约 3 GB
moe	MoE 专家分发	不定	不定
moe_act	MoE 激活函数	低	少量
shared_experts	共享专家层	中等	中等
mla_up_proj	多潜变量注意力上投影	中等	中等

性能测试框架 CLI

python scripts/performance/run_performance_workload.py \
  --recompute_granularity selective \
  --recompute_modules core_attn layernorm \
  ...

兼容性和限制

• recompute_granularity=selective 要求 recompute_modules 列表非空
• recompute_granularity=full 要求设置 recompute_method 和 recompute_num_layers
• 层级重计算（recompute_granularity="full" + recompute_num_layers）与 TE 作用域的 CUDA 图不兼容。 MCore 将此称为「full」粒度——这里的名称指重计算完整的 Transformer 层，而不是整个模型。即使你选择了要重计算的层数，MCore 对它的处理方式也与子模块重计算不同。任何 TE 作用域（attn、mlp、moe_router 等）都会触发断言。这常见于默认启用 TE 作用域 CUDA 图的 FP8 配置（例如，LLAMA3_70B_SFT_CONFIG_H100_FP8_CS_V1 设置了 cuda_graph_impl="transformer_engine"、cuda_graph_scope="mlp"）。可选方案：
  • 使用子模块重计算（recompute_granularity="selective" + recompute_modules）——与 TE 作用域的 CUDA 图兼容
  • 禁用 CUDA 图（cuda_graph_impl="none"）并使用层级重计算
  • 切换为 cuda_graph_impl="local"、cuda_graph_scope="full_iteration"
• distribute_saved_activations=True 不能与 sequence_parallel=True 同时使用
• 组合使用 mlp + core_attn 重计算的效果略差于仅使用 mlp， 原因是双重重计算开销

测量结果

Llama3 70B SFT，运行于 32x H100 80GB，FP8（当前扩展配置）：

• 基线：TP=4、PP=4、VPP=5、DP=2、MBS=1、GBS=32、seq_len=4096
• 黄金 GPU 利用率：709.93 TFLOP/s/GPU
• 性能回退阈值：5%

实验	recompute_modules	TFLOP/s/GPU	相对黄金值	峰值内存（GB）	结果
基线	[core_attn]	~704	-0.8%	58.8（rank0 OOM）	OOM
实验 1	[mlp]	593.6	-16.4%	55.6	性能回退
实验 2	[mlp, core_attn]	586.8	-17.3%	55.6	性能回退
实验 3	[core_attn, layernorm]	~702	-1.1%	59.6（rank0 OOM）	OOM

主要结论：

• layernorm 重计算几乎不增加计算开销，但节省的内存微乎其微
• mlp 重计算可将峰值内存降低约 3 GB，但会带来约 16% 的性能损失，因为 Llama3 70B FFN（hidden=28672）的重计算开销很高
• 组合使用 mlp + core_attn 的效果略差于仅使用 mlp
• 对于此工作负载，实际解决 OOM 的方法是设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True （原因是内存碎片，而非容量不足）。参见 @skills/nemo-mbridge-perf-memory-tuning/SKILL.md。

代码锚点

重计算模块枚举和 selective 检查点逻辑

# 3rdparty/Megatron-LM/megatron/core/transformer/transformer_block.py
# _checkpointed_forward() applies selective recompute based on recompute_modules

重计算配置验证

# 3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
# Validates recompute_granularity, recompute_method, recompute_num_layers

Llama3 配方默认值

    # Memory saving (recompute & offloading)
    cfg.model.recompute_granularity = None
    cfg.model.recompute_modules = None
    cfg.model.fine_grained_activation_offloading = False
    cfg.model.offload_modules = None

Full 重计算 + CUDA 图断言（MCore）

            if self.recompute_granularity:
                if self.recompute_granularity != "selective":
                    assert self.cuda_graph_scope == [
                        CudaGraphScope.full_iteration
                    ], "full recompute is only supported with full iteration CUDA graph."

CPU 卸载与 PP 不兼容（MCore）

        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
            raise ValueError(
                "Currently there is no support for Pipeline parallelism with CPU offloading"
            )

故障诊断

现象	原因	确认方法	修复方法
GPU 利用率下降 >15%	对大型 FFN 使用 mlp 重计算	检查 recompute_modules 是否包含 mlp	检查是否设置了 expandable_segments:True；考虑减小 MBS
添加 layernorm 后仍然 OOM	layernorm 激活值过小	比较添加前后的峰值内存	添加 mlp 重计算或检查是否设置了 expandable_segments:True
AssertionError: full recompute is only supported with full iteration CUDA graph	使用 TE 作用域 CUDA 图进行层级重计算（recompute_granularity=full + recompute_num_layers）。FP8 CS 配置默认使用 cuda_graph_impl=transformer_engine、scope=mlp。	检查 cuda_graph_impl 和 cuda_graph_scope	使用子模块重计算（selective + recompute_modules），或设为 cuda_graph_impl=none，或使用 local + full_iteration
ValueError：PP + CPU 卸载	在 pipeline_model_parallel_size > 1 时启用了 cpu_offloading=True	检查 PP 配置	禁用 CPU 卸载或将 PP 设为 1
mlp+core_attn 的效果差于仅使用 mlp	双重重计算开销	比较实验 1 和实验 2	仅使用 mlp

已知限制

• 每个模块的内存节省效果因模型架构和隐藏维度而有很大差异
• 不会自动选择模块——用户必须选择要重新计算的模块
• 单独使用 layernorm 重新计算几乎从来不值得
• PP > 1 时，CPU 卸载（零计算成本的替代方案）会受到阻止

验证

uv run python -m pytest \
  tests/unit_tests/training/test_config.py -k "recompute" -q

成功标准：

• 重新计算配置验证的单元测试通过
• 配置验证不会触发断言错误