Chevauchement de communication MoE

Pour un aperçu de plus haut niveau, voir :

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

Décision rapide

Utilise le chevauchement de communication MoE quand :

EP > 1
le temps de dispatch de tokens ou de combine est visible dans le profile
la run est déjà correcte et tu accordes maintenant le débit

Évite de l'activer comme étape initiale de mise en service. Il est plus facile de valider après que le dispatcher, le mode de routage et le plan de recompute 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ésactive le chevauchement d'expert partagé lors de l'utilisation du dispatch overlap
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

Définir moe_flex_dispatcher_backend seul n'active pas le dispatch flex. Tu dois aussi définir moe_token_dispatcher_type = "flex".

Interaction entre recompute et CUDA Graph

Le recompute complet n'est pas un bon compagnon pour le chemin overlap.
delay_wgrad_compute ajoute d'autres contraintes si les portées CUDA-graph incluent l'attention ou le travail de routeur MoE.
En pratique, le recompute sélectif est l'appairage plus sûr quand le overlap est activé.

Ancres de code

Validation du overlap : src/megatron/bridge/training/comm_overlap.py
Backend flex dispatcher : 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ésactive le chevauchement d'expert partagé lors de l'utilisation du chemin dispatch overlap.
PP sans VPP : Le overlap MoE requiert VPP quand le pipeline parallelism est actif. Sans lui, la planification du overlap ne peut pas s'entrelacer correctement.
Flex ≠ flag backend : moe_flex_dispatcher_backend="deepep" seul ne fait rien si moe_token_dispatcher_type est encore "alltoall".
Défauts conservateurs des recettes publiques : La plupart des recettes publiques laissent le MoE overlap désactivé. Tu dois l'activer explicitement via des overrides.
Les gains de performance dépendent de la charge de travail : le overlap aide surtout quand la communication de dispatch est déjà une part visible du temps de step. Ce n'est pas garanti pour aider chaque run EP petite ou légèrement chargée.

Vérification

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