Hierarchical Context Parallel Skill

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