Operational guide for enabling hierarchical context parallelism in Megatron-Bridge, including config knobs, code anchors, pitfalls, and verification.
This NVIDIA skill is an operational guide for enabling hierarchical context parallelism (CP) in Megatron-Bridge. It solves the problem of scaling context parallelism beyond the number of key-value heads by using nested context-parallel process groups driven by the a2a+p2p communication type and configured with hierarchical_context_parallel_sizes. It gives engineers the exact config knobs, upstream code anchors, pitfalls, and verification steps to do this correctly.
The guide provides a minimal Bridge override (context_parallel_size, cp_comm_type set to a2a+p2p, the hierarchical sizes list, and decentralized process-group flag) and the required constraints: the product of the hierarchical sizes must equal the context parallel size, sequence length must be divisible by twice the context parallel size, and Transformer Engine must be 1.12.0 or newer. It cites specific upstream source ranges in Megatron-LM and Bridge to ground the config declarations, argument validation, process-group creation, and Transformer Engine integration. A pitfalls section warns about MPU-only support today, the lack of a checked-in recipe exercising HCP, single-GPU load helpers clearing the sizes, and a subtle correctness trap where older stacks silently disabled CP communication (producing artificially high throughput with broken gradients) versus the newer assertion. It is Apache-2.0 licensed with published evaluation results.
It targets ML engineers and performance specialists scaling long-context training in Megatron-Bridge, or investigating a commit that changed CP config and caused an out-of-memory error or regression. It produces configuration guidance and analysis with code references; it performs no destructive or privileged actions.
When scaling context parallelism beyond the number of KV heads, using nested CP process groups with the a2a+p2p communication type configured via hierarchical_context_parallel_sizes.
The product of hierarchical_context_parallel_sizes must equal context_parallel_size, sequence length must be divisible by twice the context parallel size, and Transformer Engine must be at least 1.12.0.
Yes. Hierarchical CP is MPU-only today (with use_decentralized_pg=True it falls back to flat CP groups), no checked-in Bridge recipe exercises it directly, and single-GPU load helpers clear the hierarchical sizes.
Older versions would silently disable CP communication if a2a+p2p was used without setting the hierarchical sizes, so each rank attended only to its local chunk, yielding artificially high throughput but broken gradients. Newer MCore asserts instead.
Check the process-group initialization output in the logs for the hierarchical context parallel group entries.
Quick Setup:
.claude/skills/Repository
nvidia/skills