nemo-mbridge-perf-activation-recompute

アクティベーションの再計算

安定版ドキュメント：@docs/training/activation-recomputation.md カード：@skills/nemo-mbridge-perf-activation-recompute/card.yaml

アクティベーションの再計算（アクティベーションチェックポイント）は、保持するアクティベーションのメモリを削減する代わりに、バックワード時のフォワード計算を追加で行います。有効なチェックポイント境界は、モデルアーキテクチャ、アテンションバックエンド、並列化、および各ランクのピークを実際に左右するテンソルによって異なります。

すばやく判断するためのガイド

1. メモリ不足の原因がアロケーターの断片化ではなく、実際の割り当てであることを確認します。すべてのランクで max_memory_allocated() と max_memory_reserved() を比較します。
2. ワークロードが収まる場合は、再計算なしの明示的な比較対象を維持します。選択的粒度では、recompute_modules=[] が有効であり、この比較に役立ちます。
3. アーキテクチャと観測されたピークに基づいて、最初の境界を選びます。
   • 標準アテンション： core_attn が一般的な最初の候補です。融合されていないアテンションがスコア／確率テンソルを具体化する場合に、最も効果があります。Transformer Engine の融合アテンションまたは Flash Attention を使用する場合は、これらのバックエンドがすでにアテンション内部の値を再計算するため、[] と比較してください。
   • Multi-Latent Attention（MLA）： 展開された Q/K/V プロジェクションが支配的な場合は、mla_up_proj から始めます。アテンションコアの状態も引き続き重要な場合に限り、core_attn を追加します。
   • Grouped MoE： エキスパートの中間アクティベーションが支配的な場合は、moe_act から始めます。ノルム出力の占める量が大きい場合は、layernorm を追加します。再実行される追加のエキスパート計算と通信を考慮したうえで、moe 全体の再計算を使用します。
   • Dense FFN： mlp は dense-MLP のアクティベーション領域全体を削減できますが、通常は出力を破棄する範囲を絞った境界よりも計算コストが高くなります。
4. 一度に変更するラベルは1つだけにします。ランクごとの割り当て／予約ピークと、定常状態のステップ時間またはスループットを記録します。1つのレシピからモジュール全体の順位を推測しないでください。
5. 対象を絞った選択的境界ではワークロードが収まらない場合に限り、全レイヤーの再計算を使用します。全再計算は、メモリへの影響が最も広く、再実行のコストも最大です。
6. CUDA グラフ、FP8、コンテキスト並列通信、オーバーラップ機能は、後回しにせず互換性の制約として扱います。

Megatron Core の cpu_offloading=True は、再実行される計算よりも PCIe／NVLink 転送のオーバーヘッドを許容できる場合の代替手段です。アクティベーションの再計算とは併用できず、パイプライン並列度が1を超える場合は互換性がありません。

有効化

選択的再計算

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

MLA、MoE、dense-MLP、または GDN のワークロードでは、以下の判断表を使ってこのリストを置き換えるか、項目を追加してください。

全レイヤーの再計算

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

• uniform：固定数の recompute_num_layers 個のトランスフォーマーレイヤーをグループ単位でチェックポイント化します。
• 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、または layernorm と moe_act を併用します。
• メモリー圧力がより高い MoE のレシピでは、ときに moe と layernorm のような広範囲の組み合わせを使用します。

これらは候補となるパターンであり、優先順位が保証されるものではありません。ピークの要因分析と条件を揃えた測定によって、最終的なリストを決定してください。

測定の条件

各候補について、以下を記録してください。

• Bridge と Megatron Core の正確なリビジョン
• モデル、シーケンス長、マイクロ／グローバルバッチサイズ、精度、アテンションバックエンド、並列化方式
• 正確な recompute_granularity、モジュールリスト、方式、レイヤー数
• ランクごとの max_memory_allocated() と max_memory_reserved()
• ウォームアップ後の定常状態におけるステップ時間またはスループット
• タスクに応じた、簡単な収束確認または数値的な健全性チェック

再計算なしの条件を揃えた対照実験を行い、再計算の選択肢は一度に1つだけ変更してください。異なるジョブ、バックエンド、並列化構成で測定したピークメモリーは、モジュールの優劣を比較するためのベンチマークにはなりません。

