nemo-mbridge-perf-memory-tuning

メモリ調整

安定版ドキュメント：@docs/parallelisms.md カード：@skills/nemo-mbridge-perf-memory-tuning/card.yaml

概要

学習中の GPU OOM は、メモリ容量そのものではなく、メモリの断片化が原因で発生することがよくあります。PyTorch のデフォルトの CUDA アロケーターでは、確保領域の間に利用できない隙間が残る場合があります。最も効果的な解決策は次のとおりです。

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

これにより、PyTorch は拡張可能な（固定サイズではない）メモリセグメントを使用するようになり、断片化を大幅に減らして、モデルや並列化の設定を変更せずに、限界に近い OOM を解消できることがよくあります。

断片化以外に、実際のピークメモリは次の要素で決まります。

• パラメーターとオプティマイザ状態のメモリ — TP、PP、DP のシャーディング（分散オプティマイザー、FSDP）で制御します
• 活性化メモリ — 活性化再計算、シーケンス長、マイクロバッチサイズ、PEFT 固有のギャザー済み入力の保持で制御します
• 一時メモリ／ワークスペースメモリ — CUDA カーネル、NCCL バッファー、CUDA グラフ

設定を計画する際は、大規模ジョブを起動する前に Bridge 理論推定ツールを使用してください。

from megatron.bridge.training.utils.theoretical_memory_utils import estimate_training_memory

estimate = estimate_training_memory(cfg, num_microbatches=num_microbatches)

推定ツールは、最も負荷の高い GPU シャードを報告し、dense／埋め込み、ルーティングされた MoE エキスパート、活性化の各要素を分けて表示します。アロケーターの断片化、CUDA／NCCL ワークスペース、CUDA グラフバッファー、トークンの不均衡、ディスパッチャーのワークスペースは含まれないため、実行時のメモリ指標で最終設定を検証してください。

すばやい判断

学習が OOM になる、またはメモリ上限に近い場合：

1. まず PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True を設定します。 パフォーマンスコストなしで、断片化による OOM を解消できます。ほとんどの Slurm 起動テンプレートには、すでに含まれています。
2. シーケンス並列化を使用する LoRA では、入力の再ギャザーを有効にします （LoRA(sequence_parallel_input_regather=True)）。これにより、対象レイヤーごとに LoRA-A のギャザー済み入力全体を保持する必要がなくなります。SP が無効な場合は効果がありません。
3. 選択的な活性化再計算を追加します（recompute_modules=[core_attn]）。まだ有効になっていない場合に適用してください。@skills/nemo-mbridge-perf-activation-recompute/SKILL.md を参照してください。
4. メモリ対策として TP を増やすのは避けてください — TP を 2 倍にすると NVLink の all-reduce 通信量が大幅に増え、スループットが低下することがよくあります（Llama3 70B で -28%）。
5. DP を減らして PP を増やすのは避けてください — DP を半分にすると勾配累積ステップが 2 倍になり、スループットが低下します（約 6%）。
6. それでも OOM になる場合は、mlp の再計算を検討してください。約 3 GB 節約できますが、大規模な dense モデル（Llama3 70B）では GPU 使用率が約 16% 低下します。
7. PP > 1 の場合、CPU オフロードは使用できません。

有効化

拡張可能なセグメント（最初に推奨する手順）

ジョブを起動する前に、環境変数に設定します。

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Slurm スクリプトでは、通常、ほかの環境変数と並べて設定します。

export CUDA_DEVICE_MAX_CONNECTIONS=1
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

モデル設定を変更する必要はありません。スループットへの影響もありません。

並列化の調整

モデルが実際にメモリに収まらない場合（断片化が原因ではない場合）は、並列化を調整してください。

戦略	メモリへの効果	スループットへのコスト	備考
PP を増やす（DP を維持）	ステージごとのレイヤー数が減少	中程度（DP を半分にした場合、約 6%）	GPU 数に余裕がある場合のみ
TP を増やす	GPU あたりのパラメーター数が減少	大きい（70B で -28%）	最終手段
分散オプティマイザー	DP ランク間でオプティマイザー状態をシャーディング	約 1～2%	大規模モデルに推奨
FSDP	パラメーター、勾配、オプティマイザーをシャーディング	変動します	@skills/nemo-mbridge-perf-megatron-fsdp/SKILL.md を参照してください

活性化再計算

詳しくは @skills/nemo-mbridge-perf-activation-recompute/SKILL.md を参照してください。

PEFT + シーケンス並列入力の再ギャザー

シーケンス並列を使用した LoRA トレーニングでは、対象となる列並列の linear_qkv および linear_fc1 アダプターは、ギャザーされた LayerNorm の出力を 入力として使用します。LoRA-A は学習可能なため、デフォルトの経路では、LoRA-A の重み勾配を 計算するために、ギャザー済みの入力全体をバックワードまで保持します。

PEFT 設定の作成時に入力の再ギャザーを有効にします。

from megatron.bridge.peft.lora import LoRA

