🤖 machine learning

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

El artículo presenta DMuon, una implementación distribuida de código abierto del optimizador Muon que se integra sin problemas en los flujos de trabajo de entrenamiento existentes para reducir drásticamente la latencia del paso del optimizador y lograr una eficiencia cercana a AdamW, permitiendo así el uso escalable de la optimización basada en ortogonalización de matrices en grandes modelos fundacionales.

Autores originales: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Publicado 2026-06-26

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás entrenando el cerebro de un robot gigante (un Modelo de Lenguaje Grande o una IA con cuerpo físico) para que aprenda. Para enseñarle, necesitas un "optimizador": un entrenador que observe los errores del robot y ajuste sus conexiones internas (pesos) para hacerlo más inteligente.

Durante años, el entrenador estándar ha sido AdamW. Es como un equipo de miles de trabajadores diminutos, cada uno arreglando un pequeño tornillo del robot a la vez. Es rápido y eficiente, pero trata cada tornillo de forma independiente.

Recientemente, se descubrió un nuevo y más inteligente entrenador llamado Muon. En lugar de arreglar tornillos uno por uno, Muon observa paneles enteros del robot (matrices completas de datos) y los ajusta todos juntos. Este enfoque de "terapia de grupo" ayuda al robot a aprender más rápido y a alcanzar un mayor rendimiento. Sin embargo, hay un inconveniente: Muon es increíblemente lento de ejecutar en un sistema distribuido (un grupo de muchas computadoras trabajando juntas).

El Problema: El Cuello de Botella de la "Reunión de Todo el Personal"
En una configuración de entrenamiento distribuido, el cerebro del robot se divide entre muchas computadoras (GPUs).

AdamW funciona como un equipo remoto: La Computadora A arregla sus tornillos, la Computadora B arregla los suyos, y no necesitan hablar mucho entre sí.
Muon funciona como un comité: Para arreglar un panel, cada computadora necesita ver el panel entero primero.

En una implementación ingenua de Muon, cada una de las computadoras tiene que detener lo que está haciendo, descargar los datos completos de todos los demás, realizar la matemática compleja y luego enviar los resultados de vuelta. Esto es como celebrar una reunión masiva donde todos leen un informe de 1,000 páginas antes de tomar una sola decisión. ¡Toma tanto tiempo que la "reunión" (el paso del optimizador) tarda más que el trabajo real (los pasos de aprendizaje)! De hecho, el artículo señala que esto puede costar 2 veces más tiempo que el propio aprendizaje.

La Solución: DMuon (Muon Distribuido)
El Equipo de Robots X Square construyó DMuon, un nuevo sistema que hace que Muon funcione casi tan rápido como AdamW. Lo lograron cambiando cómo trabaja el equipo, no cambiando la matemática del entrenador.

Así es como resolvieron el cuello de botella usando tres trucos principales:

1. La Estrategia del "Experto Designado" (Ejecución Centrada en el Propietario)

En lugar de que todos intenten arreglar cada panel, DMuon asigna a una computadora específica para ser el "Propietario" de cada panel.

Antes: Todos reunían los datos, todos hacían la matemática, todos perdían tiempo.
Ahora: Solo la computadora "Propietaria" reúne los datos y realiza la matemática pesada. Las otras computadoras simplemente esperan o realizan otro trabajo.
Analogía: Imagina un sitio de construcción. En lugar de que cada trabajador intente construir toda la casa, se asigna un experto para el techo y otro para la plomería. Los demás no intentan construir el techo; solo le pasan al experto las herramientas que necesita. Esto evita que todos hagan el mismo trabajo dos veces.

2. El Sistema de la "Cinta Transportadora" (Solapamiento de Comunicación)

Mientras el "Propietario" está haciendo la matemática, las otras computadoras no están simplemente inactivas. DMuon crea una tubería:

