AutoGineer
Generated banner for Nemo Mbridge Perf Tp Dp Comm Overlap
Generated logo for Nemo Mbridge Perf Tp Dp Comm Overlap

Nemo Mbridge Perf Tp Dp Comm Overlap

NVIDIAOfficialPublished by NVIDIA

  • Skill
  • Tool integration
  • Software development
  • Apache-2.0
  • Medium risk
Download ZIP

From the author

Operational guide for enabling TP, DP, and PP communication overlap in Megatron-Bridge, including config knobs, code anchors, pitfalls, and verification.
— the description by NVIDIA

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides a comprehensive guide for enabling and tuning communication overlap (TP, DP, and PP) within the Megatron-Bridge framework. It covers essential configuration knobs, code-level gating mechanisms, common pitfalls to avoid, and verification procedures to ensure optimal performance during large-scale model training.

What it does

  • Provides configuration templates for enabling TP, DP, and PP communication overlap.
  • Identifies key code anchors that control overlap gating and P2P communication logic.
  • Highlights common pitfalls regarding sequence parallelism, Transformer Engine dependencies, and bucket sizing.
  • Details environment variable tuning for launch-time performance optimization.
  • Outlines verification steps using unit tests to confirm overlap configuration success.

When to use it

  • When configuring communication overlap for Megatron-Bridge training jobs.
  • When investigating throughput regressions related to communication overlap settings.
  • When troubleshooting why TP or PP overlap features are not activating as expected.
  • When tuning environment variables like CUDA_DEVICE_MAX_CONNECTIONS for performance.

Advantages

  • Centralizes configuration knowledge for complex parallel training setups.
  • Provides clear guidance on the interaction between model parallelism and overlap features.
  • Includes specific verification commands to validate configuration integrity.
  • Warns against common misconfigurations in mixed precision and environment settings.

Usage examples

“How do I enable TP communication overlap in my Megatron-Bridge config?”
The skill provides the CommOverlapConfig snippet to set tp_comm_overlap=True and notes the requirement for sequence_parallel=True and Transformer Engine availability.
“Why is my PP overlap not working?”
The guide explains that PP overlap is only auto-selected when pipeline_model_parallel_size > 1 and virtual_pipeline_model_parallel_size > 1.

Install

  1. Download the skill as a ZIP:

    https://autogineer.com/agent-skills/nvidia/nemo-mbridge-perf-tp-dp-comm-overlap/download

  2. Unzip it into the cross-client skills folder of your project:

    mkdir -p .agents/skills && unzip nemo-mbridge-perf-tp-dp-comm-overlap.zip -d .agents/skills

Documentation by NVIDIA

TP / DP / PP Communication Overlap Skill

For stable background and recommendation level, see:

  • @docs/training/communication-overlap.md

Enablement

Minimal Bridge override:

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

Optional TP preset:

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

Precision knobs belong to mixed precision:

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

Code Anchors

Bridge overlap gating:

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 overlap selection:

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 overlap defaults:

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

Launch-time env tuning:

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)

Pitfalls

  1. TP overlap silently disables itself if sequence_parallel=False or Transformer Engine is unavailable.

  2. PP overlap is not enabled for all PP cases. Bridge only auto-selects overlap_p2p_comm=True when PP > 1 and VPP > 1.

  3. bucket_size is a parameter-count knob, not a byte-size knob.

  4. grad_reduce_in_fp32 and fp8_param_gather should be set through mixed precision, not as standalone DDP tuning first.

  5. CUDA_DEVICE_MAX_CONNECTIONS and LayerNorm SM margin are launch-time plugin settings, not CommOverlapConfig fields.

Verification

Use the checked-in overlap unit coverage first:

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

Optional second check if nemo_run is available:

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

Success criteria:

  • first command reports 26 passed

  • second command validates plugin-owned env wiring when not skipped

Files

evals/evals.json

[
  {
    "id": "tp-dp-comm-overlap-positive-smoke",
    "question": "Use the nemo-mbridge-perf-tp-dp-comm-overlap skill. For Megatron Bridge with TP=4, sequence_parallel=True, PP=4, and VPP=2, what exact TP/DP/PP communication overlap settings should I enable and how should I verify Bridge wired them correctly?",
    "expected_skill": "nemo-mbridge-perf-tp-dp-comm-overlap",
    "expected_script": null,
    "ground_truth": "The answer should use the TP/DP/PP communication overlap skill. It should show CommOverlapConfig(tp_comm_overlap=True), require tensor_model_parallel_size > 1 and sequence_parallel=True for TP overlap, and set DDP overlap_grad_reduce=True plus overlap_param_gather=True with use_distributed_optimizer. It should explain that with PP > 1 and VPP > 1 Bridge selects overlap_p2p_comm=True and batch_p2p_comm=False, and recommend verifying tests or logs such as tests/unit_tests/training/test_comm_overlap.py.",
    "expected_behavior": [
      "Read the nemo-mbridge-perf-tp-dp-comm-overlap skill before answering.",
      "Identify TP, DP, and PP communication overlap as the target feature.",
      "List CommOverlapConfig(tp_comm_overlap=True).",
      "Mention sequence_parallel=True as a TP overlap requirement.",
      "List overlap_grad_reduce and overlap_param_gather for DP overlap.",
      "Explain PP overlap selection for PP > 1 and VPP > 1.",
      "Include a concrete verification path."
    ]
  }
]

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python

Security

What our scanner noticed· Medium risk

  • Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Published by
NVIDIA
Publisher
NVIDIA (Official)
Source
NVIDIA/skills · skills/nemo-mbridge-perf-tp-dp-comm-overlap · commit fce3cca
Licence
Apache License 2.0 (redistribution allowed)
  • LICENSE (added by AutoGineer from LICENSE-APACHE)
  • ATTRIBUTION.md (written by AutoGineer)
Synced

AutoGineer is not affiliated with or endorsed by NVIDIA or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download nemo-mbridge-perf-tp-dp-comm-overlap.zip