Nemo-MBridgeTP/DP/PP通信重叠性能调优Skill nemo-mbridge-perf-tp-dp-comm-overlap

名称	nemo-mbridge-perf-tp-dp-comm-overlap
描述	在Megatron-Bridge中启用TP、DP和PP通信重叠的操作指南，包括配置项、代码锚点、注意事项和验证方法。
开源协议	Apache-2.0 when_to_use: 启用TP/DP/PP通信重叠，或将吞吐量回归追溯到通信重叠配置更改时；‘overlap_param_gather’，‘overlap_grad_reduce’，‘序列并行重叠’，‘TP重叠’，‘DP重叠’，‘通信重叠’。

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 仅在 PP > 1 且 VPP > 1 时自动选择 overlap_p2p_comm=True。
3. bucket_size 是参数数量控制的旋钮，而不是字节大小的旋钮。
4. grad_reduce_in_fp32 和 fp8_param_gather 应通过混合精度设置，而不是首先作为独立的 DDP 调优项。
5. CUDA_DEVICE_MAX_CONNECTIONS 和 LayerNorm SM margin 是启动时插件设置，不是 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

成功标准：

• 第一个命令报告 26 passed
• 第二个命令在未跳过时验证插件拥有的环境变量接线