NVIDIA impulsa optimizadores avanzados para mejorar el entrenamiento de LLMs con Megatron

Los algoritmos de optimización de orden superior, como Shampoo, han demostrado ser efectivos en el entrenamiento de redes neuronales durante al menos una década. Su éxito ha aumentado recientemente al aplicarse a los modelos LLM más avanzados. En particular, Muon (MomentUm Orthogonalized by Newton-Schulz) se ha utilizado para entrenar algunos de los mejores modelos de código abierto actuales, incluyendo Kimi K2 y GLM-5.

Este artículo explica cómo NVIDIA ofrece un soporte integral para Muon y otros optimizadores emergentes, así como las tecnologías que permiten entrenar modelos a gran escala.

Rendimiento del entrenamiento de Muon en NVIDIA GB300 NVL72

La Tabla 1 resume el rendimiento de entrenamiento de los modelos Kimi K2 y Qwen3 30B utilizando Muon y el optimizador AdamW en el sistema NVIDIA GB300 NVL72. Los resultados muestran que la pérdida de rendimiento al usar el optimizador Muon es mínima en comparación con AdamW. La utilización de FLOPs del modelo (MFU) es mayor con Muon al considerar los FLOPs de la multiplicación de matrices en las iteraciones de Newton-Schulz.

Estas mediciones se lograron utilizando NVIDIA NeMo Megatron Bridge 26.02, una biblioteca nativa de PyTorch dentro del NeMo Framework que proporciona preentrenamiento, SFT y LoRA para modelos populares de LLM y VLM. El equipo de NVIDIA utilizó 256 GPUs NVIDIA GB300 con PP4DP64EP64 para el entrenamiento de Kimi K2 y ocho GPUs NVIDIA GB300 con DP8EP8 para el entrenamiento de Qwen3 30B-A3B.

Para más detalles sobre la configuración experimental y los pasos para reproducir la medición, consulte la sección de inicio rápido a continuación.

Tecnologías habilitadoras para el entrenamiento a gran escala de Muon

Aunque los optimizadores de orden superior como Muon prometen mejorar la eficiencia del entrenamiento de LLM, su implementación práctica a gran escala enfrenta varios desafíos. El paso de precondicionamiento, que utiliza iteraciones de Newton-Schulz o descomposición en valores propios, aumenta significativamente el costo computacional y el consumo de memoria. Además, puede surgir inestabilidad numérica durante el entrenamiento de precisión mixta y acumulación de gradientes. Distribuir eficientemente las actualizaciones sincronizadas y ortogonalizadas entre miles de GPUs puede introducir cuellos de botella de comunicación que disminuyan los beneficios de eficiencia generales.

Esta sección introduce tecnologías para superar estos obstáculos y escalar a miles de GPUs. Las tecnologías seleccionadas están equilibradas entre generalidad, rendimiento y complejidad de implementación. Se favorecen tecnologías que se pueden generalizar a Muon, SOAP (ShampoO with Adam in the Preconditioner’s eigenbasis), y otros optimizadores complejos en lugar de centrarse únicamente en Muon, para el cual se aplicarían más optimizaciones debido a su simplicidad.

Optimizador distribuido por capas

Un optimizador distribuido por elementos (comúnmente aplicado a optimizadores como AdamW a gran escala) ofrece lo siguiente:

• Particionamiento de estados: En lugar de que cada GPU almacene todo, los estados del optimizador se distribuyen uniformemente entre todas las GPUs disponibles.
• Reducción y dispersión de gradientes: Se realiza una reducción y dispersión sobre todos los gradientes y cada GPU recibe una porción de los gradientes correspondientes a los parámetros que “posee”.
• Actualizaciones locales: Cada GPU actualiza solo la porción específica de los parámetros del modelo que “posee”.
• AllGather de parámetros: Después de la actualización, las GPUs se comunican para asegurar que todos los dispositivos tengan los pesos actualizados del modelo completo para el siguiente pase hacia adelante.

Una versión más avanzada descompone operaciones y superpone comunicación con cálculos.

