nemo-mbridge-perf-cpu-offloading

CPU オフロード

参考資料

• 安定版ドキュメント：@docs/training/cpu-offloading.md
• 構造化メタデータ：@skills/nemo-mbridge-perf-cpu-offloading/card.yaml

概要

GPU から CPU メモリへデータを移動する、独立した 2 つの仕組みです。

仕組み	設定名前空間	オフロード対象	PP の制限
アクティベーションのオフロード	model.cpu_offloading*	Transformer レイヤーごとのアクティベーション（およびオプションで重み）	PP は 1 でなければなりません
オプティマイザーのオフロード	optimizer.optimizer_cpu_offload	HybridDeviceOptimizer を介した Adam オプティマイザーの状態（モメンタム＋分散）	なし

すばやく選ぶ

状況	推奨
大規模な MoE モデル（30B 以上）で、PP > 1 が必要	オプティマイザーのオフロード — PP=1 の制約により、アクティベーションのオフロードは利用できません
小〜中規模モデルで、PP=1 で収まり、アクティベーションメモリが支配的	アクティベーションのオフロード
メモリと速度のトレードオフを調整したい	optimizer_offload_fraction を指定したオプティマイザーのオフロード
スループットを最優先する	有効にしないでください — オフロードには常にオーバーヘッドが伴います
CUDA グラフが必要	オプティマイザーのオフロードのみ — アクティベーションのオフロードとは互換性がありません
メモリ負荷が中程度	効率を最大化するには、オプティマイザーのオフロードを 25～50% の割合で設定します

有効化

オプティマイザーの CPU オフロード（大規模モデルに推奨）

cfg.optimizer.optimizer_cpu_offload = True
cfg.optimizer.optimizer_offload_fraction = 1.0
cfg.optimizer.overlap_cpu_optimizer_d2h_h2d = True

CLI オーバーライド：

optimizer.optimizer_cpu_offload=True \
optimizer.optimizer_offload_fraction=0.5 \
optimizer.overlap_cpu_optimizer_d2h_h2d=True

アクティベーションの CPU オフロード（小〜中規模モデルのみ）

cfg.model.cpu_offloading = True
cfg.model.cpu_offloading_num_layers = 16
cfg.model.cpu_offloading_activations = True
cfg.model.cpu_offloading_weights = False

cfg.model.pipeline_model_parallel_size = 1
cfg.model.recompute_granularity = None
cfg.model.cuda_graph_impl = "none"

設定パラメーターのリファレンス

オプティマイザーのオフロード

パラメーター	デフォルト	説明
optimizer_cpu_offload	False	マスタースイッチ
optimizer_offload_fraction	0.0	CPU 上に置くオプティマイザー状態の割合（0.0～1.0）
overlap_cpu_optimizer_d2h_h2d	False	GPU↔CPU 転送と計算をオーバーラップさせます
use_torch_optimizer_for_cpu_offload	False	CPU 部分で、融合オプティマイザーの代わりに torch.optim を使用します

アクティベーションのオフロード

パラメーター	デフォルト	説明
cpu_offloading	False	マスタースイッチ
cpu_offloading_num_layers	0	オフロードする Transformer レイヤーの数（0 から num_layers-1）
cpu_offloading_activations	True	アクティベーションをオフロードします
cpu_offloading_weights	False	重みをオフロードします
cpu_offloading_double_buffering	False	再読み込み時にレイヤー間でダブルバッファリングします

互換性と制約

アクティベーションのオフロード

• pipeline_model_parallel_size は 1 でなければなりません
• recompute_granularity は None でなければなりません
• fine_grained_activation_offloading とは併用できません
• CUDA グラフとは併用できません
• cpu_offloading_num_layers は [0, num_layers-1) の範囲内でなければなりません

オプティマイザーのオフロード

• use_distributed_optimizer = True が必要です（ほとんどのレシピでデフォルト）
• PP、再計算、CUDA グラフに関する制約はありません
• optimizer_offload_fraction は [0.0, 1.0] の範囲内でなければなりません

実践上の注意：大規模 MoE モデル

Qwen3-30B-A3B や同様の大規模 MoE モデルでは、アクティベーションのオフロードは利用できません。PP=1 の制約により、各 GPU が 48 レイヤーすべてを保持します。モデルの重みとオプティマイザー状態だけで約 70 GB に達し、H100 の 80 GB の容量を超えます。

最小限の実行可能なコマンド

uv run python scripts/training/run_recipe.py \
  --recipe qwen3_30b_a3b_pretrain_config \
  optimizer.optimizer_cpu_offload=True \
  optimizer.optimizer_offload_fraction=0.5 \
  train.train_iters=20 \
  train.global_batch_size=8 \
  train.micro_batch_size=1

