Scalable Training of Mixture-of-Experts Models with Megatron Core

Este informe presenta Megatron Core, un marco de código abierto que aborda los desafíos de escalabilidad en el entrenamiento de modelos de expertos mezclados (MoE) mediante optimizaciones integradas en memoria, comunicación y computación, logrando un alto rendimiento en hardware NVIDIA GB300/GB200 para modelos desde miles de millones hasta billones de parámetros.

Lo esencial

¡Claro que sí! Imagina que este informe es el manual de instrucciones definitivo para construir y entrenar a los "gigantes" de la inteligencia artificial del futuro.

Aquí tienes la explicación de cómo NVIDIA está entrenando estos modelos masivos (llamados MoE o "Mezcla de Expertos") usando su tecnología Megatron Core, explicada como si estuviéramos hablando en una cafetería.

---

🧠 ¿Qué es un modelo "Mezcla de Expertos" (MoE)?

Imagina que tienes que resolver un problema muy difícil, como escribir un libro de historia o diagnosticar una enfermedad rara.

• El modelo antiguo (Dense): Era como tener un solo genio en una habitación que intentaba saberlo todo. Para resolver cualquier cosa, ese genio tenía que pensar con toda su mente. Era lento y costoso.
• El modelo nuevo (MoE): Es como tener una biblioteca gigante de especialistas. Hay un "recepcionista" (el enrutador) que, cuando llega una pregunta, la envía solo a los 2 o 3 expertos que realmente saben de ese tema.
  • Ejemplo: Si preguntas sobre "fútbol", el recepcionista llama al experto en deportes, no al experto en cocina.
  • La ventaja: Puedes tener miles de expertos (mucha inteligencia total) sin tener que hacer trabajar a todos al mismo tiempo (ahorro de energía).

🚧 El Gran Problema: Las "Tres Paredes"

Aunque esta idea suena genial, entrenar a miles de expertos a la vez es una pesadilla logística. El informe dice que hay tres paredes que intentan detenernos:

1. La Pared de la Memoria (El "Cerebro" se llena)

Imagina que cada experto necesita su propia libreta de notas. Si tienes 100 expertos, necesitas 100 libretas. Pero como solo trabajan 3 a la vez, ¡la mayoría de las libretas están guardadas en el armario!

• El problema: Las tarjetas gráficas (GPUs) tienen un tamaño de armario limitado. Si intentas guardar las notas de todos los expertos, el armario explota.
• La solución de NVIDIA:
  • Comprimir las notas: En lugar de escribir con tinta gruesa (precisión alta), escriben con tinta muy fina (FP8 o FP4), ahorrando espacio sin perder mucho detalle.
  • Reescribir en el momento: En lugar de guardar todas las notas intermedias, las borran y las vuelven a calcular si las necesitan después (como si un chef no guardara los ingredientes cortados, sino que los cortara de nuevo si hace falta).
  • Alquilar espacio extra: Si el armario se llena, mueven las libretas menos usadas a un trastero en la CPU (memoria del ordenador) y las traen solo cuando son necesarias.

2. La Pared de la Comunicación (El "Teléfono Roto")

Como los expertos están en diferentes habitaciones (diferentes tarjetas gráficas), el recepcionista tiene que correr para llevar las preguntas a cada uno y traer las respuestas.

• El problema: Si hay 100 expertos en 100 habitaciones, el recepcionista pasa más tiempo corriendo que trabajando. Si las habitaciones están en edificios diferentes (diferentes servidores), el tiempo de viaje es enorme.
• La solución de NVIDIA:
  • Carriles exclusivos: Han creado "autopistas" especiales (llamadas DeepEP y HybridEP) para que los mensajes viajen a la velocidad de la luz entre las tarjetas.
  • Hacer dos cosas a la vez: Mientras el recepcionista corre llevando una pregunta, el experto ya está trabajando en la anterior. Se superponen las tareas para que el tiempo de espera sea cero.

3. La Pared de la Eficiencia Computacional (El "Motor que se calienta")

Los expertos son muy pequeños y rápidos. A veces, el ordenador pasa más tiempo encendiendo y apagando el motor (lanzando instrucciones) que conduciendo.

• El problema: Hay tantos expertos pequeños que el cerebro del ordenador (la CPU) se agota intentando dar órdenes a cada uno por separado.
• La solución de NVIDIA:
  • Agrupar tareas: En lugar de pedirle a un experto que haga una tarea, le piden que haga 100 a la vez en un solo bloque.
  • Grabar un guion: Usan una técnica llamada CUDA Graphs que es como grabar un video de una coreografía perfecta. En lugar de que el director (CPU) grite "¡haz esto, luego aquello!" en cada repetición, simplemente le dice al equipo "¡Ejecuta el video!". Esto elimina el tiempo de espera.