Los optimizadores distribuidos por elementos han estado funcionando bien con optimizadores Adam. Sin embargo, muchos optimizadores emergentes, incluido Muon, requieren gradientes de toda la capa para calcular actualizaciones a los pesos de esa capa. Si los pesos y los estados del optimizador están distribuidos uniformemente entre los rangos de paralelismo de datos (DP), no se pueden calcular actualizaciones basadas en los datos disponibles en cada GPU. Se necesitará comunicación adicional para recopilar datos para calcular la actualización completa.

Se introduce un optimizador distribuido por capas para soportar este tipo de optimizador. Los parámetros de diferentes capas se distribuyen a diferentes rangos DP. Cada GPU tiene parámetros completos de capas para que se pueda calcular el precondicionador.

Un cambio que viene con la distribución por capas es la comunicación de tamaño variable. Debido a que las capas completas varían en tamaño, cada GPU necesita recopilar actualizaciones de parámetros de diferentes tamaños de diferentes GPUs a través de all_gatherv.

Este optimizador distribuido por capas está completamente integrado en NVIDIA Megatron Core, una biblioteca de código abierto para construir y entrenar modelos a gran escala con paralelismo avanzado, precisión mixta y núcleos de GPU optimizados. La implementación se puede encontrar en layer_wise_optimizer.py.

Newton-Schulz distribuido

El entrenamiento de LLM a gran escala a menudo utiliza paralelismo de tensores (TP) junto con otros paralelismos, lo que introduce desafíos únicos. TP divide matrices de pesos individuales entre múltiples GPUs a lo largo de dimensiones específicas, lo que significa que ningún dispositivo tiene el tensor completo de parámetros. El paso crítico de ortogonalización del optimizador Muon requiere acceso a toda la matriz de impulso para su cálculo. Cuando este búfer de impulso se divide entre dispositivos, se necesita comunicación adicional para manejar el impulso y los pesos fragmentados.

Se proporcionan dos métodos para manejar tales escenarios en TensorParallelMuon: modo duplicado y modo distribuido.

Modo duplicado

Los momentos se recopilan primero entre el dominio TP para que cada GPU tenga todos los datos para la iteración de Newton-Schulz. Luego, cada GPU realiza iteraciones completas de Newton-Schulz y actualiza los pesos que posee mediante el momento ortogonalizado correspondiente. Este modo está optimizado para la latencia de red a costa de que cada GPU realice las mismas iteraciones de NS. Para cada actualización, independientemente de cuántas interacciones de NS se utilicen (comúnmente cinco), solo se necesita una comunicación de recopilación al principio.

Modo distribuido

Como su nombre indica, el cálculo de las iteraciones de NS se distribuye entre GPUs en el dominio TP. Cada interacción de NS realiza tres multiplicaciones de matrices, y se necesita una reducción de todos después de la primera multiplicación de matrices de cada iteración. Este modo está optimizado para el cálculo a costa de una comunicación más frecuente.

Además, se admite un modo por bloques, que realiza la ortogonalización y actualiza solo con el momento y los pesos que posee cada GPU. Es más rentable en términos de cálculo y no requiere comunicación en comparación con los otros dos modos. Sin embargo, no es equivalente a ortogonalizar toda la matriz de impulso al mismo tiempo. Matemáticamente, es conceptualmente una ortogonalización por bloques (de ahí el nombre blockwise), similar al enfoque de bloques introducido por primera vez en Scalable Second Order Optimization for Deep Learning.

Optimización adicional

Otras optimizaciones pueden aumentar aún más el rendimiento de entrenamiento. Algunas importantes se discuten brevemente en esta sección y se introducirán gradualmente a través del proyecto de investigación de NVIDIA Emerging Optimizers.

Ocultamiento de comunicación

En la versión actual del proyecto NVIDIA Emerging Optimizers, los parámetros se recopilan inmediatamente después del paso del optimizador y se exponen completamente, lo que podría convertirse en un cuello de botella en ciertas circunstancias. La misma técnica utilizada en optimizadores distribuidos por elementos puede aplicarse; la recopilación de parámetros se puede retrasar hasta el paso hacia adelante del siguiente lote y superponerse con el cálculo.