検証

単体テスト

uv run python -m pytest \
  tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py -k "cpu_offload" \
  tests/unit_tests/peft/test_utils.py -k "cpu_offload" -q

成功条件

• 選択したオフロードモードで設定の検証に合格する
• OOM や NCCL エラーが発生せず、トレーニングが完了する
• 損失がオフロードなしのベースラインと一致する（最大差 < 0.001）
• オフロード率に比例してメモリ使用量が減少する

コードの参照箇所

MCore のアクティベーションオフロード制約

        if self.cpu_offloading and (
            self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers
        ):
            raise ValueError(...)

        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
            raise ValueError(
                "Currently there is no support for Pipeline parallelism with CPU offloading"
            )

        if self.cpu_offloading and self.recompute_granularity is not None:
            raise ValueError(
                "CPU offloading does not work when activation recomputation is enabled"
            )

MCore の CUDA グラフ非互換性

            if self.cpu_offloading:
                raise ValueError("CUDA graphs not supported with CPU offloading.")

MCore のきめ細かなオフロードとの相互排他性

        if self.fine_grained_activation_offloading:
            assert (
                not self.cpu_offloading
            ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading."

MCore の HybridDeviceOptimizer インスタンス化

        if config.optimizer_cpu_offload:
            # ... setup cpu/gpu optimizer classes ...
            optimizer = HybridDeviceOptimizer(
                param_groups,
                offload_fraction=config.optimizer_offload_fraction,
                cpu_optimizer_cls=cpu_optimizer_cls,
                gpu_optimizer_cls=gpu_optimizer_cls,
                overlap_cpu_optimizer_d2h_h2d=config.overlap_cpu_optimizer_d2h_h2d,
                pin_cpu_grads=config.pin_cpu_grads,
                pin_cpu_params=config.pin_cpu_params,
            )

Bridge の CUDA グラフガード

        assert not config.cpu_offloading and config.recompute_granularity is None, "Cudagraphs not supported"

PEFT における Bridge のアクティベーションオフロード

        if self.config.cpu_offloading and self.config.cpu_offloading_activations:
            x.activation_offloading = True
        x, _ = self.linear_in(x)
        x = self.activation(x)
        if self.config.cpu_offloading and self.config.cpu_offloading_activations:
            x.activation_offloading = True
        x, _ = self.linear_out(x)

障害の診断

症状	考えられる原因	確認方法	対処方法
Currently there is no support for Pipeline parallelism with CPU offloading	Activation offload と PP > 1 の併用	pipeline_model_parallel_size を確認します	PP=1 に設定するか、オプティマイザーのオフロードを使用します
CPU offloading does not work when activation recomputation is enabled	Activation offload と再計算の併用	recompute_granularity を確認します	recompute_granularity=null に設定します
fine_grained_activation_offloading cannot be enabled with cpu_offloading	両方のオフロードモードが有効	両方のフラグを確認します	どちらか一方を使用します
CUDA graphs not supported with CPU offloading	CUDA グラフと activation offload の併用	cuda_graph_impl を確認します	cuda_graph_impl="none" に設定します
Activation offloading 使用時に OOM が発生する	PP=1 ではモデルが大きすぎます	割り当て済みメモリと 80 GB を比較します	PP > 1 でオプティマイザーのオフロードを使用します
極端に遅くなる（4 倍超）	オプティマイザーのオフロード率が 100% で、CPU Adam がボトルネックになっています	異なるオフロード率でイテレーション時間を比較します	オフロード率を下げるか、overlap_cpu_optimizer_d2h_h2d を有効にします
部分的なオプティマイザーのオフロードで OOM が発生する	この設定ではオフロードが不十分です	異なるオフロード率でメモリ使用量を確認します	オフロード率を上げるか、PP を追加します

既知の制限事項

• Activation offloading には PP=1 が必要なため、パイプライン並列処理が必要な大規模モデル （30B 以上の MoE）では実用的ではありません。
• オプティマイザーのオフロードによるスループット低下は、オフロード率に比例して大きくなります（Qwen3-30B-A3B では 25% で約 1.9 倍、 100% で約 4.2 倍）。
• D2H/H2D のオーバーラップによる高速化は約 7% にとどまります。これは、CPU Adam の計算が 主なボトルネックだからです。
• fine_grained_activation_offloading はモジュール単位の独立したアプローチであり、 PP > 1 でも機能しますが、レイヤー単位の cpu_offloading と併用できません。