Ai2 abre su stack de entrenamiento MoE, probado con más de un billón de parámetros

El Allen Institute for AI (Ai2) lanzó el jueves Olmo-core 3, una nueva versión del framework open source que usa para entrenar sus modelos de lenguaje Olmo. El cambio principal es un sistema de entrenamiento rehecho para modelos mixture-of-experts (MoE), la arquitectura dispersa en la que cada token pasa solo por un puñado de subredes especializadas. Ai2 dice que probó el stack con más de un billón de parámetros totales y que la próxima generación de Olmo se va a construir sobre él.

Esto es infraestructura, no un modelo. Hoy no hay nada nuevo para llamar desde una API. Lo que hace es abrir al público el tipo de plomería de entrenamiento que los laboratorios de frontera suelen guardarse.

Qué cambió

Los modelos MoE prometen más capacidad por aproximadamente el mismo cómputo por token, pero a medida que crecen la logística se come esa ventaja: cada experto igual tiene que vivir en algún lado de la memoria de las GPU, y mandar cada token al experto correcto a través de un clúster cuesta comunicación. Olmo-core 3 es el intento de Ai2 de mantener ese costo bajo control.

El giro de arquitectura más importante es dejar el paralelismo de datos totalmente fragmentado (FSDP), que en la implementación MoE anterior de Ai2 reunía y volvía a fragmentar los pesos en cada micro-batch, y pasar a un sistema basado en paralelismo de datos distribuido (DDP) que deja a los expertos fijos en sus GPU y mueve los datos hacia ellos. Encima de eso, el stack combina tres formas de repartir el trabajo:

• Paralelismo de expertos: cada GPU guarda solo una parte del conjunto de expertos.
• Paralelismo de pipeline: las capas del modelo se reparten entre grupos de GPU.
• Un optimizador distribuido: el estado del optimizador se reparte entre GPU en vez de copiarse en cada una.

Ai2 también enumera optimizaciones de más bajo nivel: ubicar los tokens ruteados directamente en los buffers de entrada de cada experto, mantener los metadatos de ruteo en la GPU para que la CPU no se frene esperándolos, kernels de grouped GEMM que agrupan muchos cálculos chicos de expertos, y soporte para el formato numérico de menor precisión MXFP8.

Los números, y qué muestran y qué no

Todas las cifras que siguen salen de benchmarks propios de Ai2; todavía no hay reproducción independiente.

• En ocho GPU NVIDIA B300, un MoE de 47.000 millones de parámetros procesó unos 52.000 tokens por segundo por GPU con el stack nuevo, contra 19.400 con la implementación anterior basada en FSDP: alrededor de 2,7 veces más. Ai2 lo describe como una prueba preliminar.
• Pasar de 8 a 128 expertos, ruteando igual cada token a cuatro, mantuvo los parámetros activos en torno a 3.200 millones, llevó los parámetros totales de 4.600 a 47.000 millones y bajó el throughput de entrenamiento menos de 5%.
• Activar MXFP8 donde más ayudaba dio cerca de 21% más de throughput que BF16 en cuatro B300, con la memoria activa máxima bajando de 103 GiB a 95 GiB.
• Una configuración de 1,2 billones de parámetros, con 58.360 millones activos, corrió en 512 GPU y llegó a un pico de 858 TFLOP/s por GPU. Otra prueba corta con DeepEP v2 alcanzó una configuración de 2,38 billones de parámetros.

Las cifras de billones hay que leerlas con cuidado. Ai2 aclara que esas pruebas usaron ruteo aleatorio para medir el rendimiento del sistema, no la calidad de un modelo entrenado, y que la corrida de 2,38 billones fue un test corto de capacidad y no un entrenamiento sostenido. Muestran que la plomería aguanta un modelo de ese tamaño; no dicen nada sobre si un modelo entrenado ahí sería bueno. Ai2 también ubica su stack al lado de Megatron-Core de NVIDIA, al que llama una opción establecida para MoE grandes, pero no publica una comparación directa.

Lo más útil quizás sean los resultados negativos

El informe técnico que acompaña el lanzamiento documenta varios hallazgos que van contra supuestos comunes, justo el tipo de detalle que un lanzamiento abierto puede compartir:

• Un puntaje de balanceo pensado para fomentar un ruteo parejo podía mejorar mientras la carga real se volvía *menos* balanceada. Ai2 lo llama token gerrymandering.
• Bajar la tasa de aprendizaje de los expertos porque ven menos tokens no mejoró los resultados en la familia de modelos que probaron.
• Los kernels de GPU tardaban distinto según los valores procesados, aun con matrices de idénticas dimensiones, así que una comparación justa de rendimiento necesita entradas iguales, no solo formas iguales.
• Superponer comunicación y cómputo en streams de GPU separados a veces hizo más lento el entrenamiento de punta a punta.

Qué significa para quienes desarrollan

Si hoy estás eligiendo entre modelos, esta semana no cambia nada. La relevancia es de más largo plazo. Ai2 dice que su próximo Olmo va a ser un MoE, entrenado con su dataset más grande y con su ventana de contexto más larga hasta ahora, y que apunta a que sea su Olmo más capaz. No dio fecha. Si ese modelo sale con pesos, datos y código de entrenamiento abiertos, les daría a los desarrolladores un modelo disperso de estilo frontera cuyo pipeline completo se puede inspeccionar y reproducir, algo que los laboratorios cerrados no ofrecen.

Para equipos que entrenan o hacen fine-tuning de sus propios MoE, Olmo-core 3 ya es algo utilizable: Ai2 dice que el código está abierto en GitHub y que se puede adaptar a otro hardware y a otros esquemas de ruteo. Si las cifras de throughput se sostienen en clústeres distintos del setup de B300 de Ai2 es la pregunta abierta que conviene testear antes de adoptarlo.

Fuente: el anuncio de Ai2.