Récacheminement Top‑1 et Élagage d’Experts Réduisent les Calculs FFN de MoE de 50%

Un simple passage du routage top‑2 au top‑1 peut presque réduire de moitié les FLOPs des FFN experts dans les transformateurs mixture‑of‑experts (MoE) — pourtant, de nombreuses équipes échouent à voir l’accélération en production en raison des retards, de l’ajustement et de la planification déséquilibrée qui récupèrent le gain théorique. Avec les bonnes mathématiques de routage, un ajustement de capacité et des runtimes MoE conscients, la parcimonie architecturale se transforme en un flux stable de tokens par seconde.

Pourquoi maintenant: des modèles ouverts à la pointe comme Mixtral et DeepSeek‑MoE ont rendu le MoE courant, et des guides pratiques de DeepSpeed‑MoE et MegaBlocks montrent comment router et planifier les experts sans fondre le débit. Cet article démontre comment la réduction top‑k et l’élagage d’experts réduisent le calcul, ce qui casse lorsque vous modifiez le routage, et comment les noyaux et les planificateurs rétablissent l’équilibre.

Nous passerons en revue: où vont les FLOPs de MoE (spoiler: les FFN experts dominent), la réduction exacte de 2× du FFN de top‑2 → top‑1 et les chemins de communication qu’elle élimine, l’ajustement des facteurs de capacité et des pertes auxiliaires pour éviter l’effondrement du routage, l’élagage d’experts basé sur l’utilisation, les mécanismes de runtime dans DeepSpeed‑MoE et MegaBlocks pour un débit stable, la récupération de la qualité via un léger LoRA, et une liste de vérification pratique avec du code pour implémenter et vérifier les changements.

Détails de l’Architecture/Implémentation

Où vont concrètement les FLOPs dans MoE

Dans les blocs Transformer MoE, l’attention reste dense tandis que le MLP/FFN devient épars via des experts. Par token, le routage sélectionne k experts; chaque expert sélectionné applique son FFN. Avec k=2, vous payez deux passes FFN par token; avec k=1, vous payez une. Comme les FFN surpassent les FLOPs d’attention dans de nombreux paramètres de décodeur, réduire k de 2 à 1 réduit le calcul des FFN experts d’environ 50 % et réduit également le trafic interappareils lorsque les experts sont fragmentés.

De top‑2 à top‑1: mathématiques de la passation et chemins supprimés

Soit x une représentation de token, W_r la projection du routeur produisant les logits g = W_r x. Le routeur sélectionne les indices top‑k et applique un softmax sur ces k entrées pour produire des poids de mélange p. Avec k ∈ {1,2}:

# Pseudocode pour le passage MoE
logits = router(x) # [num_experts]
indices = topk(logits, k) # k = 1 (top‑1) ou 2 (top‑2)
weights = softmax(logits[indices]) # mélange sur les experts sélectionnés
outputs = sum_j weights[j] * expert_ffn[indices[j]](x)

Opérationnellement, top‑2 introduit deux distributions d’experts par token, souvent via deux collectives tout‑à‑tout: les tokens sont partitionnés par expert, envoyés aux appareils hébergeant ces experts, traités, puis renvoyés. Top‑1 supprime une distribution et divise par deux les appels FFN. Cela réduit également le rembourrage: dans chaque micro‑lot d’expert, moins de tokens signifient moins de créneaux inactifs lorsque les lots sont limités par la capacité.

Chemins de communication simplifiés par top‑1:

• Moins de phases tout‑à‑tout et de charges utiles plus petites par étape.
• Réduction des files d’attente entrantes/sortantes d’experts à vidanger.
• Moins de rembourrage par micro‑lot d’expert lorsque la capacité est respectée.

Ces effets se multiplient avec des planificateurs conscients de MoE qui fusionnent le travail et atténuent les retards.

Facteur de capacité et perte auxiliaire: éviter l’effondrement du routage

Le MoE de type commutateur utilise un facteur de capacité C pour limiter les tokens par expert: capacité = floor(C × tokens_par_lot / nombre_d’experts). Les tokens au-delà de la capacité sont soit abandonnés (en temps d’entraînement) soit réacheminés, et une perte auxiliaire d’équilibrage de charge encourage le routeur à distribuer les tokens uniformément tout en préservant les attributions à haute probabilité. Après réduction de k, deux choses échouent couramment:

• Effondrement: le routeur se concentre excessivement sur quelques experts, saturant la capacité et entraînant débordement/rembourrage.
• Sous-utilisation: certains experts se refroidissent, nuisant à la spécialisation et à la qualité.

Atténuations:

• Augmenter le coefficient de perte d’équilibrage de charge et démarrer à chaud à partir d’un point de contrôle qui utilise déjà cette perte auxiliaire.
• Réajuster C. Avec top‑1, augmenter légèrement C (par exemple, 1,0 → 1,2) offre de la marge pour éviter le débordement à des horizons courts; pour des contextes longs, vous pouvez abaisser C pour réduire le rembourrage.
• Appliquer un court passage de récupération uniquement sur le routeur ou LoRA‑sur‑le‑routeur pour réétaler le trafic.

Élagage d’experts via l’historique de l’utilisation

Une fois le routage stable sous top‑1, profiler l’utilisation des experts sur des charges de travail représentatives. Calculer la fraction de tokens (ou de tokens pondérés par score de passage) traitée par expert. Les experts en dessous d’un seuil (par exemple, <0,5–1,0 % sur de longues fenêtres) sont candidats pour l’élagage. Après suppression, rééquilibrer le routeur et effectuer une légère récupération pour recentrer la spécialisation.

Sauvegardes clés:

• Utiliser de longues traces et des invites diverses pour éviter d’élaguer des experts qui servent des créneaux rares mais essentiels.
• Privilégier un élagage étagé avec de petits lots d’experts retirés à la fois, avec validation entre les étapes.

Pièges à débit: retardataires, rembourrage, déséquilibre

Les changements de routage révèlent souvent des inefficacités cachées:

• Retardataires: un expert lourdement chargé allonge l’étape tandis que d’autres restent inactifs.
• Rembourrage: les noyaux à forme fixe remplissent jusqu’à la capacité; une assignment déséquilibrée token‑à‑expert amplifie les créneaux perdus.
• Dérive de communication: le trafic déséquilibré tout‑à‑tout produit une latence en queue.

MegaBlocks aborde ces problèmes en regroupant des micro‑lots en blocs équilibrés, réduisant le rembourrage et aplanissant la charge au niveau du dispositif; il reste efficace sous des distributions biaisées et des ensembles d’experts post-élagage. DeepSpeed‑MoE fournit des paramètres de runtime pour la capacité, les politiques de débordement et la superposition communication/calcul qui stabilisent p50/p99 sous top‑1.

Runtimes conscients de MoE: DeepSpeed‑MoE et MegaBlocks

• DeepSpeed‑MoE: un moteur prêt pour la production avec passage top‑k, contrôles du facteur de capacité, pertes d’équilibrage de charge, et partitionnement parallèle d’experts. Il raccorde all‑to‑all avec les calculs experts et expose la configuration pour le routage, la capacité, et le partitionnement d’experts.
• MegaBlocks: une approche de noyau et de planification qui partitionne les calculs en blocs uniformes, limitant le rembourrage et atténuant les retardataires. Elle fournit un débit robuste lorsque le routage biaisé ou l’élagage d’experts perturbe l’uniformité.

Les deux approches sont compatibles avec les pipelines d’entraînement/inférence type Mixtral‑style MoE et DeepSeek‑MoE de même que les formulations router/perte auxiliaire dérivées des Transformateurs Commutés.

Notes d’implémentation sur les piles GPU sans toucher aux noyaux denses

• CUDA/NVIDIA: Si vos chemins denses attention/MLP utilisent déjà des noyaux fusionnés, vous pouvez adopter top‑1 et l’élagage d’experts sans modifier les noyaux denses. Concentrez-vous sur les noyaux d’envoi MoE et la superposition tout‑à‑tout dans DeepSpeed‑MoE, ou substituez MegaBlocks pour le calcul expert groupé en blocs.
• Triton/ROCm: MegaBlocks compile via Triton et s’exécute sur ROCm, offrant aux utilisateurs AMD une voie portable pour réaliser les gains de routage MoE sans réécrire les noyaux denses.
• Environnements mixtes: Gardez les noyaux denses intacts; les changements de routage se produisent à la frontière de délégation/agrégation MoE. Validez la performance collective et l’alignement de la disposition mémoire, non les calculs attention/FFN eux-mêmes.

Maintenir la qualité après les changements structurels

Top‑1 et l’élagage d’experts altèrent la spécialisation. Une récupération légère stabilise les métriques:

• LoRA sur le routeur (et éventuellement les FFN experts) pendant quelques milliers d’étapes suffit souvent à rétablir l’équilibre d’utilisation et à récupérer 0,5–2 points sur des évaluations communes, pour une fraction du coût de réglage complet.
• Une distillation de connaissances brève du modèle pré-élagage peut préserver davantage la connaissance des experts sur les motifs rares après l’élagage.

Tableaux Comparatifs

