Skill de chevauchement Expert-Parallel MoE

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

Références

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

Qu'est-ce que c'est

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

Bridge supporte deux chemins de dispatcher :

Décision rapide

Utiliser le chevauchement EP quand :

• le modèle est MoE avec EP > 1
• la communication de dispatch/combine des experts représente une partie significative du temps d'étape
• vous avez de la marge mémoire et accordez le tuning pour le débit

Préférer :

• le dispatcher alltoall pour le premier déploiement (compatibilité plus large)
• flex + DeepEP/HybridEP lors de l'exécution sur des GPUs supportés et en cherchant des gains supplémentaires

Éviter le chevauchement EP quand :

• la réactivation complète d'activation est activée
• moe_shared_expert_overlap est activé
• l'exécution est encore en train d'être mise au point pour la correction
• PyTorch < 2.6.0

Résultat attendu :

• si le dispatch all-to-all est un goulot d'étranglement clair du profile, le chevauchement peut produire une accélération modeste à significative
• si l'exécution est minuscule, 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 centré sur la correction. Ajoutez wgrad retardé, dispatch flex, et les interactions CUDA-graph seulement après avoir confirmé que le chemin de chevauchement simple fonctionne.

Preuves mesurées sur court terme

Une exécution de préentraînement simulé Qwen3 30B-A3B H100 x16 du main actuel du 2026-05-18 utilisait EP=16, alltoall, BF16, taille de batch global 1024, graphes CUDA désactivés, et moe_permute_fusion=false. Avec les itérations 3-8 comme fenêtre stable :

C'est une preuve pour activer le chevauchement EP simple sur cette forme all-to-all inter-nœud. Cela ne montre pas un gain indépendant significatif de wgrad retardé, et ne valide pas la permutation MoE fusionnée car ce chemin était désactivé pour la pile runtime.

Commande minimale exécutable

Exemple de harnais de performance à l'intérieur d'une allocation Slurm. Gardez le modèle, le parallélisme, le dispatcher et le runtime fixes, et variez uniquement les deux surcharges de chevauchement :

uv run python scripts/performance/run_script.py \
  -m qwen \
  -mr qwen3_30b_a3b \
  --task pretrain \
  -g h100 \
  -c bf16 \
  -ng 16 \
  -gn 8 \
  --max_steps 8 \
  --config_variant v1 \
  --cuda_graph_impl none \
  --moe_flex_dispatcher_backend None \
  --moe_a2a_overlap false \
  --tokenizer_type NullTokenizer \
  comm_overlap.overlap_moe_expert_parallel_comm=true \
  comm_overlap.delay_wgrad_compute=false \
  model.moe_shared_expert_overlap=false

N'utilisez pas --moe_a2a_overlap true pour séparer le chevauchement EP simple de wgrad retardé : le helper du harnais de performance active overlap_moe_expert_parallel_comm et delay_wgrad_compute.

Vérification du 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 des logs

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

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

Critères de succès

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

Ancres 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"), ...
    # ... vérification PP + VPP, vérifications de réactivation, vérification shared_expert_overlap ...

Validation wgrad retardé

if self.user_comm_overlap_cfg.delay_wgrad_compute is True:
    # Vérifications de version TE pour overlap_grad_reduce et gradient_accumulation_fusion
    # Validations de scope CUDA graph pour wgrad retardé
    assert overlap_moe_expert_parallel_comm, ...

Activation dispatcher flex

def apply_flex_dispatcher_backend(...):
    # Vérification d'architecture GPU pour 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

Surcharge 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 des défaillances

Limitations connues

• Définir seul moe_flex_dispatcher_backend n'active pas le dispatch flex — vous devez appeler apply_flex_dispatcher_backend(...).
• Les recettes publiques sont souvent conservatrices 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 du code est plus forte qu'une affirmation de performance universelle unique.
• 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 soigneuse de la version TE et du scope.