Skill Chevauchement Expert-Parallèle MoE

Docs stables : @docs/training/communication-overlap.md Carte : @skills/perf-expert-parallel-overlap/card.yaml

Références

• Docs stables : @docs/training/communication-overlap.md
• Métadonnées structurées : @skills/perf-expert-parallel-overlap/card.yaml

Description

Le chevauchement expert-parallèle (EP) masque le coût de la communication all-to-all de dispatch/combine de tokens en l'exécutant concurremment avec le calcul FFN des experts. Optionnellement, le calcul retardé du gradient de poids d'expert (delay_wgrad_compute) fournit un chevauchement supplémentaire en reportant le wgrad pour chevaucher le forward de la couche suivante.

Bridge supporte deux chemins dispatcher :

Décision Rapide

Utilisez le chevauchement EP quand :

• le modèle est MoE avec EP > 1
• la communication de dispatch/combine d'experts est une part significative du temps d'étape
• vous avez de la marge mémoire et accordez la priorité au débit

Préférez :

• le dispatcher alltoall pour le premier déploiement (compatibilité plus large)
• flex + DeepEP/HybridEP quand vous utilisez des GPUs pris en charge et recherchez des gains supplémentaires

Évitez le chevauchement EP quand :

• la recomputation complète d'activation est activée
• moe_shared_expert_overlap est activé
• l'exécution est encore en phase de vérification de correction
• PyTorch < 2.6.0

Résultat attendu :

• si le dispatch all-to-all est un goulot d'étranglement clair en profilage, le chevauchement peut produire une accélération modérée à significative
• si l'exécution est petite, légère en communication, ou dominée par un autre mur, le gain peut être négligeable

Activation

Dispatcher alltoall

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
cfg.comm_overlap.delay_wgrad_compute = True
cfg.model.moe_shared_expert_overlap = False

cfg.model.expert_model_parallel_size = 8
cfg.model.num_moe_experts = 64
cfg.model.moe_token_dispatcher_type = "alltoall"
cfg.model.bf16 = True
cfg.model.fp16 = False

Dispatcher flex (DeepEP ou HybridEP)

from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
cfg.comm_overlap.delay_wgrad_compute = True
cfg.model.moe_shared_expert_overlap = False

apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend="deepep")
# ou : apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend="hybridep")

Compatibilité et Contraintes

• expert_model_parallel_size > 1
• num_moe_experts > 1
• moe_token_dispatcher_type doit être "alltoall" ou "flex"
• moe_shared_expert_overlap = False
• La précision de base est BF16 ou FP16
• PyTorch >= 2.6.0
• Si PP > 1, virtual_pipeline_model_parallel_size doit être défini
• recompute_granularity != "full", recompute_method = None, recompute_num_layers = None
• mtp_num_layers doit être None ou 1
• delay_wgrad_compute nécessite overlap_moe_expert_parallel_comm comme prérequis
• delay_wgrad_compute avec overlap_grad_reduce nécessite TE >= 2.7.0
• delay_wgrad_compute avec gradient_accumulation_fusion nécessite TE >= 2.7.0
• CUDA graph scope attn + delay_wgrad_compute nécessite TE >= 2.12.0, gradient_accumulation_fusion = True, et pas de biais d'attention
• DeepEP : GPUs Ampere, Hopper, B200, B300 uniquement
• HybridEP : Ampere, Hopper, B200, B300, GB200/GB300 avec NVL72

Configuration Minimale Fonctionnelle

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
cfg.comm_overlap.delay_wgrad_compute = False
cfg.model.expert_model_parallel_size = 4
cfg.model.num_moe_experts = 64
cfg.model.moe_token_dispatcher_type = "alltoall"
cfg.model.moe_shared_expert_overlap = False
cfg.model.bf16 = True

Utilisez ceci comme point de départ en priorité correction. Ajoutez le wgrad retardé, le dispatch flex, et les interactions CUDA-graph uniquement après que le chemin de chevauchement simple soit confirmé fonctionnel.

Commande Minimale Exécutable

Exemple de harnais de performance :

python scripts/performance/setup_experiment.py \
  --model qwen3-30b-a3b \
  --moe_a2a_overlap \
  --num_nodes 2 \
  --gpus_per_node 8 \
  --max_steps 20

Vérification par test unitaire :

uv run python -m pytest \
  tests/unit_tests/training/test_comm_overlap.py -k "moe" \
  tests/unit_tests/training/test_deepep.py -q

Vérification

Tests unitaires

uv run python -m pytest \
  tests/unit_tests/training/test_comm_overlap.py \
  tests/unit_tests/training/test_deepep.py -q

Vérifications de logs

Après une exécution réussie avec chevauchement EP :

1. Confirmer l'absence d'erreurs d'assertion lors de la finalisation de CommOverlapConfig
2. Confirmer que overlap_moe_expert_parallel_comm apparaît en tant que True dans la config loggée
3. Si vous utilisez le dispatcher flex, confirmer moe_token_dispatcher_type = "flex" et le backend correct dans les logs

Critères de succès

• La validation config réussit pour le dispatcher et les paramètres de chevauchement sélectionnés
• L'entraînement se termine sans blocages ou pannes d'assertion
• Le débit s'améliore ou au moins ne régresse pas pour la charge de travail cible
• La trajectoire de loss correspond à la baseline (le chevauchement ne doit pas affecter la convergence)

Ancrages de Code

Validation de chevauchement Bridge

if self.user_comm_overlap_cfg.overlap_moe_expert_parallel_comm is True:
    assert model_cfg.expert_model_parallel_size > 1, ...
    assert model_cfg.num_moe_experts > 1, ...
    assert model_cfg.moe_token_dispatcher_type in ["alltoall", "flex"], ...
    assert model_cfg.bf16 or model_cfg.fp16, ...
    assert is_torch_min_version("2.6.0"), ...
    # ... PP + VPP check, recompute checks, shared_expert_overlap check ...

Validation du wgrad retardé

if self.user_comm_overlap_cfg.delay_wgrad_compute is True:
    # TE version checks for overlap_grad_reduce and gradient_accumulation_fusion
    # CUDA graph scope validations for delayed wgrad
    assert overlap_moe_expert_parallel_comm, ...

Activation flex-dispatcher

def apply_flex_dispatcher_backend(...):
    # GPU architecture check for DeepEP / HybridEP
    model_config.moe_token_dispatcher_type = "flex"
    model_config.moe_flex_dispatcher_backend = moe_flex_dispatcher_backend
    model_config.moe_shared_expert_overlap = False

Override du harnais perf

def _set_moe_a2a_overlap_overrides(recipe, moe_a2a_overlap=False):
    if moe_a2a_overlap:
        recipe.comm_overlap.overlap_moe_expert_parallel_comm = True
        recipe.comm_overlap.delay_wgrad_compute = True
        recipe.model.moe_shared_expert_overlap = False

Tests

Diagnostic de Défaillance

Limitations Connues

• Définir moe_flex_dispatcher_backend seul n'active pas le dispatch flex — vous devez appeler apply_flex_dispatcher_backend(...).
• Les recettes publiques sont souvent prudentes et laissent le chevauchement MoE désactivé par défaut.
• Les gains de débit end-to-end n'ont pas encore été mesurés dans une expérience Bridge contrôlée pour chaque famille de modèles. La validation de code est plus solide qu'une seule affirmation de performance universelle.
• Le chevauchement MoE et le chevauchement d'expert partagé s'excluent mutuellement.
• CUDA graph plus wgrad retardé est un chemin multi-contrainte qui nécessite une validation prudente de version TE et de scope.