Résumé Technique : DMuon – Un entraînement distribué efficace de Muon avec un surcoût proche d'Adam

1. Énoncé du Problème

Les optimiseurs sensibles aux matrices, particulièrement Muon, ont démontré des propriétés de convergence supérieures par rapport aux optimiseurs élément par élément comme AdamW en appliquant des itérations de Newton-Schulz (NS) aux gradients agrégés par moment des matrices de poids. Cette approche pousse les valeurs singulières des mises à jour vers un spectre uniforme, offrant des gains d'efficacité de calcul significatifs lors du préentraînement de grands modèles (ex. : LLM et modèles de fondation incarnés).

Cependant, le déploiement de Muon dans les environnements d'entraînement distribués modernes fait face à un goulot d'étranglement critique au niveau des systèmes. Les piles d'entraînement distribuées contemporaines (ex. : ZeRO, FSDP) opèrent sous un contrat implicite d'optimisation élément par élément, où les mises à jour sont calculées indépendamment sur les fragments de paramètres. Muon viole ce contrat car son itération NS centrale nécessite la présence de la matrice de poids complète sur un seul dispositif pour calculer les produits matriciels nécessaires (spécifiquement $XX^\top$).

Dans une implémentation distribuée naïve (« Gather-then-Compute » ou « Rassembler-puis-Calculer »), chaque rang doit :

1. Effectuer un All-gather de la matrice de gradient réduite complète à partir de tous les fragments.
2. Effectuer l'itération NS localement.
3. Écarter les résultats pour les fragments qu'il ne possède pas.

Cette approche introduit deux coûts majeurs :

• Surcharge de Communication : La reconstruction de matrices complètes à chaque étape introduit une communication collective qui évolue avec la taille du modèle et la largeur distribuée.
• Calcul Redondant : Chaque rang effectue la même orthogonalisation coûteuse, multipliant le coût de calcul par le nombre de rangs ($D$).

En pratique, ces coûts font que le temps de l'étape de l'optimiseur dans le cas de Muon distribué rivalise avec ou dépasse le temps de calcul combiné des passes avant et arrière, érodant les avantages algorithmiques de Muon. Cette pénalité est particulièrement sévère pour l'entraînement de modèles incarnés (ex. : VLA), où les passes avant-arrière sont plus légères par rapport à l'étape de l'optimiseur comparativement au préentraînement standard des LLM.

2. Méthodologie : Conception Système de DMuon

DMuon est une implémentation distribuée open-source conçue pour combler l'écart de performance entre Muon et AdamW. Elle y parvient grâce à une architecture système co-conçue qui intègre l'optimisation de la communication, l'équilibrage de charge et l'efficacité au niveau des noyaux (kernels), sans modifier le framework d'entraînement sous-jacent (ex. : PyTorch/FSDP2).

2.1 Stratégie d'Exécution Centrée sur le Propriétaire (Owner-Centric)

Au lieu de répliquer le calcul sur tous les rangs, DMuon assigne chaque paramètre matriciel à un rang propriétaire (owner rank) unique.

• Rôle du Propriétaire : Le propriétaire maintient l'état d'optimisation faisant autorité et calcule la mise à jour Muon (itération Newton-Schulz) exactement une fois par matrice.
• Rôle du Non-Propriétaire : Les autres rangs ne matérialisent le paramètre que lorsqu'il est requis pour le calcul avant/arrière et ne réalisent pas l'orthogonalisation.
• Flux de Communication :
  • Forward (Avant) : Les propriétaires publient les paramètres mis à jour vers les non-propriétaires de manière asynchrone, en pipeline avec l'exécution des couches.
  • Backward (Arrière) : Les gradients sont accumulés dans des tampons compactés et réduits vers le propriétaire.
  • Chevauchement (Overlap) : DMuon emploie une hiérarchie à deux étages (intra-nœud et inter-nœud) avec des structures de poids fines pour faire chevaucher la communication avec le calcul, évitant ainsi la contention des collectifs.

2.2 Itération de Newton-Schulz de type Gram Efficace

Pour réduire le coût de calcul de la mise à jour côté propriétaire, DMuon adopte la formulation Gram Newton-Schulz :

• Récurrence dans l'Espace Gram : Au lieu d'opérer dans l'espace $m \times n$ original, l'itération est effectuée dans l'espace Gram $m \times m$ ($G = XX^\top$). Cela réduit la complexité de $O(m^2n)$ à $O(m^3)$ lorsque $m < n$.
• Exploitation de la Symétrie : Puisque $G$ est symétrique, DMuon utilise un chemin d'exécution de type SYRK pour ne calculer que la partie triangulaire inférieure, divisant par deux le travail arithmétique des produits matriciels dominants.
• Batching (Regroupement) : Pour remédier aux problèmes d'occupation GPU avec de petites matrices, DMuon regroupe les matrices de formes similaires et exécute la récurrence Gram NS conjointement, amortissant ainsi le surcoût de lancement des noyaux.

2.3 Équilibrage de Charge Sensible au Calcul

Parce que les coûts de mise à jour de Muon varient considérablement selon la forme des matrices, une assignation naïve (ex. : tour de rôle) entraîne des rangs retardataires (stragglers). DMuon formule l'assignation des propriétaires comme un Problème de Programmation Linéaire en Nombres Entiers Mixtes (MILP) :