🧩 El Truco Maestro: "El Doblaje de Paralelismo" (Parallel Folding)

Aquí viene la parte más ingeniosa. En un modelo grande, tienes dos tipos de capas:

1. Capas de Atención: Necesitan ver todo el contexto (como leer todo un libro).
2. Capas de Expertos: Necesitan dividir el trabajo en muchos trozos pequeños.

Antes, tenías que usar la misma estrategia para ambos, lo cual era como intentar usar el mismo mapa para conducir por una autopista y por un laberinto. No funcionaba bien.

La solución de NVIDIA: Han inventado el "Parallel Folding".
Imagina que tienes un equipo de construcción.

• Para las capas de atención, organizas a los albañiles en grandes grupos grandes que trabajan juntos (alta eficiencia en grandes bloques).
• Para las capas de expertos, divides al equipo en muchos grupos pequeños que trabajan en paralelo (alta eficiencia en tareas pequeñas).
• El truco: Permiten que estos dos grupos usen diferentes mapas y estrategias al mismo tiempo dentro del mismo edificio. Esto permite escalar modelos de billones de parámetros sin que el sistema se rompa.

🏁 Los Resultados: ¿Qué tan rápido es?

Gracias a todas estas trucos, NVIDIA ha logrado entrenar modelos gigantes (como DeepSeek-V3 con 685 mil millones de parámetros) a velocidades increíbles:

• En sus nuevas tarjetas GB300/GB200, logran una velocidad de procesamiento que es 3 veces más rápida que las tarjetas anteriores.
• Han logrado entrenar modelos que antes tardarían años en meses, o incluso semanas.

🚀 ¿Y para qué sirve todo esto?

Esto no es solo para hacer chatbots más inteligentes. Esto permite:

1. Modelos más baratos: Al ser más eficientes, cuestan menos dinero entrenarlos.
2. Contextos largos: Pueden leer libros enteros o ver horas de video y recordar todo.
3. Aprendizaje por Refuerzo: Ayuda a entrenar robots o agentes de IA que aprenden probando cosas (como jugar videojuegos o conducir coches) de forma mucho más rápida y estable.

En resumen

NVIDIA ha tomado un problema logístico imposible (entrenar a miles de expertos simultáneamente) y lo ha resuelto como un director de orquesta genial:

1. Ahorra espacio comprimiendo las partituras.
2. Mejora el tráfico creando autopistas para los músicos.
3. Agrupar las notas para que la orquesta toque en bloque en lugar de nota por nota.
4. Permite que cada sección (cuerdas vs. vientos) use su propio ritmo y partitura sin chocar.

El resultado es una máquina capaz de crear inteligencias artificiales que antes eran solo ciencia ficción.

Resumen técnico

Resumen Técnico: Entrenamiento Escalable de Modelos Mixture-of-Experts (MoE) con Megatron Core

Este documento presenta un análisis técnico detallado del informe de NVIDIA sobre Megatron-Core MoE, un framework de código abierto diseñado para abordar los desafíos sistémicos únicos del entrenamiento de modelos de gran escala basados en la arquitectura Mixture-of-Experts (MoE). El informe detalla cómo superar las limitaciones de memoria, comunicación y eficiencia computacional para entrenar modelos con billones de parámetros en clusters de miles de GPUs.

1. El Problema: Los "Tres Muros" del MoE

El entrenamiento de modelos MoE a gran escala introduce desafíos sistémicos que no existen en los modelos densos tradicionales. La naturaleza dispersa de MoE (donde solo un subconjunto de expertos se activa por token) crea una desincronización entre parámetros y cómputo: el número total de parámetros crece mucho más rápido que el cómputo por token. Esto genera tres barreras fundamentales ("Tres Muros") que limitan la escalabilidad:

1. Muro de Memoria: Aunque solo una fracción de los expertos se activa, todos los parámetros de los $E$ expertos, sus gradientes y estados del optimizador deben residir en la memoria de la GPU. En modelos como DeepSeek-V3 (685B parámetros totales, 37B activos), la presión de memoria excede por mucho la capacidad de una sola GPU (requiriendo ~200 GB/GPU sin optimización).
2. Muro de Comunicación: La paralelización de expertos (EP) requiere operaciones all-to-all para enrutar tokens a los expertos correspondientes en diferentes GPUs. A medida que aumenta el grado de EP, este tráfico puede saturar el ancho de banda, especialmente si cruza límites de nodos (inter-nodo), consumiendo hasta el 60% del tiempo de entrenamiento.
3. Muro de Eficiencia Computacional: Los expertos finos generan muchas operaciones GEMM (multiplicación de matrices) pequeñas que no aprovechan al máximo las unidades de cómputo de la GPU (Tensor Cores). Además, la sobrecarga del host (lanzamiento de kernels, sincronización) y el desequilibrio de carga entre expertos dejan a las GPUs inactivas.