cfg.peft = LoRA(
    # Keep the recipe's existing LoRA settings here.
    sequence_parallel_input_regather=True,
)

このオプションを有効にしても、フォワードでは LoRA-A の GEMM 用に入力全体を一時的に 実体化しますが、MCore の自動微分ではシーケンスのローカルシャードのみを保持します。バックワードでは 入力全体を非同期で再ギャザーし、可能な場合は dgrad と集合通信をオーバーラップさせて、 LoRA-A の重み勾配を計算した後、一時的な通信バッファーを再利用します。

これはメモリと通信のトレードオフであり、一般的な活性化チェックポイントではありません。 LayerNorm、アテンション、MLP、LoRA GEMM は再実行されません。スループットの低下が予想され、 対象となる LoRA-A の活性化の保持量が多いほど効果が大きくなります。シーケンス並列が無効な場合、 このオプションは効果がありません。

CPU オフロード

cfg.model.cpu_offloading = True

PP > 1 とは互換性がありません。 pipeline_model_parallel_size = 1 の場合にのみ使用できます。

VPP に関する注意

仮想パイプライン並列（VPP）は主にスループットを最適化する機能で、小さなモデルチャンクを インターリーブしてパイプラインのバブルによるオーバーヘッドを削減します。ピークメモリへの影響は ごくわずかです。VPP を変更しても、GPU 上の活性化、パラメーター、オプティマイザーの合計メモリ量は 実質的に変化しません。

以前の実験では、誤って VPP の調整（VPP 5→10）によって OOM が解消したと判断していました。 実際の修正は PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True で、これにより メモリの断片化が解消されました。VPP=10 の実行ではピークメモリが実際にはわずかに 増加（60.2 GB vs 58.8 GB）しましたが、expandable segments によって断片化が防止されたため、 OOM は発生しませんでした。

VPP はパイプラインのバブルを削減するために調整してください（@docs/parallelisms.md を参照）。 メモリ対策として使用するものではありません。

互換性と制約

• expandable_segments:True は --use-nccl-ub（NCCL ユーザーバッファー登録）と互換性がありません。 Megatron-FSDP のドキュメントを参照してください。
• CUDA グラフと expandable_segments:True を使用する場合は、 NCCL_GRAPH_REGISTER=0 を設定してください（Blackwell より前の GPU では必須であり、 MCore の CudaGraphManager によって強制されます）。
• CPU オフロードには pipeline_model_parallel_size = 1 が必要です。
• 分散オプティマイザーを使用するには、オプティマイザー設定で use_distributed_optimizer = True が必要です。
• sequence_parallel_input_regather は、対象となる非エキスパートの列並列 LoRA-A 投影にのみ適用されます。 行並列アダプター、エキスパートアダプター、TP=1、CUDA グラフ、CPU 活性化オフロード、 およびレイヤー全体または選択的な MLP 活性化再計算とのオーバーラップでは、既存の経路にフォールバックします。

測定結果

Llama3 70B SFT、32x H100 80GB、FP8（Current Scaling）:

• ベースライン: TP=4、PP=4、VPP=5、DP=2、MBS=1、GBS=32、seq_len=4096
• Golden GPU 使用率: 709.93 TFLOP/s/GPU
• 性能低下のしきい値: 5%

戦略の比較: メモリ削減のための並列化設定の変更

実験	TP	PP	VPP	DP	TFLOP/s/GPU	Golden 比	ピークメモリ (GB)	結果
ベースライン	4	4	5	2	~704	-0.8%	58.8	OOM（断片化）
PP 増加	4	8	5	1	668.0	-5.9%	53.2	性能は境界線上
TP 増加	8	4	5	1	508.7	-28.4%	50.2	大幅な性能低下
ベースライン + expandable_segments	4	4	5	2	~704	-0.8%	~59	合格

主なポイント:

• expandable_segments:True が最適です。 ベースラインの OOM は容量不足ではなく、 メモリの断片化が原因でした。この環境変数を設定すると、スループットを損なわず、 並列化設定を変更することなく OOM が解消しました。
• PP=8 はメモリには有効ですが、DP が失われます（2→1）。つまり、バッチあたりの 勾配累積ステップ数が 32 となり、スループットが約 6% 低下します。
• TP=8 は壊滅的です（-28%）。TP を 2 倍にすると NVLink 全体の all-reduce 通信量も 比例して増加し、DP=1 ではマイクロバッチのオーバーラップがなくなるためです。

CPU オフロード：利用不可

実験	offload_layers	結果
Exp 4	2	非対応（PP > 1）
Exp 5	4	非対応（PP > 1）
Exp 6	6	非対応（PP > 1）

ValueError: Currently there is no support for Pipeline parallelism with CPU offloading. CPU オフロードを使用するモデルでは、PP > 1 の構成は利用できません。

アクティベーション再計算：コストの高い代替手段

mlp を使用した選択的アクティベーション再計算により、ピークメモリを約 3 GB 削減できましたが、このワークロードでは GPU 使用率が約 16% 低下しました。詳細な結果については、@skills/nemo-mbridge-perf-activation-recompute/SKILL.md を参照してください。

