mcpField
N
nemo-mbridge-perf-moe-long-context

nemo-mbridge-perf-moe-long-context

github.com/nvidia開発3.8k1.8k 回インストール今日更新
リポジトリ

nemo-mbridge-perf-moe-long-context について

長文脈の MoE モデル学習を調整するエンジニア向けのスキルです。並列化、再計算、精度、オフロードの選択肢を、複数のモデルとハードウェア構成例とともに説明します。

主な機能

  1. 01コンテキスト並列度の目安
  2. 02選択的再計算の推奨
  3. 03精度とオプティマイザーのオフロード
  4. 04H100 と GB200 の構成例
  5. 05CUDA Graph の利用指針

活用シーン

  1. 01長文脈 MoE 学習の計画
  2. 02CP、TP、EP、PP の設定選択
  3. 03活性化またはオプティマイザーのメモリー削減
  4. 04必要な GPU 数の実現可能性評価

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/nvidia/skills --skill nemo-mbridge-perf-moe-long-context

インストール後、Claude Code で/nemo-mbridge-perf-moe-long-contextを入力して読み込まれたか確認します。

同梱ファイル

SKILL.md と一緒にインストールされる参考資料やスクリプト。

BENCHMARK.md
card.yaml
skill-card.md
skill.oms.sig
evals/
evals.json

詳細

MoE ロングコンテキストトレーニング

安定版ドキュメント：@docs/training/moe-optimization.md カード：@skills/nemo-mbridge-perf-moe-long-context/card.yaml

ロングコンテキストで何が変わるか

シーケンス長が4Kクラスの範囲を大きく超えると、アテンションメモリと アクティベーションの常駐が主な制約になります。MoE モデルでは、通常、 次のいくつかを組み合わせる必要があります。

  • コンテキスト並列化
  • 選択的再計算
  • 低精度化
  • オプティマイザー状態の CPU オフロード
  • 残された少ない DP 予算を無駄にしないディスパッチャーと PP レイアウト

概算のスケーリング傾向

H100 上の DSV3

DSV3 のロングコンテキスト実行では、安定した傾向が見られます。

  • 最短のコンテキストを超えると、完全再計算より選択的再計算のほうが効果的です
  • CP を適切に増やせば、中程度の長さから非常に長いコンテキストまで、スループットはかなり狭い範囲に収まります
  • CP が増えるにつれて、トレードオフの焦点は「メモリに収まるか」から「必要な GPU 台数を確保できるか」へ移ります

つまり、適切なレイアウトを選べば、ロングコンテキストですぐに使用率が低下するわけではありませんが、DP 予算は急速に消費されます。

GB200 上の Qwen3-Next

Qwen3-Next は、メモリに敏感な中規模モデルに近い動作をします。

  • 8K と 32K は、適度な CP で引き続き実用的です
  • 64K も可能ですが、スループットの低下が目立ち、メモリの余裕も大幅に小さくなります
  • パイプラインレイアウトと Grouped GEMM の改善は、CP とほぼ同じくらい重要です
GB200 上の Qwen3 235B

Qwen3 235B は、TP、CP、HybridEP を連携させれば、NVL72 システムでロングコンテキストを効率的に処理できることを示しています。最適な 128K クラスの構成は、単に「収まるだけ」の設定ではありません。ルーティング、並列化、再計算のバランスを取れば、高い効率を維持できます。

CP の目安

  1. 4K のシャードサイズを目標に始める：最初の目安としては CP ~= seq_len / 4096 とし、実用的な2の累乗のレイアウトに丸めます。

  2. 可能であれば DP を維持する：CP、EP、TP、PP によって DP が限界まで減ると、ロングコンテキストのスケーリングは不安定になります。

  3. 選択的再計算を優先する：完全再計算を選ぶ前に、up_proj、norm、 moe、moe_act、mlp などのモジュールを再計算します。

  4. 非常に長いコンテキストでは SDPA を多用する再計算を避ける：アテンションの内部処理を再計算すると、より小さな MoE や MLP 側のモジュールを再計算する場合に比べて、メモリ削減効果が小さい一方で、処理量が大幅に増えることがあります。

  5. NVL72 システムでは TP も別の調整手段として使う：GB200 および GB300 の実行では、効率を維持しながら CP を一部 TP に置き換えられる場合があります。

  6. GBS を小さくする必要があると想定する：CP が増えて DP が減るにつれて、グローバルバッチサイズを小さくするか、GA の増加を受け入れる必要が生じることがあります。

代表的な構成ファミリー

H100 上の DSV3、128K
TP=1  CP=32  EP=32  PP=8  VPP=4
Precision: FP8-class
Dispatcher: DeepEP
Recompute: up_proj, norm, moe, mlp
Extra memory help: optimizer CPU offload
H100 上の DSV3、256K
TP=1  CP=64  EP=32  PP=8  EDP=2  VPP=4
Precision: FP8-class
Dispatcher: DeepEP
Recompute: up_proj, norm, moe, mlp
Extra memory help: optimizer CPU offload
GB200 上の Qwen3 235B、128K
TP=4  CP=4  EP=32  PP=4  VPP=12
Precision: BF16 or MXFP8
Dispatcher: HybridEP
Recompute: moe_act, norm
CUDA Graph: attn + moe_router + moe_preprocess

再計算と CUDA グラフのガイダンス

ロングコンテキストの MoE トレーニングでは、次のようにします。

  • まず選択的再計算から始めます
  • CUDA グラフを追加するのは、形状とルーティング経路が安定してからにします
  • CUDA グラフを使用する場合は、シーケンス長と MBS を固定します
  • 実行が大きく変動するバッチに依存する場合は、Eager 実行を優先します

参考資料：

  • @docs/training/activation-recomputation.md
  • @skills/nemo-mbridge-perf-cuda-graphs/SKILL.md

注意点

  1. CP は EP や PP の代わりにはなりません：並列化の次元が増えるのであって、ほかの次元が不要になるわけではありません。

  2. 4K の優れたベースラインが、ロングコンテキストでも優れたベースラインとは限りません：ルーティング方式、再計算の選択、オフロード戦略を変更する必要があることがよくあります。

  3. 必要な GPU 台数の確保が実際の制約になります：非常に長いコンテキストは単一の構成では問題なく見えても、モデル全体で EP と PP を適切に加えると実現不可能になることがあります。

  4. CUDA グラフには静的な形状が必要です：可変長バッチや状況に応じたパディング戦略によって、処理経路が気付かないうちに壊れることがあります。

  5. 128K 以上では、コンテナーとカーネルのサポートがより重要です：ロングコンテキストの処理経路は、短いコンテキストの立ち上げよりも、新しいカーネルやバグ修正に依存する傾向があります。

コメント

Markdown 対応