Allen Institute publica Olmo-core 3, su infraestructura abierta para entrenar modelos gigantes de IA
El nuevo sistema permite entrenar modelos de mezcla de expertos de más de un billón de parámetros sin perder eficiencia computacional, democratizando el desarrollo de IA avanzada.
Olmo-core 3 es una infraestructura de código abierto del Allen Institute for AI para entrenar modelos de mezcla de expertos (arquitecturas de IA con componentes especializados que solo activan una parte por entrada) de más de un billón de parámetros. Usa paralelismo de datos distribuido en lugar de fragmentado, manteniendo expertos residentes en GPU y enrutando datos a ellos, lo que multiplica por 2,7 el rendimiento frente a versiones anteriores.
El Allen Institute for AI ha publicado Olmo-core 3, una actualización importante de su infraestructura de código abierto para desarrollar modelos de lenguaje. La novedad principal es un sistema rediseñado para entrenar modelos de mezcla de expertos o MoE (arquitecturas de IA que contienen muchos componentes especializados, pero solo activan una parte de ellos para cada entrada, ahorrando recursos). Según el instituto, este sistema está pensado para escalar hasta modelos de más de un billón de parámetros (los valores numéricos ajustables que determinan el comportamiento del modelo) sin perder eficiencia.
Entrenar modelos de IA grandes consume enormes cantidades de cómputo y energía, lo que aleja esta tecnología de universidades y laboratorios pequeños. Los modelos MoE son más eficientes porque no usan todos sus parámetros para cada cálculo, pero siguen necesitando almacenar el modelo completo en la memoria de las GPU (los procesadores especializados que aceleran el entrenamiento de IA) y coordinar qué datos van a qué experto a través de múltiples máquinas, lo que genera sus propios costes. En una prueba de referencia, Olmo-core 3 aumentó el número de expertos de 8 a 128 mientras seguía seleccionando solo cuatro por token (las unidades mínimas de texto que procesa un modelo de lenguaje), manteniendo los parámetros activos por token en torno a 3.200 millones. La capacidad total de parámetros creció de 4.600 millones a 47.000 millones, pero la velocidad de entrenamiento solo cayó un 5%.
La clave del avance está en cómo distribuye el trabajo. Las versiones anteriores de Olmo-core usaban paralelismo de datos totalmente fragmentado o FSDP (una técnica que reagrupaba los pesos del modelo para cada pequeño lote de datos de entrenamiento), pero Olmo-core 3 cambia a paralelismo de datos distribuido o DDP (que mantiene a los expertos residentes en las GPU y les envía los datos relevantes, evitando ese reagrupamiento continuo). En pruebas preliminares con ocho GPU NVIDIA B300, un modelo MoE de 47.000 millones de parámetros procesó 52.000 tokens por segundo y por GPU con el nuevo sistema, frente a 19.400 con la versión anterior, 2,7 veces más rendimiento.
El sistema combina varias técnicas de distribución: paralelismo de expertos (reparte los expertos entre GPU para que cada una almacene solo parte del conjunto), paralelismo de canalización (divide las capas sucesivas del modelo entre grupos de GPU) y un optimizador distribuido (reparte el estado del optimizador, los datos adicionales para calcular actualizaciones durante el entrenamiento, en lugar de guardar una copia completa en cada GPU). También reduce costes de enrutamiento con técnicas como paralelismo de expertos por filas, que coloca los datos directamente en los búferes de entrada de cada experto, y enrutamiento residente en GPU, que mantiene los metadatos de enrutamiento en las GPU para que la CPU pueda encolar trabajo sin esperar copias. Además, soporta MXFP8, un formato numérico de menor precisión que representa algunos valores con menos bits, reduciendo cómputo y transferencia de datos entre GPU.
El Allen Institute ha probado Olmo-core 3 con configuraciones de hasta 1,2 billones de parámetros en 512 GPU NVIDIA B300, alcanzando 858 TFLOP/s/GPU (operaciones de coma flotante por segundo y por GPU, una medida de potencia de cálculo útil). También experimentó con DeepEP v2, un método alternativo de comunicación entre expertos, llegando a 2,38 billones de parámetros en una prueba corta de capacidad. El informe técnico documenta además lecciones aprendidas: por ejemplo, una métrica para balancear el enrutamiento entre expertos podía mejorar mientras la carga real se desequilibraba (un fenómeno que llaman gerrymandering de tokens), y que superponer comunicación y cómputo en flujos separados de GPU no siempre aceleraba el entrenamiento. Olmo-core 3 es la base para la próxima generación de Olmo, que usará arquitectura MoE y apunta a ser el modelo más capaz del instituto hasta la fecha.
