huggingface.co détecté sur le web

Pré-entraînement async : Muon survit là où AdamW s'effondre

TL;DR

  • Sur SmoLLM-360M, le délai d'un pas coûte 0,012 de perte à Muon contre 0,278 à AdamW, selon les benchmarks du papier.
  • La correction Error Feedback (x_{t+1} = x_t - 2·u_{t-1} + u_{t-2}) récupère 50 à 70 % de la dégradation, 85 à 90 % pour AdamW.
  • Sur un MoE de 10 milliards entraîné sur 200 milliards de tokens, async + Error Feedback égale exactement la perte synchrone de 1,906.

Il y a dans la dernière fournée de papiers sur l'entraînement distribué un résultat qui mérite qu'on s'y arrête, parce qu'il déplace une hypothèse de fond. Le délai d'un pas dans un pipeline parallèle asynchrone, longtemps considéré comme une limite intrinsèque qui dégrade la qualité du modèle, dépendrait en réalité surtout du choix d'optimiseur. C'est la thèse défendue par Zmushko, Petrov, Abdullaev, Khrushchev et Horváth dans « One-Step Gradient Delay is Not a Barrier for Large-Scale Asynchronous Pipeline Parallel LLM Pretraining ».

Le contraste empirique est net. Sur un SmoLLM-360M, AdamW encaisse une dégradation de perte de 0,278 sous staleness d'un pas, là où Muon ne perd que 0,012. Adan, SOAP, Lion et plusieurs variantes Muon restent eux aussi sous la barre des 0,03. Les auteurs proposent en complément une correction inspirée de l'Error Feedback qui s'applique au niveau de l'update, de la forme x_{t+1} = x_t - 2·u_{t-1} + u_{t-2}, et qui récupère 50 à 70 % de l'écart pour les optimiseurs déjà robustes, 85 à 90 % pour AdamW et MARS, pour un surcoût mémoire jugé négligeable.

L'intérêt pratique se joue à grande échelle. Pour valider la méthode, l'équipe a entraîné un MoE de 10 milliards de paramètres sur 200 milliards de tokens. La perte de validation synchrone de référence est de 1,906 ; l'async standard atteint 1,911 ; l'async avec Error Feedback retombe exactement sur 1,906. Les auteurs présentent ce point comme la première démonstration d'async pipeline parallel à cette échelle sans écart de qualité, avec les mêmes hyperparamètres que le run synchrone.

Le caveat honnête, c'est que le mécanisme exact par lequel le momentum amortit la staleness n'est pas élucidé : c'est listé comme limite par les auteurs eux-mêmes, au même titre que l'absence de validation au-delà de 200 milliards de tokens, un balayage de taille de batch restreint au 135M, et l'exploration limitée des schedules WPipe.

Pour les équipes qui s'entraînent déjà sous Muon, le message est exploitable tout de suite : la fenêtre pour basculer en pipeline asynchrone sans payer en qualité finale s'ouvre, à condition de partir sur PipeDream-2BW à délai constant. Pour celles qui restent sur AdamW, le coût d'opportunité d'une migration d'optimiseur vient de gagner une ligne nouvelle dans la feuille de calcul.