• Profilage : Pendant l'initialisation, le système mesure le temps d'exécution de diverses tailles de lots et de programmes de noyaux pour chaque forme de matrice distincte.
• Optimisation : Le solveur MILP minimise l'étendue maximale (makespan) prédite à travers les rangs propriétaires, garantissant que les matrices les plus lourdes sont distribuées pour équilibrer la charge de travail.
• Repli (Fallback) : Pour les espaces de recherche extrêmement vastes, une heuristique gloutonne est utilisée pour limiter le temps d'initialisation.

2.4 Autotuning de Noyau Adaptatif à la Forme

DMuon utilise un flux de travail d'autotuning piloté par un DSL (utilisant TileLang et le DSL CUTE) pour générer des noyaux de haute performance.

• Il recherche parmi les formes de tuiles (tiles), les tailles de blocs et les configurations de pipeline logiciel spécifiques au matériel cible et aux formes de matrices.
• La configuration optimale est mise en cache et réutilisée pour les étapes suivantes, garantissant que le programme du noyau est toujours proche de l'optimal pour les dimensions de problèmes spécifiques rencontrées lors de l'entraînement de transformers.

3. Contributions Clés

1. Optimisation de Communication à Grain Fin : Un runtime centré sur le propriétaire qui minimise le surcoût de communication via un ordonnancement hiérarchique et des structures de poids à grain fin, préservant la sémantique exacte de l'optimiseur.
2. Pile d'Exécution Adaptative à la Forme : Un moteur de haute performance combinant le Newton-Schulz de type Gram par lot, des noyaux sensibles à la symétrie et un autotuning pour gérer efficacement les charges de travail matricielles hétérogènes.
3. Équilibrage de Charge Sensible au Calcul : Une approche d'optimisation basée sur le coût mesuré (MILP) qui assigne les propriétaires de matrices en fonction du profilage de l'exécution réelle plutôt que des simples comptes théoriques de FLOPs.
4. Intégration Transparente (Drop-in) : Un module qui s'intègre dans les pipelines d'entraînement existants (ex. : FSDP2) avec un effort minimal, ne nécessitant aucune modification du framework.
5. Validation à Échelle de Production : Déploiement réussi sur des charges de travail de modèles de langage et de robotique (modèles incarnés).

4. Résultats Expérimentaux

Les auteurs ont évalué DMuon sur quatre charges de travail de production : Wall-OSS-0.5, Pi0, Wall-WM, et Qwen2.5-7B, en passant de 8 à 256 GPU A800.

• Temps de l'Étape de l'Optimiseur : DMuon atteint une accélération de 6,85× à 163,00× du temps de l'étape de l'optimiseur par rapport à une implémentation naïve de Muon distribuée de type « Gather-then-Compute ».
• Temps de l'Étape de Bout en Bout : DMument atteint une accélération de 1,48× à 3,01× du temps total de l'étape.
• Surcoût vs AdamW : La latence par étape de DMuon est ramenée à des niveaux proches d'AdamW, avec un surcoût moyen de seulement +2 % par rapport à AdamW sur l'ensemble des charges de travail évaluées.
• Passage à l'Échelle (Scaling) : Sur Wall-OSS, DMuon maintient un débit proche de celui d'AdamW à mesure que le nombre de GPU augmente, là où un Muon naïf souffrirait de graves goulots d'étranglement de communication.

Répartition de l'Accélération :

• Noyau Gram Symétrique : 48 % de l'accélération (division par deux de l'arithmétique dans les produits dominants).
• Ordonnancement des Propriétaires et Équilibrage de Charge : 32 % (suppression du calcul redondant $D$-fois).
• Autotuning et Batching : 16 % (réduction du surcoût de lancement de noyau pour les petites matrices).

5. Signification et Revendications

L'article affirme que DMuon résout avec succès l'incompatibilité système entre les optimiseurs sensibles aux matrices et les infrastructures d'entraînement distribuées. En éliminant le calcul redondant et en optimisant la communication, DMuon permet aux avantages de convergence de Muon (tels qu'une meilleure efficacité de calcul et une meilleure géométrie d'optimisation) de se traduire en accélérations concrètes du temps de calcul réel.

Les auteurs positionnent DMuon non pas comme un nouvel algorithme, mais comme un facilitateur système nécessaire qui rend Muon un remplacement direct (drop-in replacement) viable pour AdamW dans l'entraînement de production à grande échelle. Ce travail démontre que la « pénalité » de l'optimisation sensible aux matrices peut être réduite à des niveaux négligeables, permettant une mise à l'échelle efficace tant pour les Grands Modèles de Langage que pour les Modèles de Fondations Incarnés (VLA).

Le document note explicitement que DMuon ne modifie pas la règle de mise à jour ni les propriétés de convergence de Muon ; il traite strictement du surcoût système. Par conséquent, les bénéfices sont plus marqués dans les contextes distribués, tandis que l'entraînement sur un seul GPU voit une amélioration plus modeste (environ 2× l'accélération de l'optimiseur) en raison de l'absence de redondance distribuée à éliminer.