MoonEP : l'équilibrage parfait¶
Le composant d'infrastructure le plus abouti du rapport, avec une preuve mathématique en annexe et un code ouvert.
Le problème¶
Dans les schémas de parallélisme d'experts (EP) conventionnels, les charges de jetons sont déséquilibrées entre rangs. Deux conséquences :
| Conséquence | Effet |
|---|---|
| Déséquilibre de calcul | Le débit d'entraînement est dicté par le rang le plus chargé ; les autres attendent |
| Formes dynamiques | Les tailles des activations d'experts routés varient d'un pas à l'autre, provoquant une fragmentation mémoire substantielle |
À cela s'ajoute un troisième coût, plus subtil : comme les tailles varient, l'hôte doit se synchroniser avec le GPU à chaque couche pour connaître les formes réelles avant de lancer le calcul — ce qui bloque le pipeline entre couches.
L'idée : les experts redondants dynamiques¶
MoonEP préserve le flux de calcul général des schémas conventionnels (comme DeepEP) et y ajoute la planification et la migration en ligne d'experts redondants.
Passe avant :
planifier les experts redondants à partir des sorties du routeur
du micro-lot et de la couche COURANTS
│
▼
les préchercher AVANT le calcul des experts routés
Passe arrière :
mettre en scène leurs gradients dans un tampon de réduction local
│
▼
une fois le calcul terminé, les réduire vers les tampons de gradient
de leurs rangs d'origine
Ce qu'est un « expert redondant »
Une copie temporaire d'un expert, placée sur un rang qui n'est pas son propriétaire. Si le rang A est surchargé et le rang B sous-chargé, on duplique un expert de A vers B, et une partie des jetons destinés à cet expert est traitée par B.
Le coût : la copie des poids de l'expert (33 M de paramètres). Le gain : l'équilibre parfait.
L'exigence : l'équilibre exact¶
MoonEP exige que chaque rang reçoive exactement \(S \times K\) jetons, où \(S\) est la longueur de séquence et \(K\) le nombre d'experts sélectionnés par jeton — de sorte que tous les rangs effectuent exactement la même quantité de calcul.
La question critique devient : combien d'experts redondants suffisent pour garantir un tel équilibre ?
Le théorème¶
Soit \(E\) le nombre d'experts et \(R\) la taille du groupe EP.
Théorème 1 — borne supérieure générale
Pour toute sortie de routeur, il existe un plan d'équilibrage utilisant au plus \(E/R\) experts redondants par rang.
Preuve (esquisse), telle que donnée en annexe :
Le lemme clé est qu'il existe un plan \(P^*\) tel que (a) chaque rang reçoit exactement \(S\times K\) jetons, et (b) les jetons distants de chaque rang proviennent d'un seul autre rang.
Construction : initialement, chaque rang ne détient que ses jetons locaux ; les rangs sont classés en sous-chargés ou surchargés. On choisit répétitivement un rang sous-chargé et un rang surchargé, et on migre des jetons du surchargé vers le sous-chargé jusqu'à le remplir exactement à \(S\times K\). Le rang surchargé peut rester surchargé, devenir exactement équilibré, ou devenir sous-chargé, et il est remis dans l'ensemble correspondant.
Chaque remplissage rend un rang sous-chargé définitivement équilibré, donc le processus se termine en au plus \(R-1\) remplissages. Et chaque rang n'est rempli qu'une seule fois, donc ses jetons distants proviennent d'un rang unique — ce qui prouve le lemme.
Conséquence : si tous les jetons distants du rang \(r\) viennent du rang \(s\), ces jetons appartiennent à au plus \(E/R\) experts locaux de \(s\), d'où \(m_r(P^*) \le E/R\), et donc
Théorème 2 — la borne est essentiellement atteinte
Il existe des sorties de routeur pour lesquelles \(M = \lceil E(R-1)/R^2 \rceil \approx E/R\).
Construction du pire cas : les experts du rang 0 ne reçoivent aucun jeton, et tous les experts des \(R-1\) autres rangs se partagent équitablement tous les jetons. Le rang 0 doit alors recevoir \(S\times K\) jetons, tous distants, qui impliquent au moins \(E(R-1)/R^2\) experts distincts.
Pourquoi ces deux théorèmes comptent en pratique
Réserver \(E/R\) emplacements d'experts redondants par rang garantit que la planification admet toujours une solution réalisable. Donc :
l'entraînement n'est jamais interrompu.
Comparaison donnée par le rapport avec les travaux antérieurs (ECHO, UltraEP) : ils présélectionnent un nombre d'experts redondants ou imposent un plafond de jetons par rang. L'entraînement est alors forcé de s'arrêter dès qu'aucun plan réalisable n'existe dans ce plafond — et le plafond lui-même exige un réglage manuel tout en laissant un déséquilibre résiduel.
MoonEP transforme un paramètre à régler en garantie prouvée.
La planification en ligne¶
Calculer l'optimum exact à chaque pas serait prohibitif. La démarche de Moonshot :
- calculer hors ligne des solutions exactes par programmation linéaire en nombres entiers (ILP) pour des cas représentatifs, servant de références ;
- concevoir un noyau de planification GPU qui est near-optimal, dont le surcoût est négligeable, et qui respecte toujours la borne \(E/R\).
La méthode générale
Résoudre exactement hors ligne, s'en servir de référence, puis concevoir une heuristique rapide et la mesurer contre cette référence. C'est un schéma d'ingénierie robuste, applicable bien au-delà du MoE.
Les trois bénéfices en cascade¶
L'équilibre parfait ne fait pas qu'équilibrer : il simplifie tout ce qui suit.
1. Communication sans copie¶
Un opérateur fusionné permute/unpermute : le noyau de planification précalcule la destination de chaque jeton, de sorte que les jetons sont envoyés directement à leurs positions groupées par expert sur les rangs distants, et que des vues du tampon de communication sont rendues directement au calcul. Aucune copie intermédiaire.
Le gain quantifié
Sous déséquilibre du pire cas, supporter le même chemin de données sans copie dans DeepEP exigerait un tampon de communication de taille \(S \times K \times R\).
MoonEP n'a besoin que d'un tampon fixe de \(S \times K\) — un facteur \(R\) d'économie, où \(R\) est la taille du groupe EP.
2. Exécution sans synchronisation, à formes statiques¶
Ce qui se passe sans équilibre parfait
Les comptes de jetons par expert varient d'un pas et d'une couche à l'autre. L'hôte doit se synchroniser avec l'appareil à chaque couche pour obtenir les formes réelles avant de lancer le calcul des experts — bloquant le pipeline entre couches.
Avec l'équilibre parfait, chaque rang reçoit exactement \(S\times K\) jetons : les formes de calcul de toutes les couches sont statiquement connues.
Résultat : suppression de la synchronisation hôte–MoE par couche, et réduction du surcoût de lancement de noyaux côté hôte.
3. Ordonnancement des GEMM d'experts¶
Un déséquilibre résiduel subsiste
Même avec la charge agrégée parfaitement équilibrée entre rangs, les comptes de jetons par expert à l'intérieur d'un rang restent asymétriques. Un ordonnancement à ordre fixe, aveugle à la charge, transformerait cette asymétrie en makespan déséquilibré entre les travailleurs SM.
Solution : un ordonnanceur conscient de la charge qui adapte ses paramètres à la distribution de jetons courante avant le lancement, puis les garde fixes pendant l'exécution. Une heuristique légère les sélectionne à partir d'un modèle de coût analytique des métriques matérielles, dont les coefficients clés sont calibrés par autotuning hors ligne.
Pour les experts partagés, les GEMM sont dispatchés sur un flux séparé, pour recouvrir d'autres noyaux.
Ce qui est ouvert¶
MoonEP est publié : github.com/MoonshotAI/MoonEP.
Pourquoi c'est significatif
C'est l'un des rares composants d'infrastructure de Kimi K3 qu'un tiers peut réellement auditer et réutiliser. Contrairement aux données ou aux hyperparamètres, c'est une contribution vérifiable.
Vérification de compréhension¶
Avec \(E = 896\) et \(R = 64\), combien d'experts redondants par rang ?
Au plus \(E/R = 14\). Chaque rang détient nativement 14 experts et doit réserver de la place pour 14 copies temporaires — soit un doublement de l'empreinte mémoire des experts dans le pire cas. C'est le coût de la garantie.
Que se passerait-il sans experts redondants du tout ?
On ne pourrait pas équilibrer : un jeton destiné à l'expert \(j\) doit être traité par un rang qui possède une copie de \(j\). Sans duplication, la charge de chaque rang est entièrement dictée par le routeur, sur lequel on n'a pas la main. Les experts redondants sont le seul degré de liberté.
Pourquoi la migration des gradients est-elle plus délicate que celle des poids ?
Les poids se copient (lecture seule). Les gradients doivent être réduits (additionnés) vers le propriétaire de l'expert, sinon les contributions calculées sur les copies seraient perdues. D'où le tampon de réduction local puis la réduction vers les rangs d'origine, décrits par le rapport.
Chapitre précédent : KDA Context Parallelism · Chapitre suivant : Mémoire et parallélismes