Paso hacia adelante (Forward Pass): Mientras el robot está "pensando" (procesando datos), el sistema envía silenciosamente el siguiente conjunto de herramientas a los expertos en segundo plano.
Paso hacia atrás (Backward Pass): Mientras el robot está "aprendiendo de sus errores", el sistema ya está reuniendo el siguiente lote de datos para los expertos.
Analogía: Es como la cocina de un restaurante. El chef (el Propietario) está cortando vegetales. En lugar de esperar a que el chef termine antes de traer la siguiente orden, el mesero trae la siguiente orden mientras el chef todavía está cortando. El tiempo de espera se oculta dentro del tiempo de corte.

3. La "Línea de Ensamblaje Inteligente" (Agrupación y Ajuste)

La matemática que hace Muon es compleja. A veces las piezas son enormes (como una pared gigante) y otras veces son diminutas (como una pequeña baldosa).

El Problema: Si intentas procesar una baldosa diminuta en una máquina gigante, la máquina se queda ociosa. Si procesas una pared gigante, toma una eternidad.
La Solución: DMuon agrupa muchas baldosas diminutas en un solo "lote" para que la máquina se mantenga ocupada. También utiliza un "ajustador" que encuentra automáticamente la forma más rápida de cortar cada forma específica de baldosa.
Analogía: En lugar de enviar un camión de reparto para un solo sobre, DMuon espera hasta tener un camión lleno de sobres y los envía todos a la vez. También cambia la ruta del camión dependiendo de si el destino es una ciudad o una granja.

Los Resultados

El artículo probó DMuon en modelos del mundo real, incluyendo un modelo de entrenamiento de robots (Wall-OSS) y un modelo de lenguaje (Qwen2.5).

Velocidad: DMuon hizo que el "paso del optimizador" (la reunión) fuera de 6 a 163 veces más rápido que la antigua forma ingenua.
Eficiencia: El tiempo total para entrenar un paso es ahora solo un 2% más lento que el método estándar AdamW.
Conclusión: DMuon permite a los equipos utilizar el estilo de aprendizaje superior de "terapia de grupo" de Muon sin pagar la enorme penalización de tiempo. Convierte un proceso lento y tosco en uno fluido y eficiente, haciéndolo apto para su uso en la producción de modelos de IA masivos.

En resumen, DMuon es el sistema de control de tráfico que permite que el entrenador más inteligente (Muon) corra en una autopista de computadoras sin causar un embotellamiento, haciendo que finalmente sea práctico entrenar la próxima generación de IA.

Resumen Técnico: DMuon – Entrenamiento de Muon Distribuido Eficiente con una Sobrecarga Cercana a Adam

1. Declaración del Problema

Los optimizadores conscientes de las matrices, particularmente Muon, han demostrado propiedades de convergencia superiores en comparación con los optimizadores elemento a elemento como AdamW, al aplicar iteraciones de Newton–Schulz (NS) a los gradientes agregados de momento de las matrices de pesos. Este enfoque impulsa los valores singulares de las actualizaciones hacia un espectro uniforme, ofreciendo ganancias significativas en la eficiencia computacional en el preentrenamiento de modelos grandes (por ejemplo, LLMs y modelos fundacionales para robótica/embodied).

Sin embargo, el despliegue de Muon en entornos de entrenamiento distribuido modernos enfrenta un cuello de botella crítico de sistemas. Los stacks de entrenamiento distribuido contemporáneos (por ejemplo, ZeRO, FSDP) operan bajo un contrato implícito de optimización elemento a elemento, donde las actualizaciones se computan de forma independiente en los fragmentos (shards) de los parámetros. Muon viola este contrato porque su iteración central de NS requiere que la matriz de peso completa esté presente en un solo dispositivo para computar los productos matriciales necesarios (específicamente $XX^\top$ ).

En una implementación distribuida ingenua ("Gather-then-Compute"), cada rango debe:

Realizar un All-gather de la matriz de gradiente reducida completa desde todos los fragmentos.
Realizar la iteración NS localmente.
Descartar los resultados para los fragmentos que no le pertenecen.

Este enfoque introduce dos costos principales:

Sobrecarga de Comunicación: La reconstrucción de matrices completas en cada paso introduce comunicación colectiva que escala con el tamaño del modelo y la anchura distribuida.
Computación Redundante: Cada rango realiza la misma ortogonalización costosa, multiplicando el costo computacional por el número de rangos ( $D$ ).

