Skip to main content
skill CX Codex

nemo-mbridge-perf-tp-dp-comm-overlap — Claude Code skill: TP, DP, PP

A skill that walks through enabling TP, DP and PP communication overlap in Megatron-Bridge: which knobs to turn, where the code gates it, how to verify.

Install into Claude Code · Skills CLI (Vercel) All methods and platforms ↓
npx skills add NVIDIA/skills --skill nemo-mbridge-perf-tp-dp-comm-overlap -a claude-code
Open on GitHub

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

This skill covers turning on communication overlap in Megatron-Bridge. It ships a minimal config override, code anchors, a list of pitfalls and the commands that confirm the setup.

Overlap is configured through CommOverlapConfig: tp_comm_overlap=True, plus use_distributed_optimizer, overlap_grad_reduce and overlap_param_gather for DP. A ready TP preset is available as userbuffers_bf16_h100_h12288_tp4_mbs1_seqlen2048. Precision settings live in mixed precision instead: grad_reduce_in_fp32 and fp8_param_gather.

The trap worth knowing: TP overlap silently disables itself when sequence_parallel=False or Transformer Engine is missing. PP overlap is only auto-selected when PP > 1 and VPP > 1. Also, bucket_size counts parameters, not bytes. CUDA_DEVICE_MAX_CONNECTIONS and the LayerNorm SM margin are launch-time plugin settings, so check where you set them.

Verification runs uv run python -m pytest tests/unit_tests/training/test_comm_overlap.py -q, which should report 26 passed. A second command, tests/unit_tests/recipes/test_run_plugins.py, checks the plugin env wiring when nemo_run is available.

What it helps with

  • enable TP, DP and PP communication overlap in Megatron-Bridge
  • trace a throughput regression back to a comm overlap config change
  • find out why TP overlap turned itself off
  • set bucket_size for DP overlap without mixing up the units
  • confirm plugin env vars are wired at launch time

Example prompts for the agent

  • enable TP, DP and PP overlap in Megatron-Bridge and show me the minimal config
  • why did throughput drop after I changed the comm overlap settings
  • run test_comm_overlap.py and tell me what's wrong with my config

How to install nemo-mbridge-perf-tp-dp-comm-overlap

1Skills CLI (Vercel)

npx skills add NVIDIA/skills --skill nemo-mbridge-perf-tp-dp-comm-overlap -a claude-code

2Manually: put the files in place

~/.claude/skills/nemo-mbridge-perf-tp-dp-comm-overlap/SKILL.md

Copy the whole skill folder; for one project use .claude/skills/nemo-mbridge-perf-tp-dp-comm-overlap/

Step by step, with commands for each platform: How to install skills in Claude Code: five ways

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

How do I install nemo-mbridge-perf-tp-dp-comm-overlap in Claude?

Run npx skills add NVIDIA/skills --skill nemo-mbridge-perf-tp-dp-comm-overlap -a claude-code. The skill lands in ~/.claude/skills/nemo-mbridge-perf-tp-dp-comm-overlap/SKILL.md.

What does nemo-mbridge-perf-tp-dp-comm-overlap do?

It is an operational guide to TP, DP and PP communication overlap in Megatron-Bridge: config knobs, code anchors, pitfalls and verification.

Why does TP overlap stay off?

It disables itself when sequence_parallel=False or Transformer Engine is unavailable. Check both before editing the config.

Source text

SKILL.md · Apache-2.0
---
name: nemo-mbridge-perf-tp-dp-comm-overlap
description: Operational guide for enabling TP, DP, and PP communication overlap in Megatron-Bridge, including config knobs, code anchors, pitfalls, and verification.
license: Apache-2.0
when_to_use: Enabling TP/DP/PP comm overlap, or tracing a throughput regression to a comm overlap config change; 'overlap_param_gather', 'overlap_grad_reduce', 'sequence-parallel overlap', 'TP overlap', 'DP overlap', 'comm overlap'.
---

# TP / DP / PP Communication Overlap Skill

For stable background and recommendation level, see:

- @docs/training/communication-overlap.md

## Enablement

Minimal Bridge override:

```python
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:

```python
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:

```python
cfg.mixed_precision.grad_reduce_in_fp32 = False
cfg.mixed_precision.fp8_param_gather = False
```

## Code Anchors

Bridge overlap gating:

```439:449:src/megatron/bridge/training/comm_overlap.py
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:

```451:458:src/megatron/bridge/training/comm_overlap.py
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:

```572:579:src/megatron/bridge/training/comm_overlap.py
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:

```570:609:src/megatron/bridge/recipes/run_plugins.py
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:

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

Optional second check if `nemo_run` is available:

```bash
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

© NVIDIA · Apache-2.0 · NVIDIA/skills

Similar skills

Everything for AI and ML →

From the same set178

mcore-create-issue Investigates failing GitHub Actions runs for NVIDIA/Megatron-LM, extracts the error cause, and files a structured bug report. CX Codex +1 1,800installs mcore-run-on-slurm Prepare a SLURM launch for distributed Megatron-LM training, from an sbatch skeleton to rank-level failure checks. Useful when configuring or debugging a multi-node job. CX Codex +1 1,800installs nemo-rl-brev-etiquette A skill that tells the agent where files belong on a Brev box: code stays in the repo, heavy run output goes to /ephemeral. CX Codex +1 1,800installs mcore-linting-and-formatting Sets up linting and code formatting for Megatron-LM projects. It helps fix ruff and mypy violations and sort imports before opening a pull request. CX Codex +1 1,800installs nemo-mbridge-perf-moe-dispatcher-selection Picks between alltoall, DeepEP and HybridEP for MoE token dispatch based on your hardware, EP degree and optimization stage. CX Codex +1 1,800installs nemo-mbridge-perf-megatron-fsdp Configures Megatron FSDP in Megatron-Bridge to replace standard DDP, helping resolve out-of-memory errors and verify training checkpoints. CX Codex +1 1,800installs nemo-rl-docs This skill hands your agent the documentation rules for NeMo-RL: keeping docs/index.md current and writing docstrings in Google style. Useful when you touch the docs. CX Codex +1 1,800installs nemo-mbridge-resiliency A skill covering resiliency setup in Megatron Bridge: fault tolerance, straggler detection, preemption and re-run state machine. CX Codex +1 1,800installs

Project sponsors

Companies that keep this analytics open
Gold sponsor

Impressions, clicks and code copies — open statistics. We keep the numbers public for the same reason NeBlask itself is public. Become a sponsor