Adicionalmente, existe una desincronización densa-dispersa: las capas de atención (densas) y las capas MoE (dispersas) tienen configuraciones óptimas de paralelismo conflictivas (ej. alta Tensor Parallelism para atención vs. alta Expert Parallelism para MoE), lo que obliga a los frameworks tradicionales a elegir configuraciones subóptimas.

2. Metodología y Arquitectura

Megatron-Core aborda estos desafíos mediante un diseño integrado que abarca desde la arquitectura de la capa MoE hasta la optimización a nivel de sistema.

2.1. Arquitectura de la Capa MoE

El framework descompone la capa MoE en cuatro etapas secuenciales:

1. Enrutamiento (Route): Un router asigna tokens a los $k$ mejores expertos.
2. Despacho (Dispatch): Comunicación all-to-all para mover tokens a las GPUs que alojan los expertos asignados.
3. Cómputo (Compute): Ejecución local de los expertos mediante kernels GEMM agrupados.
4. Combinación (Combine): Comunicación inversa para devolver los tokens procesados a sus GPUs originales.

2.2. Paralelismo Multidimensional y "Parallel Folding"

La contribución central para resolver la desincronización densa-dispersa es Parallel Folding.

• Desacoplamiento: Permite que las capas de atención y las capas MoE utilicen configuraciones de paralelismo independientes. Mientras la atención puede usar alta Tensor Parallelism (TP) y Context Parallelism (CP), las capas MoE pueden usar alta Expert Parallelism (EP) con TP=1.
• Flexibilidad: Rompe la restricción tradicional donde el grado de EP debe ser menor o igual al de Data Parallelism (EP ≤ DP). Ahora, el EP puede "doblarse" a través de grupos de TP y CP, permitiendo configuraciones optimizadas para topologías de hardware específicas (ej. NVL72 en GB200).

2.3. Optimizaciones para Romper los Tres Muros

A. Rompiendo el Muro de Memoria

• Permutación Eficiente de Memoria: Reorganiza algebraicamente los cálculos para eliminar tensores intermedios redundantes sin costo computacional.
• Recomputación de Granularidad Fina: En lugar de recomputar capas enteras, se recomputan solo operaciones costosas en memoria pero baratas en cómputo (ej. funciones de activación, LayerNorm).
• Offloading de Activaciones: Transfiere activaciones a la memoria del CPU de forma asíncrona, solapando la transferencia con el cómputo.
• Precisión Reducida (FP8/FP4): Almacena activaciones y estados del optimizador en FP8 o FP4, reduciendo drásticamente el uso de memoria.
• FSDP para MoE: Implementa un diseño de Dual DeviceMesh que comparte parámetros y estados del optimizador de expertos dentro de grupos de paralelismo de expertos (EDP), no globalmente.

B. Rompiendo el Muro de Comunicación

• Despachadores Optimizados (DeepEP y HybridEP): Kernels personalizados que eliminan la etapa de permutación redundante y utilizan primitivas de hardware (como TMA e IBGDA) para maximizar el ancho de banda, especialmente en topologías NVLink.
• Solapamiento de Comunicación y Cómputo: Utiliza un esquema 1F1B (Forward-Backward) fusionado y dividido (W/D split) para ocultar la latencia del all-to-all detrás del cómputo de expertos de micro-lotes adyacentes.

C. Rompiendo el Muro de Eficiencia Computacional

• GEMM Agrupado (Grouped GEMM): Agrupa múltiples cálculos de expertos pequeños en un solo kernel para mejorar la utilización de los Tensor Cores.
• Fusión de Kernels: Combina operaciones de enrutamiento, permutación y cálculo de pérdidas auxiliares en kernels únicos para reducir la sobrecarga del host.
• CUDA Graphs: Captura secuencias de kernels para eliminar la latencia de lanzamiento por iteración. Para MoE sin caída de tokens (dropless), se introducen técnicas de ejecución sin sincronización (Sync-Free):
  • Kernels iniciados por dispositivo: Los kernels leen la información de forma (número de tokens) directamente desde la GPU, eliminando la sincronización CPU-GPU.
  • ECHO (Elastic Cloning): Clona dinámicamente expertos populares en slots inactivos para equilibrar la carga y reducir la fragmentación de memoria.
  • Paged Stashing: Gestiona la memoria de manera dinámica dentro del gráfico CUDA para evitar la asignación de buffers en el peor caso.

2.4. Entrenamiento de Precisión Reducida (FP8/FP4)

Se implementa una estrategia de precisión selectiva:

