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

card.yaml

title: hierarchical_context_parallel
validated_on: "2026-03-14"
summary: >
  Megatron-Bridge currently supports hierarchical context parallelism
  (`cp_comm_type="a2a+p2p"` plus `hierarchical_context_parallel_sizes`) only
  through the MPU initialization path. The decentralized process-group path
  remains flat and does not create hierarchical CP groups.
validation_status:
  upstream_a2a_p2p_core:
    - code_verified
  bridge_mpu_passthrough:
    - code_verified
  bridge_mpu_runtime_groups:
    - code_verified
  bridge_decentralized_pg_hcp:
    - code_verified
  bridge_hcp_recipes_examples:
    - unclear
  bridge_hcp_docs:
    - doc_only
  bridge_hcp_end_to_end_training:
    - unclear
feature_meaning:
  a2a_p2p: >
    Megatron-Core hierarchical context-parallel transport path used by
    Transformer Engine attention and enabled by cp_comm_type="a2a+p2p".
  hierarchical_context_parallel_sizes: >
    Per-level subgroup sizes for hierarchical context parallelism. The product
    must equal context_parallel_size.
  inner_outer_cp_groups: >
    Hierarchical CP creates inner and outer context-parallel groups. With
    a2a+p2p, the inner group uses a2a and the outer group uses p2p.
recommended_path:
  model.context_parallel_size: 4
  model.cp_comm_type: a2a+p2p
  model.hierarchical_context_parallel_sizes:
    - 2
    - 2
  dist.use_decentralized_pg: false
known_constraints:
  - Transformer Engine must be >= 1.12.0 for a2a+p2p.
  - hierarchical_context_parallel_sizes must be set when cp_comm_type contains a2a+p2p.
  - The product of hierarchical_context_parallel_sizes must equal context_parallel_size.
  - seq_length must be divisible by 2 * context_parallel_size when CP > 1.
  - Bridge HCP is MPU-path only today.
known_limitations:
  - The decentralized-PG path initializes flat CP groups and leaves HCP unset.
  - No checked-in Bridge recipe sets cp_comm_type=a2a+p2p.
  - No checked-in Bridge functional test runs an end-to-end HCP training step.
evidence:
  - src/megatron/bridge/training/initialize.py
  - src/megatron/bridge/training/config.py
  - src/megatron/bridge/training/model_load_save.py
  - docs/performance-guide.md
  - tests/unit_tests/training/test_decentralized_pg.py
  - 3rdparty/Megatron-LM/megatron/core/model_parallel_config.py
  - 3rdparty/Megatron-LM/megatron/core/parallel_state.py
  - 3rdparty/Megatron-LM/megatron/core/extensions/transformer_engine.py
  - 3rdparty/Megatron-LM/megatron/training/arguments.py
follow_up_validation:
  - Add a positive Bridge functional test that completes at least one HCP training step.
  - Add Bridge-side validation that rejects HCP-looking config on use_decentralized_pg=true.
  - Add a checked-in Bridge recipe or example for a2a+p2p.
  - Validate model-family-specific HCP correctness beyond group initialization.

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