nemo-mbridge-perf-hierarchical-context-parallel

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()