Chevauchement de communication MoE

Pour un aperçu de niveau supérieur, voir :

@docs/training/communication-overlap.md
@skills/nemo-mbridge-perf-moe-comm-overlap/card.yaml

Décision rapide

Utilisez le chevauchement de communication MoE quand :

EP > 1
le temps de dispatch ou de combinaison des tokens est visible dans le profil
l'exécution est déjà correcte et vous affinez maintenant le débit

Évitez de l'activer comme étape initiale de mise en place. Il est plus facile de valider après que le dispatcher, le mode de routage et le plan de recalcul soient déjà stables.

Activation

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True

# Optionnel : wgrad retardé pour un chevauchement supplémentaire
cfg.comm_overlap.delay_wgrad_compute = True

# IMPORTANT : désactiver le chevauchement d'expert partagé lors de l'utilisation du chevauchement de dispatch
cfg.model.moe_shared_expert_overlap = False

Prérequis

expert_model_parallel_size > 1
num_moe_experts > 1
moe_token_dispatcher_type doit être "alltoall" ou "flex"
Précision : BF16 ou FP16
Si PP est utilisé, VPP (virtual_pipeline_model_parallel_size) doit être défini (non-None)

Activation du dispatcher Flex

La définition de moe_flex_dispatcher_backend seule n'active pas le dispatch flex. Vous devez aussi définir moe_token_dispatcher_type = "flex".

Interaction avec recalcul et CUDA Graph

Le recalcul complet n'est pas un bon compagnon pour le chemin de chevauchement.
delay_wgrad_compute ajoute des contraintes supplémentaires si les scopes CUDA-graph incluent le travail d'attention ou de routeur MoE.
En pratique, le recalcul sélectif est l'appairage plus sûr quand le chevauchement est activé.

Mise en garde sur les mesures de courte durée

Un smoke test H100 x16 du 2026-05-18 sur Qwen3 30B-A3B mock pretraining du main courant utilisait EP=16, alltoall, global batch size 1024, CUDA graphs désactivés, et moe_permute_fusion=false car la pile PyTorch 25.11 / TE / Triton a échoué dans la permutation fusionnée de Transformer Engine lors de la mise en place antérieure.

Les résultats étaient directionnels plutôt que de qualité release :

sans chevauchement EP : moyenne en régime établi de 41,25s sur les itérations 3-8
chevauchement EP : moyenne en régime établi de 31,31s sur les itérations 3-8
chevauchement EP plus delay_wgrad_compute : moyenne en régime établi de 31,20s sur les itérations 3-8

Considérez ceci comme une preuve que le chevauchement EP peut aider une forme MoE alltoall inter-nœud quand la communication est exposée. Ce n'est pas une preuve que le wgrad retardé est un gain séparé, et cela ne valide pas le chemin de permutation fusionnée. Un smoke test antérieur du 2026-05-16 sur la même forme a montré le même motif.

Ancres de code

Validation du chevauchement : src/megatron/bridge/training/comm_overlap.py
Backend du dispatcher Flex : src/megatron/bridge/training/flex_dispatcher_backend.py
Config : src/megatron/bridge/training/config.py
Tests unitaires : tests/unit_tests/training/test_comm_overlap.py
Tests DeepEP : tests/unit_tests/training/test_deepep.py

Pièges

Conflit de chevauchement d'expert partagé : moe_shared_expert_overlap et overlap_moe_expert_parallel_comm peuvent entrer en conflit. Désactivez le chevauchement d'expert partagé lors de l'utilisation du chemin de chevauchement de dispatch.
PP sans VPP : le chevauchement MoE nécessite VPP quand le parallélisme de pipeline est actif. Sans cela, la planification du chevauchement ne peut pas s'entrelacer correctement.
Flex ≠ drapeau backend : moe_flex_dispatcher_backend="deepep" seul ne fait rien si moe_token_dispatcher_type est toujours "alltoall".
Valeurs par défaut de recette conservatrices : la plupart des recettes publiques laissent le chevauchement MoE désactivé. Vous devez l'activer explicitement via des overrides.
Les gains de performance dépendent de la charge de travail : le chevauchement aide le plus quand la communication de dispatch est déjà une tranche visible du temps d'étape. Il n'est pas garanti d'aider chaque exécution EP petit ou faiblement chargée.

Vérification

Recherchez les messages de journal liés au chevauchement lors de l'initialisation. La validation du chevauchement de comm dans comm_overlap.py lèvera une exception si les prérequis ne sont pas satisfaits, donc un démarrage propre confirme que la fonctionnalité est active.

Pour un smoke test court de harnais de performance, gardez la forme de commande explicite et ne variez qu'un seul bouton de chevauchement à la fois :

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

Si la permutation MoE fusionnée échoue lors de la mise en place, ajoutez model.moe_permute_fusion=false pour séparer le chronométrage du chevauchement de la validation de la pile runtime, puis retestez avec le conteneur de production appairé.