SuperAPI AI 工具集

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

通用与其它

查看项目

更新于 10-16 04:33

技能介绍

用途:在 Megatron-Bridge 中启用 TP(张量并行)、DP(数据并行)、PP(流水并行)的通信与计算重叠,减少通信开销对训练吞吐的影响。资料主要提供操作指南、配置开关、代码锚点、陷阱和验证方式,未提供具体性能收益数据。

主要能力:通过 CommOverlapConfig 设置 tp_comm_overlap=True 启用 TP 重叠;设置 DDP 的 overlap_grad_reduce、overlap_param_gather 并启用分布式优化器实现 DP 重叠;PP 重叠由 Bridge 在 PP>1 且 VPP>1 时自动选择 overlap_p2p_comm=True,否则用 batch_p2p_comm=True。TP 可选预设如 userbuffers_bf16_h100_h12288_tp4_mbs1_seqlen2048。精度相关开关 grad_reduce_in_fp32、fp8_param_gather 需通过 mixed_precision 设置,而不是单独作为 DDP 调参。启动时可通过 executor.env_vars 设置 CUDA_DEVICE_MAX_CONNECTIONS 及 NVTE_FWD/BWD_LAYERNORM_SM_MARGIN。

使用方式:最小覆盖示例为设置 cfg.model.tensor_model_parallel_size、sequence_parallel、pipeline_model_parallel_size、virtual_pipeline_model_parallel_size,然后 cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=True),并配置 ddp.use_distributed_optimizer=True、overlap_grad_reduce=True、overlap_param_gather=True。

限制与陷阱:sequence_parallel=False 或 Transformer Engine 不可用时 TP 重叠会静默失效;PP 重叠并非所有 PP 场景都启用,仅 PP>1 且 VPP>1 时自动选择 overlap_p2p_comm=True;bucket_size 是参数量旋钮而非字节大小旋钮;CUDA_DEVICE_MAX_CONNECTIONS 与 LayerNorm SM margin 是启动期插件设置,不属于 CommOverlapConfig 字段。

验证:先运行 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 验证插件环境变量接线(未跳过时)。