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

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 つ目のコマンドでプラグインが管理する環境変数の設定を検証できる