En la práctica, estos costos causan que el tiempo del paso del optimizador en el Muon distribuido iguale o incluso supere el tiempo de ejecución (wall-clock time) combinado de los pasos forward y backward. Esta penalización es particularmente severa para el entrenamiento de modelos para robótica (VLA), donde los pasos forward-backward son computacionalmente más ligeros en relación con el paso del optimizador comparado con el preentrenamiento estándar de LLMs.

2. Metodología: Diseño del Sistema DMuon

DMuon es una implementación distribuida de código abierto diseñada para cerrar la brecha de rendimiento entre Muon y AdamW. Logra esto mediante un sistema co-diseñado que integra optimización de comunicación, equilibrio de carga y eficiencia a nivel de kernel, sin modificar el framework de entrenamiento subyacente (por ejemplo, PyTorch/FSDP2).

2.1 Estrategia de Ejecución Centrada en el Propietario (Owner-Centric)

En lugar de replicar la computación en todos los rangos, DMuon asigna cada parámetro de matriz a un único rango propietario (owner rank).

Rol del Propietario: El propietario mantiene el estado autoritativo del optimizador y computa la actualización de Muon (iteración Newton–Schulz) exactamente una vez por matriz.
Rol del No-Propietario: Los otros rangos materializan el parámetro solo cuando es requerido para la computación forward/backward y no realizan la ortogonalización.
Flujo de Comunicación:
- Forward: Los propietarios publican los parámetros actualizados a los no-propietarios de forma asíncrona, en pipeline con la ejecución de la capa.
- Backward: Los gradientes se acumulan en buffers empaquetados y se reducen hacia el propietario.
- Solapamiento (Overlap): DMuon emplea una jerarquía de dos etapas (intra-nodo e inter-nodo) con layouts de pesos de grano fino para solapar la comunicación con la computación, evitando la contención colectiva.

2.2 Iteración Gram Newton–Schulz Eficiente

Para reducir el costo computacional de la actualización en el lado del propietario, DMuon adopta la formulación Gram Newton–Schulz:

Recurrencia en el Espacio Gram: En lugar de operar en el espacio original $m \times n$ , la iteración se realiza en el espacio Gram $m \times m$ ( $G = XX^\top$ ). Esto reduce la complejidad de $O(m^2n)$ a $O(m^3)$ cuando $m < n$ .
Explotación de Simetría: Dado que $G$ es simétrica, DMuon utiliza una ruta de ejecución tipo SYRK para computar solo la parte triangular inferior, reduciendo a la mitad el trabajo aritmético de los productos matriciales dominantes.
Batching: Para abordar los problemas de ocupación de GPU con matrices pequeñas, DMumo agrupa matrices de formas similares y ejecuta la recurrencia Gram NS conjuntamente, amortizando la sobrecarga de despacho de kernels.

2.3 Equilibrio de Carga Consciente de la Computación

Debido a que los costos de actualización de Muon varían significamente con la forma de la matriz, una asignación ingenua (por ejemplo, round-robin) genera rangos rezagados (stragglers). DMuon formula la asignación de propietarios como un Programa Lineal de Enteros Mixtos (MILP):

Perfilado (Profiling): Durante la inicialización, el sistema mide el tiempo de ejecución de varios tamaños de batch y esquemas de kernel para cada forma de matriz distinta.
Optimización: El solver MILP minimiza el makespan predicho (carga máxima) entre los rangos propietarios, asegurando que las matrices más pesadas se distribuyan para equilibrar la carga de trabajo.
Fallback: Para espacios de búsqueda extremadamente grandes, se utiliza una heurística voraz (greedy) para limitar el tiempo de inicialización.

2.4 Autotuning de Kernels Adaptativo a la Forma

DMuon utiliza un flujo de trabajo de autotuning impulsado por un DSL (usando TileLang y CUTE DSL) para generar kernels de alto rendimiento.