Equilibrio de carga

Debido a los diversos tamaños de las matrices de pesos de atención y MLP, el costo computacional del optimizador debido al precondicionamiento también varía entre capas. No se puede lograr un equilibrio de carga perfecto ni en teoría. Actualmente, las capas se distribuyen entre GPUs en el mismo dominio DP de manera rotativa. Existen métodos que pueden lograr un mejor equilibrio de carga basados en alguna estimación del costo computacional y de comunicación.

SYRK y reducción fusionada

Las dos primeras de las tres multiplicaciones de matrices de una interacción de NS se pueden mapear a SYRK (actualización de rango K simétrica), de modo que se pueden ahorrar casi la mitad de las operaciones de punto flotante. El ahorro es ligeramente inferior a la mitad porque los bloques a lo largo de la diagonal deben calcularse completamente. Los núcleos de Triton se proporcionan en la versión actual.

SYRK produce una matriz completa en lugar de triángulo (o inferior) en formato compacto. En el modo distribuido, la reducción de todos aún necesita comunicar la matriz completa en lugar de la mitad. Una versión más avanzada puede fusionar la comunicación en el núcleo de SYRK, lo que no solo ahorra ancho de banda, sino que también puede ocultar la comunicación a un nivel más fino (es decir, a nivel de bloque).

¿Qué otros optimizadores apoya NVIDIA para la investigación?

Además de Muon, NVIDIA también apoya muchos otros optimizadores para que la comunidad de investigación los explore, incluyendo:

Reproduciendo resultados de entrenamiento de Muon

Esta sección explica cómo reproducir los resultados de rendimiento de la Tabla 1 y cómo Muon está integrado en la pila de software de NVIDIA.

Inicio rápido

Siga las instrucciones en las recetas de rendimiento en el repositorio de GitHub NVIDIA-NeMo/Megatron-Bridge.

Como ejemplo, para ejecutar Kimi en 256 GPUs NVIDIA GB300, comience desde la raíz del repositorio de Megatron Bridge:

 C python scripts/performance/setup_experiment.py --account  -i ${CONTAINER} --partition  -m kimi -mr kimi_k2 --log_dir  --num_gpus 256 --gpus_per_node 4 -t "00:15:00" -g gb300 -c fp8_mx -hf  

Integraciones

Muon está integrado en Megatron Core, incluyendo una clase Muon y un envoltorio distribuido por capas.

Tenga en cuenta que para Muon, use_distributed_optimizer lo enviará automáticamente al optimizador distribuido por capas en lugar del correspondiente por elementos.

Comience con optimizadores emergentes para el entrenamiento de LLM

Los optimizadores de orden superior como Muon son esenciales para ampliar los límites de la eficiencia en el entrenamiento de LLM. A través de la optimización distribuida por capas, iteraciones de Newton-Schulz distribuidas a medida y el trabajo continuo en el ocultamiento de comunicación y núcleos fusionados de SYRK/reducción, estos poderosos métodos ahora pueden desplegarse efectivamente a gran escala. Con Megatron Core, puedes comenzar a usarlos ahora.

Cuando elija una configuración, tenga en cuenta las siguientes consideraciones:

• Modo NS distribuido: Use el modo duplicado cuando la latencia de red sea predominante. Use el modo distribuido cuando el cálculo sea el cuello de botella. El modo por bloques evita la comunicación por completo pero aproxima la ortogonalización completa.
• Escala: El optimizador distribuido por capas permite Muon a escalas DP grandes con una sobrecarga mínima. La Tabla 1 muestra un rendimiento casi equivalente al de AdamW en NVIDIA GB300.
• Mejoras futuras: El ocultamiento de comunicación, un mejor equilibrio de carga y núcleos de SYRK/reducción fusionados cerrarán cualquier brecha de rendimiento restante en lanzamientos futuros.

¿Listo para comenzar? Consulte las recetas de rendimiento de Megatron Bridge.