AutoGineer
Generated banner for Nemo Mbridge Perf Hierarchical Context Parallel
Generated logo for Nemo Mbridge Perf Hierarchical Context Parallel

Nemo Mbridge Perf Hierarchical Context Parallel

NVIDIAOfficialPublished by NVIDIA

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

From the author

Operational guide for enabling hierarchical context parallelism 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 the necessary configuration steps, code references, and verification methods for enabling hierarchical context parallelism (HCP) within the Megatron-Bridge framework. It focuses on the a2a+p2p communication type, detailing the requirements for hierarchical group sizing, dependency versions, and integration with Transformer Engine. The guide also highlights common pitfalls, such as compatibility issues with decentralized process groups and the importance of matching product sizes to total context parallel size.

What it does

  • Defines configuration knobs for hierarchical context parallelism.
  • Maps code anchors for MCore and Bridge initialization paths.
  • Provides validation logic for CP communication types.
  • Outlines requirements for Transformer Engine integration.
  • Lists common pitfalls and troubleshooting steps for HCP deployment.
  • Offers verification procedures via log inspection and manual smoke tests.

When to use it

  • When scaling context parallelism beyond standard KV head limits.
  • When investigating OOM errors or performance regressions related to CP configuration changes.
  • When implementing a2a+p2p communication strategies in Megatron-Bridge.
  • When debugging silent failures or broken gradients in older training stacks.

Advantages

  • Enables advanced scaling of sequence dimensions across GPU ranks.
  • Provides clear constraints to prevent configuration mismatches.
  • Includes specific log verification steps to confirm group initialization.
  • Offers a structured approach to testing hierarchical configurations.

Usage examples

“How do I enable hierarchical context parallelism with a size of 4 using a2a+p2p?”
The user receives the specific configuration overrides for model.context_parallel_size, model.cp_comm_type, and model.hierarchical_context_parallel_sizes.
“How can I verify that hierarchical context parallelism is active in my training logs?”
The user is instructed to check logs for the creation of HIERARCHICAL_CONTEXT_PARALLEL_GROUPS rather than just the standard CONTEXT_PARALLEL_GROUP.

Install

  1. Download the skill as a ZIP:

    https://autogineer.com/agent-skills/nvidia/nemo-mbridge-perf-hierarchical-context-parallel/download

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

    mkdir -p .agents/skills && unzip nemo-mbridge-perf-hierarchical-context-parallel.zip -d .agents/skills

Documentation by NVIDIA

Hierarchical Context Parallel Skill

This skill covers hierarchical context parallelism: nested context-parallel process groups used by cp_comm_type="a2a+p2p" and configured with hierarchical_context_parallel_sizes.

For what hierarchical CP is, when to use it, and the decision tree (a2a+p2p vs pure a2a vs p2p), see:

  • @docs/training/hierarchical-context-parallel.md

  • @skills/nemo-mbridge-perf-hierarchical-context-parallel/card.yaml

Enablement

Minimal Bridge override:

cfg.model.context_parallel_size = 4
cfg.model.cp_comm_type = "a2a+p2p"
cfg.model.hierarchical_context_parallel_sizes = [2, 2]
cfg.dist.use_decentralized_pg = False

Required constraints:

  • prod(hierarchical_context_parallel_sizes) == context_parallel_size

  • seq_length % (2 * context_parallel_size) == 0

  • Transformer Engine >= 1.12.0

Code Anchors

Upstream config and validation:

context_parallel_size: int = 1
"""Splits network input along sequence dimension across GPU ranks."""

hierarchical_context_parallel_sizes: Optional[list[int]] = None
"""Degrees of the hierarchical context parallelism. Users should provide a list to specify 
   the sizes for different levels. Taking the a2a+p2p cp comm type as example, it contains
   groups of two levels, so the first value of the list indicates the group size of the a2a
   communication type, and the second value indicates the group size of the p2p communication
   type.
"""
if args.hierarchical_context_parallel_sizes:
    from numpy import prod
    assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes)
if "a2a+p2p" in args.cp_comm_type:
    assert args.hierarchical_context_parallel_sizes is not None, \
    "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm"

Bridge MPU path:

parallel_state.initialize_model_parallel(
    ...
    context_parallel_size=model_config.context_parallel_size,
    hierarchical_context_parallel_sizes=model_config.hierarchical_context_parallel_sizes,
    ...
)
...
return ProcessGroupCollection.use_mpu_process_groups()

Bridge decentralized-PG path:

pg_collection = ProcessGroupCollection(
    ...
    cp=cp_pg,
    tp_cp=tp_cp_pg,
    hcp=None,
    ep=ep_pg,
    ...
)

Implementation Map

The code anchors above show the config declarations and argument validation.

Validation (MCore)

TransformerConfig.__post_init__ enforces that a2a+p2p requires HCP sizes and the product matches CP.

Process group creation

parallel_state.initialize_model_parallel creates hierarchical CP sub-groups when HCP sizes are provided via create_hierarchical_groups. Bridge currently gets those groups through the MPU-backed ProcessGroupCollection.

TE integration

TEDotProductAttention passes the hierarchical groups to Transformer Engine when a2a+p2p is used. Requires Transformer Engine >= 1.12.0.