Busca sobre formas de tiles, tamaños de bloques y configuraciones de pipeline de software específicas para el hardware objetivo y las formas de las matrices.
La configuración óptima se almacena en caché y se reutiliza para pasos subsecuentes, asegurando que el esquema del kernel sea siempre cercano al óptimo para las dimensiones de problema específicas encontradas en el entrenamiento de transformers.

3. Contribuciones Clave

Optimización de Comunicación de Grano Fino: Un runtime centrado en el propietario que minimiza la sobrecarga de comunicación mediante programación jerárquica y layouts de peso de grano fino, preservando la semántica exacta del optimizador.
Stack de Ejecución Adaptativo a la Forma: Un motor de alto rendimiento que combina Gram Newton–Schulz por lotes, kernels conscientes de la simetría y autotuning para manejar eficientemente cargas de trabajo de matrices heterogéneas.
Equilibrio de Carga Consciente de la Computación: Un enfoque de optimización basado en costos medidos (MILP) que asigna propietarios de matrices basándose en el perfilado de tiempo de ejecución real en lugar de conteos teóricos de FLOPs.
Integración Directa (Drop-in): Un módulo que se integra en pipelines de entrenamiento existentes (por ejemplo, FSDP2) con un esfuerzo mínimo, sin requerir modificaciones a nivel de framework.
Validación a Escala de Producción: Despliegue exitoso tanto en cargas de trabajo de robótica (modelos fundacionales para robótica) como de modelos de lenguaje.

4. Resultados Experimentales

Los autores evaluaron DMuon en cuatro cargas de trabajo de producción: Wall-OSS-0.5, Pi0, Wall-WM y Qwen2.5-7B, escalando de 8 a 256 GPUs A800.

Tiempo del Paso del Optimizador: DMuon logra una aceleración de 6.85× a 163.00× en el tiempo del paso del optimizador en comparación con una implementación ingenua de Muon "Gather-then-Compute".
Tiempo de Paso de Extremo a Extremo (End-to-End): DMuon logra una aceleración de 1.48× a 3.01× en el tiempo total del paso.
Sobrecarga vs. AdamW: La latencia por paso de DMuon se reduce a niveles cercanos a AdamW, con una sobrecarga promedio de solo +2% en comparación con AdamW en todas las cargas de trabajo evaluadas.
Escalabilidad: En Wall-OSS, DMuon mantiene un throughput cercano a AdamW a medida que aumenta el número de GPUs, mientras que el Muon ingenuo sufriría de severos cuellos de botella de comunicación.

Desglose de la Aceleración:

Kernel Gram Simétrico: 48% de la aceleración (reducción a la mitad de la aritmética en productos dominantes).
Programación de Propietarios y Equilibrio de Carga: 32% (eliminación de la computación redundante $D$ -veces).
Autotuning y Batching: 16% (reducción de la sobrecarga de lanzamiento de kernels para matrices pequeñas).

5. Significado y Reivindicaciones

El artículo afirma que DMuon resuelve con éxito el desajuste de sistemas entre los optimizadores conscientes de matrices y la infraestructura de entrenamiento distribuido. Al eliminar la computación redundante y optimizar la comunicación, DMuon permite que los beneficios de convergencia de Muon (como una mayor eficiencia computacional y una mejor geometría de optimización) se traduzcan en aceleraciones prácticas de tiempo de ejecución (wall-clock speedups).

Los autores posicionan a DMuon no como un nuevo algoritmo, sino como un habilitador de sistemas necesario que hace que Muon sea un reemplazo directo (drop-in replacement) para AdamW en el entrenamiento a gran escala de producción. El trabajo demuestra que la "penalización" de la optimización consciente de matrices puede reducirse a niveles insignificantes, permitiendo el escalado eficiente tanto para Grandes Modelos de Lenguaje como para Modelos Fundacionales para Robótica (VLA).

El artículo señala explícitamente que DMuon no cambia la regla de actualización ni las propiedades de convergencia de Muon; aborda estrictamente la sobrecarga del sistema. Por consiguiente, los beneficios son más pronunciados en entornos distribuidos, mientras que en el entrenamiento de una sola GPU se observa una mejora más modesta (aproximadamente 2× de aceleración en el optimizador) debido a la falta de redundancia distribuida para eliminar.