mcpField
N
nemo-mbridge-perf-tp-dp-comm-overlap

nemo-mbridge-perf-tp-dp-comm-overlap

github.com/promptingcompany開発039 回インストール17 週間前に更新
リポジトリ

nemo-mbridge-perf-tp-dp-comm-overlap について

Megatron Bridge の学習設定で通信オーバーラップを構成するためのガイドです。TP、PP、DP の有効化条件や既定の動作、検証方法を扱い、分散学習設定を調整する開発者を対象としています。

主な機能

  1. 01TP 通信オーバーラップの要件
  2. 02PP オーバーラップの選択ロジック
  3. 03DP の勾配集約とパラメーター収集の既定値
  4. 04起動時環境変数の調整に関する説明
  5. 05ユニットテストによる検証手順

活用シーン

  1. 01TP 通信オーバーラップの設定
  2. 02PP オーバーラップの有効状態確認
  3. 03DP 通信設定のデバッグ
  4. 04オーバーラップ設定と起動プラグインの検証

インストール

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

$npx skills add https://github.com/promptingcompany/nv-skills --skill nemo-mbridge-perf-tp-dp-comm-overlap

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

同梱ファイル

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

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

詳細

TP／DP／PP 通信オーバーラップスキル

安定した背景情報と推奨レベルについては、次を参照してください。

  • @docs/training/communication-overlap.md

有効化

最小限の Bridge オーバーライド：

from megatron.bridge.training.comm_overlap import CommOverlapConfig

cfg.model.tensor_model_parallel_size = 4
cfg.model.sequence_parallel = True
cfg.model.pipeline_model_parallel_size = 4
cfg.model.virtual_pipeline_model_parallel_size = 2

cfg.comm_overlap = CommOverlapConfig(
    tp_comm_overlap=True,
)

cfg.ddp.use_distributed_optimizer = True
cfg.ddp.overlap_grad_reduce = True
cfg.ddp.overlap_param_gather = True

オプションの TP プリセット：

from megatron.bridge.training.comm_overlap import userbuffers_bf16_h100_h12288_tp4_mbs1_seqlen2048

cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_bf16_h100_h12288_tp4_mbs1_seqlen2048

精度調整ノブは混合精度に属します。

cfg.mixed_precision.grad_reduce_in_fp32 = False
cfg.mixed_precision.fp8_param_gather = False

コードの参照箇所

Bridge のオーバーラップ制御：

if self.user_comm_overlap_cfg.tp_comm_overlap is True:
    if model_cfg.tensor_model_parallel_size < 2:
        ...
    elif not model_cfg.sequence_parallel:
        ...
    elif not HAVE_TE:
        ...

PP オーバーラップの選択：

if model_cfg.pipeline_model_parallel_size > 1:
    if vp_size > 1:
        comm_overlap_cfg.overlap_p2p_comm = True
        comm_overlap_cfg.batch_p2p_comm = False
    else:
        comm_overlap_cfg.overlap_p2p_comm = False
        comm_overlap_cfg.batch_p2p_comm = True

DP オーバーラップのデフォルト設定：

if self.data_parallel_size > 1:
    comm_overlap_cfg.bucket_size = 128 * 1024 * 1024
    comm_overlap_cfg.overlap_grad_reduce = True
    comm_overlap_cfg.overlap_param_gather = True

起動時の環境変数チューニング：

executor.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] = str(cuda_device_max_connections)
...
executor.env_vars["NVTE_FWD_LAYERNORM_SM_MARGIN"] = str(self.layernorm_sm_margin)
executor.env_vars["NVTE_BWD_LAYERNORM_SM_MARGIN"] = str(self.layernorm_sm_margin)

注意点

  1. sequence_parallel=False の場合、または Transformer Engine が利用できない場合、TP オーバーラップは警告なしに無効になります。
  2. PP オーバーラップはすべての PP ケースで有効になるわけではありません。Bridge が overlap_p2p_comm=True を自動選択するのは、PP > 1 かつ VPP > 1 の場合のみです。
  3. bucket_size はパラメーター数の調整ノブであり、バイト数の調整ノブではありません。
  4. grad_reduce_in_fp32 と fp8_param_gather は、単独で DDP をチューニングするのではなく、混合精度を通じて設定してください。
  5. CUDA_DEVICE_MAX_CONNECTIONS と LayerNorm SM マージンは起動時のプラグイン設定であり、CommOverlapConfig のフィールドではありません。

検証

まず、リポジトリに含まれるオーバーラップのユニットテストを実行してください。

uv run python -m pytest tests/unit_tests/training/test_comm_overlap.py -q

nemo_run が利用可能な場合の、任意の追加チェック：

uv run python -m pytest tests/unit_tests/recipes/test_run_plugins.py -q

成功条件：

  • 1 つ目のコマンドで 26 passed と表示される
  • スキップされていない場合、2 つ目のコマンドでプラグインが管理する環境変数の設定を検証できる

コメント

Markdown 対応