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_scope="full_iteration" を指定した cuda_graph_impl="local" です。

概要

活性化の再計算では、順伝播中に中間活性化を破棄し、逆伝播中に再計算することで、GPU の計算量と引き換えにメモリーを節約します。Megatron Bridge は、次の 2 つの粒度に対応しています。

粒度	指定する内容	再計算される内容	メモリー節約量	計算コスト
selective	recompute_modules リスト（例：core_attn、mlp）	各レイヤー内の特定のサブモジュール	中程度（モジュールによって異なります）	低～高
full	recompute_num_layers + recompute_method	トランスフォーマーのレイヤー全体（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）では GPU 使用率が約 16% 低下します。
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	attention の 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	Multi-Latent Attention のアッププロジェクション	中程度	中程度

パフォーマンスハーネス 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 スコープのグラフを有効にする 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_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
• Golden GPU 利用率：709.93 TFLOP/s/GPU
• 回帰しきい値：5%

実験	recompute_modules	TFLOP/s/GPU	Golden 比	ピークメモリ（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 節約できますが、Llama3 70B の FFN（hidden=28672）は再計算コストが高いため、性能が約 16% 低下します
• mlp と core_attn を組み合わせると、mlp 単独よりわずかに性能が低下します
• このワークロードで実際に OOM を解消したのは PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True でした（容量ではなく、メモリの断片化が原因）。@skills/nemo-mbridge-perf-memory-tuning/SKILL.md を参照してください。

コードの参照箇所

再計算モジュールの列挙型と選択的チェックポイント処理

# 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

全再計算と 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 スコープのグラフを使ったレイヤーレベルの再計算（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 offloading	pipeline_model_parallel_size > 1 のときに cpu_offloading=True が指定されている	PP 設定を確認します	CPU オフロードを無効にするか、PP=1 に設定します
mlp+core_attn が mlp 単独より低性能	再計算のオーバーヘッドが重なる	実験 1 と実験 2 を比較します	mlp 単独を使用します

既知の制限事項

• モジュールごとのメモリ削減量は、モデルのアーキテクチャと隠れ次元によって大きく異なります
• モジュールの自動選択には対応していません。再計算するモジュールはユーザーが選択する必要があります
• layernorm の再計算は、単独の対策としてはほとんどの場合、効果がありません
• CPU オフロード（計算コストがかからない代替手段）は、PP > 1 の場合にブロックされます

検証

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

成功基準：

• 再計算設定の検証に関するユニットテストが成功すること
• 設定の検証でアサーションエラーが発生しないこと