LoRA + SP 入力の再ギャザー

SQuAD を使用した実チェックポイントでの H100 トレーニングでは、テストしたすべての構成でピークメモリが低下しましたが、スループットへの影響はワークロードによって異なりました。

モデル／構成	ベースラインのピーク	入力再ギャザーのピーク	削減メモリ	スループットの変化
Qwen3-8B, TP2, seq 8192	47.545 GB	42.814 GB	4.731 GB (10.0%)	-6.74%
Qwen3-30B-A3B, TP4/EP4	29.890 GB	28.321 GB	1.569 GB (5.2%)	-2.89%
GPT-OSS-120B, TP2/EP8	52.185 GB	51.371 GB	0.814 GB (1.6%)	-0.34%

すべての実行で損失は有限値となり、スキップされたイテレーションや NaN のイテレーションはありませんでした。2 ランクの BF16 と FP32 のチェックでは、出力、入力勾配、LoRA-A と LoRA-B の勾配、および 2 マイクロバッチの fused main_grad の累積がベースラインと一致しました。

コードの参照箇所

LoRA シーケンス並列入力の再ギャザー

src/megatron/bridge/peft/lora.py
    LoRA.sequence_parallel_input_regather

src/megatron/bridge/peft/utils.py
    ParallelLinearAdapter._sequence_parallel_input_regather_eligibility()
    ParallelLinearAdapter.forward()

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"
            )

VPP 構成とレイヤーの分割可能性の検証（MCore）

            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:
                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size
                if (
                    not num_layers_per_middle_pipeline_rank
                    % self.virtual_pipeline_model_parallel_size
                    == 0
                ):
                    raise ValueError(
                        f"number of layers on each middle pipeline rank:"
                        f"{num_layers_per_middle_pipeline_rank} must be divisible by virtual"
                        f"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}"
                    )

インターリーブド・パイプラインスケジュールに関する並列化のドキュメント

To minimize the pipeline bubble, the computation on each GPU can be divided into multiple subsets of layers (referred to as model chunks), rather than a single contiguous block. Enable this by setting `virtual_pipeline_model_parallel_size`:

model_config = GPTModelProvider(
    pipeline_model_parallel_size=4,
    virtual_pipeline_model_parallel_size=2,  # 2 model chunks per pipeline stage
    # ... other model parameters
)

失敗の診断

症状	原因	確認方法	修正方法
他のランクには余裕があるのに、単一ランクで OOM が発生する	メモリの断片化	expandable_segments:True が設定されているか確認します	PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True を設定します
expandable_segments を設定済みでも OOM が発生する	容量が実際に不足している	nvidia-smi でパラメーター／オプティマイザーのメモリを確認します	PP を増やす、分散オプティマイザーを使う、または再計算を追加します
起動前の推定メモリが GPU の容量を超える	モデルの状態またはアクティベーションが実際に大きすぎる	estimate_training_memory を実行し、最大の構成要素を確認します	起動前に PP/TP/CP/EP、分散オプティマイザー、または再計算を調整します
LoRA + SP でアクティベーションメモリが予想より多く残る	すべてのギャザー済み LoRA-A 入力がバックワードまで保持される	cfg.peft.sequence_parallel_input_regather が有効で、対象が適用可能か確認します	LoRA(sequence_parallel_input_regather=True) を設定し、フォールバックの制約を確認します
ValueError: PP + CPU offloading	cpu_offloading を PP > 1 で使用している	PP の構成を確認します	CPU オフロードを無効にするか、PP=1 に設定します
--use-nccl-ub + expandable segments で RuntimeError が発生する	NCCL UB が拡張可能なアロケーターと互換性がない	環境変数を確認します	expandable_segments:True を削除するか、--use-nccl-ub を無効にします

既知の制限事項

• PP > 1 の場合、CPU オフロードは利用できません
• 並列化のサイズ変更（TP/PP）には、多くの場合、大幅なスループット低下が伴います
• 理論的な推定器は数式ベースであり、実行時のプロファイリングや CUDA のメモリレポートに代わるものではありません
• LoRA 入力の再ギャザーは、行並列またはエキスパートアダプターには対応しておらず、メモリを大きく消費する対象の LoRA-A アクティベーションが少ない場合は、効果がほとんどない可能性があります

検証

expandable_segments:True が有効かどうかを手早く確認します。

import os
assert "expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")

Slurm ジョブの場合は、起動スクリプト内でトレーニングコマンドを実行する前に、環境変数がエクスポートされていることを確認します。

LoRA + SP 入力の再ギャザーについては、対象を絞った構成テストと、実際の 2 ランク MCore バックワード同等性テストを実行します。

uv run python -m pytest \
  tests/unit_tests/peft/test_utils.py -k "sequence_parallel_input_regather" \
  tests/unit_tests/peft/test_lora.py -k "sequence_parallel_input_regather"

uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest \
  tests/unit_tests/peft/test_lora_sp_input_regather_distributed.py