Pitfalls

  1. Bridge HCP is MPU-only today: If use_decentralized_pg=True, Bridge initializes flat CP groups and leaves HCP unset.

  2. No checked-in Bridge recipe currently exercises HCP directly.

  3. Single-GPU load helpers clear hierarchical_context_parallel_sizes.

  4. Silent broken training on old stacks: If you use a2a+p2p without setting hierarchical_context_parallel_sizes, MCore now asserts. Older versions would silently disable CP communication, so each rank attended only to its local chunk and produced artificially high throughput with broken gradients.

  5. Product must match: prod(hierarchical_context_parallel_sizes) must exactly equal context_parallel_size. A mismatch triggers an assertion.

  6. Verify in logs: Look for the process group initialization output. You should see HIERARCHICAL_CONTEXT_PARALLEL_GROUPS being created. If you only see CONTEXT_PARALLEL_GROUP, HCP is not active.

Verification

No dedicated Bridge end-to-end test exists yet for HCP (see @skills/nemo-mbridge-perf-hierarchical-context-parallel/card.yaml follow_up_validation). Use the existing unit tests and log inspection instead.

Run the decentralized-PG unit test to confirm the flat-CP behavior is preserved:

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

For a manual smoke check, launch a 4-GPU run with a small recipe and cp_comm_type=a2a+p2p plus hierarchical_context_parallel_sizes=[2,2]:

CUDA_VISIBLE_DEVICES=0,1,2,3 uv run python -m torch.distributed.run --nproc_per_node=4 \
  scripts/training/run_recipe.py \
  --recipe llama32_1b_pretrain_config \
  model.context_parallel_size=4 \
  model.cp_comm_type=a2a+p2p \
  "model.hierarchical_context_parallel_sizes=[2,2]" \
  train.train_iters=2

Success criteria:

  • Logs show HIERARCHICAL_CONTEXT_PARALLEL_GROUPS being created

  • Training completes at least one step without error

  • If you only see CONTEXT_PARALLEL_GROUP, HCP is not active

Files

skill-card.md

Description: <br>

Operational guide for enabling hierarchical context parallelism in Megatron-Bridge, including config knobs, code anchors, pitfalls, and verification. <br>

This skill is ready for commercial/non-commercial use. <br>

Owner

NVIDIA <br>

License/Terms of Use: <br>

Apache 2.0 <br>

Use Case: <br>

Developers and engineers scaling context parallelism beyond KV heads in Megatron-Bridge distributed training, or investigating commits that changed CP config and caused OOM or regressions. <br>

Deployment Geography for Use: <br>

Global <br>

Known Risks and Mitigations: <br>

Risk: Review before execution as proposals could introduce incorrect or misleading guidance into skills. <br> Mitigation: Review and scan skill before deployment. <br>

Reference(s): <br>

Skill Output: <br>

Output Type(s): [Configuration instructions, Shell commands, Analysis] <br> Output Format: [Markdown with inline bash code blocks] <br> Output Parameters: [1D] <br> Other Properties Related to Output: [None] <br>

Evaluation Agents Used: <br>

  • Claude Code (claude-code) <br>

  • Codex (codex) <br>

Evaluation Tasks: <br>

Evaluated against 1 evaluation task (positive skill-activation case) with 2 attempts per task via NVSkills-Eval external profile. <br>

Evaluation Metrics Used: <br>

Reported benchmark dimensions: <br>

  • Security: Checks whether skill-assisted execution avoids unsafe behavior such as secret leakage, destructive commands, or unauthorized access. <br>

  • Correctness: Checks whether the agent follows the expected workflow and produces the correct final output. <br>

  • Discoverability: Checks whether the agent loads the skill when relevant and avoids using it when irrelevant. <br>

  • Effectiveness: Checks whether the agent performs measurably better with the skill than without it. <br>

  • Efficiency: Checks whether the agent uses fewer tokens and avoids redundant work. <br>

Underlying evaluation signals used in this run: <br>

  • security: Checks for unsafe operations, secret leakage, and unauthorized access. <br>

  • skill_execution: Verifies that the agent loaded the expected skill and workflow. <br>

  • skill_efficiency: Checks routing quality, decoy avoidance, and redundant tool usage. <br>

  • accuracy: Grades final-answer correctness against the reference answer. <br>

  • goal_accuracy: Checks whether the overall user task completed successfully. <br>

  • behavior_check: Verifies expected behavior steps, including safety expectations. <br>

  • token_efficiency: Compares token usage with and without the skill. <br>

Evaluation Results: <br>

DimensionNumclaude-codecodex
Security2100% (+0%)100% (+0%)
Correctness2100% (+0%)84% (+0%)
Discoverability2100% (+0%)59% (+0%)
Effectiveness296% (-2%)96% (+0%)
Efficiency293% (-0%)58% (+3%)

Testing Completed: <br>

[x] Agent Red-Teaming <br> [ ] Network Security <br> [ ] Product Security <br>

Skill Version(s): <br>

v0.2.0rc6-1528-gb0f64d72 (source: git describe) <br>

Ethical Considerations: <br>

NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal team to ensure this skill meets requirements for the relevant industry and use case and addresses unforeseen product misuse. <br>

(For Release on NVIDIA Platforms Only) <br> Please report quality, risk, security vulnerabilities or NVIDIA AI Concerns here. <br>

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-hierarchical-context-parallel · 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-hierarchical-context-parallel.zip