nemo-mbridge-perf-activation-recompute

该技能用于在 Megatron Bridge 中配置激活重计算，通过前向丢弃中间激活、反向重算来换取更低显存。支持两种粒度：选择性（recompute_granularity="selective" 配合 recompute_modules，重算层内子模块）和整层（recompute_granularity="full"，配合 recompute_method 与 recompute_num_layers，重算 N 个 transformer 层，注意并非整个模型）。可选模块包括 core_attn、layernorm、mlp、moe、moe_act、shared_experts、mla_up_proj。建议先设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，许多边界 OOM 实为显存碎片而非容量不足；随后优先选择性重算 core_attn，视情况加 layernorm，mlp 作为最后手段。整层重算仅在选择性重算仍不足时使用，且必须指定 recompute_method 与 recompute_num_layers。约束方面：选择性子模块重算与 TE 作用域 CUDA 图兼容；整层重算与 TE 作用域 CUDA 图不兼容，会触发断言，可选改用子模块重算、关闭 CUDA 图（cuda_graph_impl="none"）或改为 local + full_iteration。distribute_saved_activations=True 不能与 sequence_parallel=True 同用；CPU offloading 与 PP>1 不兼容。实测显示 mlp 重算约省 3 GB 但 Llama3 70B 上损失约 16% 利用率，layernorm 收益极小。可用 run_performance_workload.py 验证，并用 pytest 检查配置。