用途:本技能是 Megatron-Bridge 中启用 TP、DP、PP 通信重叠的操作指南,帮助在分布式训练中让通信与计算重叠,并给出配置项、代码锚点、陷阱与验证方式。
主要能力:通过 CommOverlapConfig 启用 TP 通信重叠;通过 DDP 配置启用梯度归约与参数收集重叠;PP 重叠由 Bridge 根据流水线并行与虚拟流水线并行自动选择。还提供可选 TP 预设、混合精度相关开关,以及启动时环境变量调优。
使用方式:最小覆盖示例为设置 tensor_model_parallel_size、sequence_parallel、pipeline_model_parallel_size、virtual_pipeline_model_parallel_size,并创建 CommOverlapConfig(tp_comm_overlap=True),同时设置 ddp.use_distributed_optimizer、overlap_grad_reduce、overlap_param_gather。可选 TP 预设如 userbuffers_bf16_h100_h12288_tp4_mbs1_seqlen2048。精度开关 grad_reduce_in_fp32 与 fp8_param_gather 属于混合精度配置。启动环境变量包括 CUDA_DEVICE_MAX_CONNECTIONS、NVTE_FWD_LAYERNORM_SM_MARGIN、NVTE_BWD_LAYERNORM_SM_MARGIN。
验证:先运行 uv run python -m pytest tests/unit_tests/training/test_comm_overlap.py -q,成功标准为 26 passed;若 nemo_run 可用,可再运行 uv run python -m pytest tests/unit_tests/recipes/test_run_plugins.py -q 验证插件环境接线。
限制与陷阱:sequence_parallel=False 或缺少 Transformer Engine 时 TP 重叠会静默禁用;PP 重叠仅在 PP>1 且 VPP>1 时自动选择 overlap_p2p_comm=True;bucket_size 是参数数量旋钮而非字节大小;grad_reduce_in_fp32 与 fp8_param_gather 应通过混合精度设置;CUDA_DEVICE_MAX_CONNECTIONS 与 LayerNorm SM margin 是启动时插件设置,不是 CommOverlapConfig 字段。