対照実験より先まで処理が進んだという理由だけで、候補が成功したと判断しないでください。オプティマイザー状態の初期化まで実行し、定常状態で複数のステップを完了してください。選択的な再計算によって、ワークロードが実行可能にならないまま、メモリーの限界がフォワードから勾配同期またはオプティマイザーへ移る場合があります。

条件を揃えた H100 の実測結果：Moonlight 16B

2026-08-12 に実施した短時間実行の調査では、Bridge のリビジョン 600d069b824dd5ce50367a311a5a3244478faf22 と Megatron Core のリビジョン 24bad8e677d22625d86ef2a54c9506b6e4992c93 を正確に使用しました。Moonlight 16B の BF16 事前学習レシピは、シーケンス長 4096、MBS=1、 GBS=4、TP=2、PP=1、CP=1、EP=8、モックデータ、20 ステップで、8 基の H100 80GB GPU 上で実行しました。このモデルは、1 つの Dense 層と 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 が挙げられます。mla_up_proj に mlp を追加しても、または moe_act に layernorm を追加しても、観測されたピークは改善せず、ステップは遅くなりました。融合アテンションでは、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 だけです。

選択したモジュール	結果	ランク 0 の測定ピーク	障害または定常状態の証拠
[]	イテレーション 1 の後に OOM	イテレーション 1 の後に 66.297 GB	イテレーション 2 で MoE ルーターの割り当てに失敗しました。メモリ使用量の多いランクでは、割り当て済みメモリが約 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 でした。

正常に完了した 2 つの行では、損失はすべて有限値で、スキップされたイテレーションも NaN のイテレーションもありませんでした。この容量が制約となるレシピでは、テストした中で最小の、学習を完了できる境界は moe 全体の再計算です。layernorm を加えると、ランク 0 のピークをさらに 1.014 GB（1.57%）削減できますが、ステップ時間は 3.07% 増加します。そのため、追加の余裕が必要な場合は、レシピのより広い組み合わせが妥当です。範囲の狭い moe_act でも実際にアクティベーションメモリは削減されましたが、学習ステップ全体を実行可能にはできませんでした。

ネイティブの 8 基の H100 構成を探索的に試したところ、シーケンス長 4096 でも FP32 オプティマイザー状態の初期化中に失敗しました。これはオプティマイザーの容量不足であり、選択的境界のスループットを比較するための基準ではありません。そのため、これらの実行結果から時間の比較は行っていません。

モデル間の結論

これらの測定結果は、単一のランキングを定義するものではありません。Moonlight は空のコントロールで適合し、狭い moe_act を好みました。Nemotron は広い whole-moe の再計算を必要としました。過去の dense Llama のエビデンスでは、whole-mlp のリプレイはコストが高く、空のコントロールが欠けていました。したがって、最初に試すべき正しい候補は、アーキテクチャとピークによって示唆される最も狭い境界であり、狭い選択が完全なステップを通過しない場合にのみ、より広いリプレイに進みます。

互換性と検証

構成のセマンティクス

• recompute_granularity="selective" は recompute_modules を使用します。空のリストは明示的なコントロールとして受け入れられます。
• recompute_granularity="full" は recompute_method と recompute_num_layers を使用します。selective のラベルは適用されません。
• full グラニュラリティは、selective のモジュール選択と組み合わせるのではなく、それを上書きします。
• 不明なラベルは Megatron Core の検証で失敗します。ラベルは開発ブランチによって異なる場合があるため、正確なリビジョンの TransformerConfig バリデータを信頼できる情報源として使用してください。

アテンションバックエンドとコンテキスト並列

• TE fused と Flash Attention はすでに内部で再マテリアライゼーションを使用しています。明示的な core_attn は保持される入力/出力を変える可能性がありますが、一致した [] 比較でその価値を証明する必要があります。
• コンテキスト並列では、アテンションのチェックポイントは計算だけでなく通信もリプレイする可能性があります。測定記録には CP サイズとトポロジーを含めてください。

MoE の制限