Choix de routage et d’élagage en bref

Approche	Effet sur le calcul	Effet sur la communication	Risques principaux	Meilleures atténuations	Runtime supporté
Top‑2 (référence)	2× FFN d’expert par token	Deux passages tout‑à‑tout	Plus de FLOPs, rembourrage	N/A (référence)	Tout runtime MoE
Top‑1	~50% de FLOPs FFN d’expert en moins	Un passage supprimé; charges utiles plus petites	Effondrement de routage; experts froids	Ajuster la perte auxiliaire, ajuster C, courte récupération du routeur	DeepSpeed‑MoE, MegaBlocks
Élagage d’expert	Réduit le nombre de paramètres et d’experts actifs	Moins de destinations; peut accentuer le biais	Sur-élagage nuit à la spécialisation	Historiques d’utilisation, élagage par étape, récupération LoRA	DeepSpeed‑MoE, MegaBlocks

Mécanismes de runtime: ce qui stabilise le débit

Runtime	Mécanisme clé	Forces sous top‑1/élagage	Remarques
DeepSpeed‑MoE	Superposition de tout‑à‑tout + calcul expert; boutons de capacité/perte auxiliaire	Configurations prêtes pour la production; réglage facile du routeur	Intégration dans les pipelines d’entraînement/inférence
MegaBlocks	Programmation et calcul groupés en blocs	Atténue rembourrage/retardataires, robuste au biais	Noyaux Triton, portable sur ROCm

Meilleures Pratiques

• Commencez avec top‑1, puis élaguez: D’abord stabiliser le routage à k=1, puis élaguez les experts froids en fonction de l’utilisation à longue portée.
• Réglez le routeur comme un module de première classe:
• Augmentez le poids de la perte d’équilibrage de charge après avoir changé k; surveillez l’entropie et le trafic par expert.
• Réglez le facteur de capacité C en fonction de la taille du lot et de la longueur de la séquence; préférez un C légèrement plus élevé pour les séquences courtes pour éviter le débordement, un C plus bas pour les séquences longues pour réduire le rembourrage.
• Utilisez des runtimes conscients de MoE:
• Activez la superposition tout‑à‑tout et le respect de la capacité de DeepSpeed‑MoE; ou adoptez MegaBlocks pour réduire le rembourrage et les retardataires.
• Profilez avant l’élagage:
• Enregistrez les histogrammes d’assignation d’experts et les cartes de chaleur de latence par expert sous un trafic semblable à la production; élaguez uniquement les experts constamment froids.
• Récupérez légèrement:
• Appliquez LoRA au routeur/FFN experts pour une récupération rapide; ajoutez éventuellement une distillation de petits mélanges pour recentrer la spécialisation.
• Vérifiez le comportement à long contexte: certains experts capturent des motifs rares à longue portée; incluez des évaluations de long contexte lors de l’élagage.

Exemples Pratiques

1) Configuration DeepSpeed‑MoE pour top‑1 et ajustement de la capacité

{
 "moe": {
 "enabled": true,
 "num_experts": 16,
 "top_k": 1,
 "capacity_factor": 1.25,
 "router_aux_loss_coef": 0.01,
 "router_topk": true,
 "moe_param_group": true,
 "drop_tokens": false,
 "alltoall": { "overlap_communication": true }
 }
}

• top_k: passer de 2 → 1 pour réduire de moitié le calcul des FFN experts.
• capacity_factor: commencez légèrement au-dessus de 1,0 pour éviter le débordement aux contextes courts; réajustez selon la charge de travail.
• router_aux_loss_coef: augmentez modérément après le changement de top‑k pour prévenir l’effondrement.
• overlap_communication: gardez les GPU occupés pendant le tout‑à‑tout.

2) Histogramme d’utilisation pour guider l’élagage des experts

# extrait de type PyTorch capturant les statistiques de routage en évaluation
util = torch.zeros(num_experts, device="cuda")
with torch.no_grad():
 for batch in loader:
 logits = router(batch.hiddens) # [B, num_experts]
 idx = torch.topk(logits, k=1, dim=-1).indices.squeeze(-1) # top-1
 util.index_add_(0, idx.flatten(), torch.ones_like(idx.flatten(), dtype=util.dtype))

util = util / util.sum() # fraction de tokens par expert
cold = (util < 0.005).nonzero().flatten() # par ex., <0.5%
print("Candidats à l'élagage:", cold.tolist())

• Exécutez sur des traces diverses et longues. Confirmez avec des cartes de chaleur de latence par expert.
• Elaguez par étapes; après chaque étape, effectuez une courte récupération LoRA sur le routeur et les experts concernés.

