Skill de Contexte Hiérarchique Parallèle

Ce skill couvre le parallélisme hiérarchique du contexte : 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 savoir 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/perf-hierarchical-context-parallel/card.yaml

Activation

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

Ancres de code

Configuration amont et validation :

context_parallel_size: int = 1
"""Divise l'entrée réseau le long de la dimension de séquence across des rangs GPU."""

hierarchical_context_parallel_sizes: Optional[list[int]] = None
"""Degrés du parallélisme hiérarchique du contexte. Les utilisateurs doivent fournir une liste pour spécifier 
   les tailles pour différents niveaux. En prenant le type de communication a2a+p2p cp comme exemple, il contient
   des groupes de deux niveaux, donc la première valeur de la liste indique la taille du groupe du type
   de communication a2a, et la deuxième valeur indique la taille du groupe du type de communication p2p.
"""

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"

Chemin MPU Bridge :

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

Chemin PG décentralisé Bridge :

pg_collection = ProcessGroupCollection(
    ...
    cp=cp_pg,
    tp_cp=tp_cp_pg,
    hcp=None,
    ep=ep_pg,
    ...
)

Plan d'implémentation

Définition de configuration

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
# Pour a2a+p2p, première valeur = taille groupe a2a, deuxième valeur = taille groupe p2p.
# Le produit doit égaler 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
# Peut être par couche (List[str]) ou uniforme (str).
# Valeurs : "p2p", "all_gather", "a2a", "a2a+p2p"

Validation (MCore)

TransformerConfig.__post_init__ impose que a2a+p2p nécessite des tailles HCP et que le produit correspond à 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 récupère actuellement ces groupes via ProcessGroupCollection soutenu par MPU.

Intégration TE

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

Pièges

Bridge HCP est MPU-only aujourd'hui : Si use_decentralized_pg=True, Bridge initialise des groupes CP plats et laisse HCP non défini.
Aucune recette Bridge enregistrée ne teste actuellement HCP directement.
Les aides de chargement mono-GPU effacent hierarchical_context_parallel_sizes.
Entraînement silencieusement cassé sur d'anciennes piles : Si vous utilisez a2a+p2p sans définir hierarchical_context_parallel_sizes, MCore lance maintenant une assertion. Les versions plus anciennes désactivaient silencieusement la communication CP, donc chaque rang n'assistait qu'à sa part locale et produisait 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. Un décalage 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 en cours de création. Si vous ne voyez que CONTEXT_PARALLEL_GROUP, HCP n'est pas actif.

Vérification

Aucun test end-to-end Bridge dédié n'existe encore pour HCP (voir @skills/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 PG décentralisé 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 une vérification manuelle rapide, lancez une 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 affichent HIERARCHICAL_CONTEXT_PARALLEL_GROUPS en cours de création
L'entraînement complète au moins une étape sans erreur
Si vous ne voyez que CONTEXT_PARALLEL_GROUP, HCP n'est pas actif