• whole-moe の再計算は、エキスパート並列のオーバーラップと互換性がありません。逆伝播のリプレイが、オーバーラップされたルーティング/通信領域を繰り返すためです。
• shared_experts の再計算は、shared-expert のオーバーラップと互換性がありません。
• moe_act は grouped-GEMM エキスパートに適用され、エキスパートのアクティベーションのみを破棄する必要がある場合の、より狭い選択です。
• mlp は dense MLP を対象とし、MoE レイヤーでは何も行いません。dense/MoE 混在モデルでは、dense レイヤーで依然として効果が得られる場合があります。

FP8 の制限

• moe_act と layernorm の再計算は、FP8 遅延スケーリングではサポートされておらず、互換性のある Transformer Engine バージョンが必要です。
• Absorbed MLA パスには、追加の FP8/FP4 制限があります。mla_up_proj を選択する前に、正確なモデル/プロバイダーのパスを検証してください。

CUDA グラフ

• selective 再計算は、チェックポイントされたモジュールが選択されたグラフスコープの完全に内側または完全に外側にある場合にのみ有効です。グラフ境界をまたぐチェックポイント境界は無効です。
• キャプチャ/ウォームアップはチェックポイントのラッパーをバイパスする可能性があるため、eager の動作がそのまま引き継がれると仮定せず、最終的なグラフスコープとリプレイパスを検証してください。
• CUDA グラフを伴う full 再計算には、固定された Megatron Core で cuda_graph_impl="full_iteration" が必要です。そうでない場合は CUDA グラフを無効にしてください。ここでは、スコープ付き/ローカルのグラフキャプチャは full-iteration キャプチャの代替にはなりません。

過去の測定: コンテキストであり、モジュールのランキングではない

Bridge PR #3107 の過去の H100 測定では、32 台の H100 80GB GPU で Llama 3 70B SFT を使用し、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	ペアのメモリ測定ではない

このエビデンスの限界:

• 一致する no-recompute の行が含まれていませんでした。
• ゴールデンの行は、ペアのモジュールのみの比較ではありませんでした。
• 測定は 1 つの dense 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
• レイヤーノーム、Dense 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 でピークがほとんど、またはまったく低下しない	Fused/Flash attention が高コストな内部処理をすでに再計算しているか、ピークが別の箇所にある	[] と比較してピークの発生箇所を特定し、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 の再計算がオーバーラップと競合している	オーバーラップを維持して互換性のある内側の境界を使用するか、オーバーラップを無効にして構成全体を再測定します。
選択したラベルに測定可能な効果がない	対象モジュールが測定対象のレイヤーに存在しないか、アクティブでない、またはグラフキャプチャがラッパーを迂回している	プロバイダーとレイヤーの構成、および最終的なグラフのスコープを確認します。たとえば、純粋な MoE レイヤーでは mlp は効果がありません。
全体の再計算と CUDA グラフの併用でアサートが発生する	グラフ実装が full-iteration ではない	cuda_graph_impl="full_iteration" を設定するか、CUDA グラフを無効にします。
予約済みメモリーは多いが、割り当て済みメモリーは安定している	アロケーターの断片化またはキャッシュ	再計算を追加する前に PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True を試します。
再計算を有効にすると、OOM が別のランクに移る	パイプライン／仮想パイプラインのレイヤー分散によってボトルネックが変化した	ランクごとのピークを比較し、実際に負荷が高いステージに合わせて、ブロック全体／均一な配置または選択的な境界を調整します。
候補の設定で処理は先に進むが、依然として OOM が発生する	再計算によってピークが勾配同期またはオプティマイザー状態の初期化に移った	変化した失敗ステージを診断の根拠として記録します。ただし、合格と判定するにはオプティマイザーの初期化と定常状態での複数ステップを確認する必要があります。

既知の制限事項

• モジュールのリストは、モデルファミリー、アテンションバックエンド、並列レイアウト、Megatron Core のリビジョンをまたいでそのまま利用できるものではありません。
• 境界が重複または入れ子になるとメモリー削減量は非線形になるため、単純な加算による計算は信頼できません。
• 全体の再計算は乱数生成の実行経路を変えるため、ドロップアウトを使用するワークロードでは数値精度と収束を確認する必要があります。
• アクティベーションの再計算では、パラメーター、オプティマイザー状態、アロケーターの断片化によるメモリー負荷には対処できません。
• 適切な結果とは、ランクごとのメモリー目標を満たす、測定された再実行コストが最小の設定です。モジュールリストを最長にすることではありません。

参考資料

• 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