Hybrid / Hierarchical Context Parallel Skill

Pour comprendre ce qu'est HCP, quand l'utiliser et l'arbre de décision (a2a+p2p vs pure a2a vs p2p), voir:

• @docs/training/hybrid-context-parallel.md
• @skills/perf-hybrid-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 upstream 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
# 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__ garantit que a2a+p2p requiert les tailles HCP et que le produit correspond à CP.

Création de process group

parallel_state.initialize_model_parallel crée des sous-groupes CP hiérarchiques quand les tailles HCP sont fournies via create_hierarchical_groups.

Intégration TE

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

Pitfalls

1. Fonctionnalités différentes: a2a+p2p et hybrid_context_parallel=True upstream sont des fonctionnalités différentes. Cette dernière sert à équilibrer les charges de travail packed/variable-length.
2. Bridge HCP est MPU-only actuellement: Si use_decentralized_pg=True, Bridge initialise des groupes CP plats et laisse HCP non défini.
3. Aucune recette Bridge checked-in n'exerce actuellement HCP directement.
4. Les helpers de chargement single-GPU effacent hierarchical_context_parallel_sizes.
5. Entraînement silencieusement cassé: Si vous utilisez a2a+p2p sans définir hierarchical_context_parallel_sizes, MCore fait maintenant une assertion. Les versions plus anciennes désactivaient silencieusement la communication CP — chaque rang ne considérait que son chunk local, produisant un throughput artificiellement élevé mais des gradients complètement cassés.
6. Le produit doit correspondre: prod(hierarchical_context_parallel_sizes) doit exactement égaler context_parallel_size. Une incompatibilité déclenche une assertion.
7. Vérifiez dans les logs: Cherchez la sortie d'initialisation du process group. Vous devriez voir HIERARCHICAL_CONTEXT_PARALLEL_GROUPS être créés. Si vous ne voyez que CONTEXT_PARALLEL_GROUP, HCP n'est pas actif.

Verification

Aucun test end-to-end Bridge dédié n'existe encore pour HCP (voir @skills/perf-hybrid-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 flat-CP est préservé:

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

Pour un smoke check manuel, 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 montrent 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