Skill Context Parallèle Hiérarchique

Cette skill couvre le parallélisme de contexte hiérarchique : groupes de processus parallèles en contexte imbriqués utilisés par cp_comm_type="a2a+p2p" et configurés avec hierarchical_context_parallel_sizes.

Pour comprendre ce qu'est le CP hiérarchique, quand l'utiliser et l'arbre de décision (a2a+p2p vs pur a2a vs p2p), voir :

@docs/training/hierarchical-context-parallel.md
@skills/nemo-mbridge-perf-hierarchical-context-parallel/card.yaml

Enablement

Override Bridge minimal :

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

Contraintes requises :

prod(hierarchical_context_parallel_sizes) == context_parallel_size
seq_length % (2 * context_parallel_size) == 0
Transformer Engine >= 1.12.0

Code Anchors

Config amont et 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

Config definition

hierarchical_context_parallel_sizes est déclaré dans ModelParallelConfig :

# 3rdparty/Megatron-LM/megatron/core/model_parallel_config.py
hierarchical_context_parallel_sizes: Optional[list[int]] = None
# For a2a+p2p, first value = a2a group size, second value = p2p group size.
# Product must equal context_parallel_size.

cp_comm_type est déclaré dans TransformerConfig :

# 3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
cp_comm_type: Optional[Union[str, List[str]]] = None
# Can be per-layer (List[str]) or uniform (str).
# Values: "p2p", "all_gather", "a2a", "a2a+p2p"

Validation (MCore)

TransformerConfig.__post_init__ impose que a2a+p2p nécessite les tailles HCP et que le produit corresponde au CP.

Création de groupes de processus

parallel_state.initialize_model_parallel crée des sous-groupes CP hiérarchiques quand les tailles HCP sont fournies via create_hierarchical_groups. Bridge obtient actuellement ces groupes via le ProcessGroupCollection soutenu par MPU.

Intégration TE

TEDotProductAttention transmet les groupes hiérarchiques à Transformer Engine quand a2a+p2p est utilisé. Nécessite Transformer Engine >= 1.12.0.

Pitfalls

Bridge HCP est actuellement MPU-only : Si use_decentralized_pg=True, Bridge initialise des groupes CP plats et laisse HCP non défini.
Aucune recette Bridge enregistrée n'exerce actuellement HCP directement.
Les assistants de chargement single-GPU effacent hierarchical_context_parallel_sizes.
Entraînement silencieusement brisé sur les anciennes piles : Si vous utilisez a2a+p2p sans définir hierarchical_context_parallel_sizes, MCore fait maintenant une assertion. Les versions plus anciennes désactiveraient silencieusement la communication CP, chaque rank ne faisant attention qu'à son chunk local et produisant un débit artificiellement élevé avec des gradients cassés.
Le produit doit correspondre : prod(hierarchical_context_parallel_sizes) doit exactement égaler context_parallel_size. Une incompatibilité déclenche une assertion.
Vérifier dans les logs : Cherchez la sortie d'initialisation du groupe de processus. Vous devriez voir HIERARCHICAL_CONTEXT_PARALLEL_GROUPS être créés. Si vous voyez seulement CONTEXT_PARALLEL_GROUP, HCP n'est pas actif.

Verification

Aucun test Bridge end-to-end dédié n'existe encore pour HCP (voir @skills/nemo-mbridge-perf-hierarchical-context-parallel/card.yaml follow_up_validation). Utilisez plutôt les tests unitaires existants et l'inspection des logs.

Exécutez le test unitaire decentralized-PG pour confirmer que le comportement CP plat est préservé :

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

Pour un smoke check manuel, lancez un exécution 4-GPU avec une petite recette et 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

Critères de succès :

Les logs montrent HIERARCHICAL_CONTEXT_PARALLEL_GROUPS en cours de création
L'entraînement complète au moins une étape sans erreur
Si vous voyez seulement CONTEXT_PARALLEL_GROUP, HCP n'est pas actif