3) LoRA uniquement sur le routeur pour une récupération rapide

from peft import LoraConfig, get_peft_model

lora_cfg = LoraConfig(
 r=8, lora_alpha=16, target_modules=["router.proj"], lora_dropout=0.05
)
model.router = get_peft_model(model.router, lora_cfg)
# Ajuster finement pour quelques K étapes sur des tâches mixtes pour recentrer la spécialisation

• Concentrez-vous d’abord sur le routeur; étendez aux FFN experts si la spécialisation a dérivé.

Conclusion

Le routage Top‑1 est le levier le plus propre pour réduire de moitié les calculs FFN de MoE, mais le gain ne se matérialise d’un bout à l’autre que si vous gérez le routeur et le planificateur comme des systèmes de première classe. Le facteur de capacité et les pertes d’équilibrage de charge maintiennent une utilisation uniforme; les histogrammes d’utilisation identifient les experts pouvant être élagués en toute sécurité; et les runtimes conscients de MoE comme DeepSpeed‑MoE et MegaBlocks convertissent la parcimonie architecturale en un flux de tokens stable par seconde en apprivoisant le rembourrage, les retardataires et la dérive de communication. Une touche finale de LoRA aide à restaurer la spécialisation après des changements structurels.

Principaux enseignements:

• Top‑2 → top‑1 supprime un passage d’expert et un trajet tout‑à‑tout par token, divisant par deux les FLOPs FFN d’experts et réduisant le trafic interappareils.
• L’ajustement de la capacité et des pertes auxiliaires prévient l’effondrement du routage; surveillez l’utilisation et l’entropie comme métriques de première classe.
• Elaguez uniquement les experts persistants froids, guidé par des histogrammes à longue portée; récupérez légèrement avec LoRA.
• DeepSpeed‑MoE et MegaBlocks stabilisent le débit par une superposition de la communication et une réduction du rembourrage/retardataires, même sous biais et elague.

Prochaines étapes:

• Passez à top‑1 dans un environnement de mise en scène, réglez la capacité et la perte auxiliaire, et validez l’utilisation et p99.
• Collectez des histogrammes sur le trafic réel; élaguez les experts par étapes avec des vérifications entre chaque cycle.
• Adoptez les fonctionnalités de runtime conscient de MoE (superposition tout‑à‑tout ou planification par bloc) et ajoutez une brève récupération LoRA.

La conséquence: avec un routage et une planification disciplinés, la parcimonie théorique du MoE devient un débit réel sans réécriture de vos noyaux denses ⚙️.

Sources

• Transformateurs Commemorialeurs: Passage à des Modèles à Billion de Paramètres avec une Parcimonie Simple et Efficace — https://arxiv.org/abs/2101.03961 — Établit le passage top‑k, les facteurs de capacité, et les pertes d’équilibrage de charge auxiliaires qui sous-tendent la stabilité du routage top‑1.
• Mixtral of Experts — https://arxiv.org/abs/2401.04088 — Démontre des architectures MoE pratiques et discute des comportements de routage et des considérations de spécialisation pertinentes pour les changements top‑k et l’élagage.
• DeepSeek-MoE (repo) — https://github.com/deepseek-ai/DeepSeek-MoE — Montre les contrôles de routage réels et la gestion des experts utilisés lors de l’application des ajustements de routage et de l’élagage.
• MegaBlocks: Formation et Inference Efficaces en MoE Épars — https://arxiv.org/abs/2211.15841 — Fournit des noyaux de planification/groupage en blocs qui atténuent le rembourrage, les retardataires et le déséquilibre sous des charges d’experts biaisées/élaguées.
• Tutoriel DeepSpeed-MoE — https://www.deepspeed.ai/tutorials/moe/ — Documente les mécanismes de runtime, l’ajustement des capacités/perte auxiliaire, et la superposition communication/calcul qui stabilisent le débit sous top‑1.
• Documentation AMD ROCm — https://rocm.docs.amd.com/ — Confirme les voies Triton/ROCm pour déployer des noyaux et MoE de type MegaBlocks sur AMD sans réécrire les noyaux denses.
• LoRA: Adaptation de Basse-Rangée des Grands Modèles de Langage — https://arxiv.org/abs/2106.09685 — Supporte un ajustement fin d’adaptateur léger pour récupérer la qualité après routage/élagage.
• AdaLoRA: Allocation Budgétaire Adaptative pour le Peaufinage Efficace des Paramètres — https://arxiv.org/abs/2303.10512 — Fournit une alternative efficiente en adaptateur pour la récupération et le recentrage de la spécialisation après l’élagage.

Sources & Références