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）で制御します
• アクティベーションメモリ — アクティベーションの再計算、シーケンス長、マイクロバッチサイズで制御します
• 一時／ワークスペースメモリ — 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 シャードを示し、密層／埋め込み、ルーティングされる MoE エキスパート、アクティベーションの各要素を個別に表示します。アロケーターの断片化、CUDA／NCCL ワークスペース、CUDA グラフバッファー、トークンの偏り、ディスパッチャーのワークスペースは含まれないため、最終的な設定は実行時のメモリ指標で検証してください。

すぐにできる判断

トレーニングの実行中に OOM が発生する場合、またはメモリ上限に近い場合：

1. まず PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True を設定します。 断片化が原因の OOM を、パフォーマンスコストなしで解消できます。ほとんどの Slurm 起動テンプレートには、すでに含まれています。
2. まだ有効にしていない場合は、選択的なアクティベーション再計算（recompute_modules=[core_attn]）を追加します。@skills/nemo-mbridge-perf-activation-recompute/SKILL.md を参照してください。
3. メモリ対策として TP を増やすことは避けてください。TP を 2 倍にすると NVLink の all-reduce 通信量が大幅に増え、スループットが低下することがよくあります（Llama3 70B では -28%）。
4. DP を減らして PP を増やすことは避けてください。DP を半分にすると勾配累積ステップが 2 倍になり、スループットが低下します（約 6%）。
5. それでも OOM が発生する場合は、mlp の再計算を検討してください。約 3 GB を節約できますが、大規模な密層モデル（Llama3 70B）では GPU 使用率が約 16% 低下します。
6. 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 を参照してください。

CPU オフロード

cfg.model.cpu_offloading = True

PP が 1 より大きい場合は使用できません。 pipeline_model_parallel_size = 1 の場合のみ使用できます。

VPP に関する注意

仮想パイプライン並列処理（VPP）は、主にスループットを最適化する手法です。小さなモデルチャンクをインターリーブすることで、パイプラインのバブルによるオーバーヘッドを削減します。ピークメモリへの影響はごくわずかで、VPP を変更しても、GPU 上のアクティベーション、パラメーター、オプティマイザーの総メモリ使用量は実質的に変わりません。

以前の実験では、OOM の解消を誤って VPP の調整（VPP 5→10）によるものと判断していました。実際の解決策は PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True で、メモリの断片化を解消しました。VPP=10 の実行では、実際にはピークメモリをわずかに多く使用していました（60.2 GB 対 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 が必要です。

測定結果

32 台の H100 80GB、FP8 での Llama3 70B SFT（現在のスケーリング）：

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

戦略の比較：メモリ削減のための並列処理の変更

実験	TP	PP	VPP	DP	TFLOP/s/GPU	ゴールデン比	ピークメモリ（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. PP > 1 を使用するモデルでは、この方法は利用できません。

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

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

コード参照箇所

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、分散オプティマイザー、または再計算を調整します
ValueError: PP + CPU offloading	PP > 1 で cpu_offloading を使用しています	PP の設定を確認します	CPU オフロードを無効にするか、PP=1 に設定します
--use-nccl-ub と expandable segments の併用時に RuntimeError が発生する	NCCL UB は拡張可能なアロケーターと互換性がありません	環境変数を確認します	expandable_segments:True を削除するか、--use-nccl-ub を無効にします

既知の制限事項

• PP > 1 の場合、CPU オフロードは使用できません
• 並列処理のサイズ変更（TP／PP）には、大幅なスループット低下が伴うことがよくあります
• 理論上の推定値は数式に基づいており、実行時のプロファイリングや CUDA メモリレポートの代わりにはなりません

検証

expandable_segments:True が有効であることを簡単に確認します。

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

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