• Se mantiene el router, las capas de embeddings y los estados del optimizador en precisión alta (FP32/BF16) para garantizar la estabilidad numérica.
• Se cuantiza agresivamente los GEMM de expertos y las activaciones a FP8 (MXFP8 en Blackwell, Blockwise en Hopper) o FP4 (NVFP4), logrando beneficios simultáneos en memoria, comunicación y velocidad de cómputo.

2.5. Entrenamiento de Contexto Largo y RL

• Context Parallelism (CP) y Empaquetado: Para secuencias largas (16K-64K+), se combina CP con Tensor Parallelism y se utiliza el formato THD (Total tokens x Heads x Dimension) para empaquetar secuencias de longitud variable sin relleno (padding), reduciendo el desperdicio de cómputo.
• Contexto Dinámico (Dynamic-CP): Ajusta adaptativamente el grado de CP por micro-lote según la longitud de la secuencia para equilibrar la carga computacional y de memoria.
• Soporte para RL: Incluye características como Router Replay (reproducir decisiones de enrutamiento de inferencia en entrenamiento) y optimizaciones para secuencias de longitud extremadamente variable en post-entrenamiento por refuerzo.

3. Contribuciones Clave

1. Parallel Folding: Un marco de paralelismo multidimensional que desacopla la configuración de capas densas y dispersas, permitiendo optimizaciones independientes y rompiendo restricciones de escalabilidad anteriores.
2. Solución Integral de los Tres Muros: Un conjunto coherente de optimizaciones que abordan simultáneamente memoria, comunicación y cómputo, demostrando que la optimización de un muro a menudo habilita o mejora la solución de los otros.
3. Soporte Nativo para Precisión Reducida: Implementación completa de FP8 y FP4 (incluyendo NVFP4) con estrategias de cuantización específicas para MoE que mantienen la estabilidad del entrenamiento.
4. Tecnologías de Ejecución Sin Sincronización: Mecanismos como Sync-Free Kernels, ECHO y Paged Stashing que permiten el uso de CUDA Graphs en MoE dinámico (dropless), eliminando cuellos de botella del host.
5. Características de Producción: Soporte para Upcycling (conversión de modelos densos a MoE), Distributed Checkpointing (cambio de configuración de paralelismo sin conversión offline) y optimizaciones específicas para Reinforcement Learning.

4. Resultados de Rendimiento

El framework ha sido validado en hardware de última generación (NVIDIA GB300, GB200 y H100) con modelos de estado del arte:

• DeepSeek-V3 (685B parámetros):
  • GB300 (256 GPUs): 1,233 TFLOPS/GPU.
  • GB200 (256 GPUs): 1,048 TFLOPS/GPU.
  • H100 (1024 GPUs): 368 TFLOPS/GPU.
• Qwen3-235B:
  • GB300 (256 GPUs): 974 TFLOPS/GPU.
  • GB200 (256 GPUs): 919 TFLOPS/GPU.
  • H100 (256 GPUs): 320 TFLOPS/GPU.

Hallazgos de Rendimiento:

• Las plataformas GB200/GB300 logran aproximadamente 3x más rendimiento en tokens por segundo en comparación con H100, gracias a un mayor ancho de banda de memoria y soporte nativo para MXFP8/FP4.
• La optimización de la comunicación (DeepEP/HybridEP + solapamiento) reduce la sobrecarga de all-to-all de hasta un 60% a menos del 5-10% del tiempo de iteración.
• El uso de FP8/FP4 reduce la memoria de activaciones en un 50-75% y acelera los GEMM, permitiendo configuraciones de paralelismo más eficientes.

5. Significancia

Este trabajo es fundamental para la industria de la IA por varias razones:

• Viabilidad de Modelos Trillonarios: Hace posible el entrenamiento de modelos con billones de parámetros en hardware actual, superando las limitaciones de memoria que antes eran infranqueables.
• Eficiencia de Costos: Al maximizar la utilización de GPU (MFU) y reducir el tiempo de entrenamiento, disminuye significativamente el costo de entrenamiento de modelos de vanguardia.
• Flexibilidad de Hardware: Demuestra cómo adaptar algoritmos a topologías de hardware específicas (como NVL72 en GB200) para obtener el máximo rendimiento, estableciendo un nuevo estándar para la co-diseño de software y hardware.
• Adopción Industrial: Al ser una solución de código abierto y lista para producción, permite a la academia y la industria entrenar modelos MoE desde prototipos hasta modelos de producción masiva, facilitando la adopción de arquitecturas MoE más allá de los modelos densos tradicionales.

En resumen, Megatron-Core MoE representa un avance sistémico que transforma el entrenamiento de modelos MoE de un desafío de ingeniería prohibitivo en un proceso escalable y eficiente, habilitando la próxima generación de modelos de lenguaje grandes.