🤖 machine learning

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

El artículo presenta a Asteria, un sistema de ejecución que habilita la optimización de segundo orden práctica y escalable para modelos de lenguaje grandes mediante la distribución dinámica de los estados del optimizador a través de jerarquías de memoria heterogéneas y la desacoplamiento de los costosos cálculos del precondicionador de la ruta crítica de entrenamiento en GPU para reducir la sobrecarga y acelerar la convergencia.

Autores originales: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Publicado 2026-05-18

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

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 intentando enseñar a un robot gigante y superinteligente (un Modelo de Lenguaje Grande) a escribir historias. Para ello, necesitas un "profesor" (un optimizador) que observe los errores del robot y le indique cómo mejorar.

Actualmente, la mayoría de los profesores utilizan un método simple llamado AdamW. Es como un estudiante que revisa sus tareas, ve algunas respuestas incorrectas y realiza correcciones pequeñas y rápidas. Es rápido y funciona bien, pero requiere muchísimas prácticas (millones de ejemplos) para volverse realmente bueno.

Existe un método de profesor más inteligente y avanzado llamado Optimización de Segundo Orden (como SOAP o Shampoo). Este profesor no solo observa qué estaba mal; analiza la forma de los errores para comprender el problema en profundidad. Aprende mucho más rápido y necesita menos ejemplos. Sin embargo, hay un gran inconveniente: este profesor inteligente es increíblemente pesado. Requiere cantidades masivas de memoria y potencia de cálculo para realizar sus complejos cálculos matemáticos, lo que a menudo provoca que la computadora se bloquee o funcione tan lentamente que el profesor simple termine el trabajo primero.

Presentamos a Asteria: El "Gerente de Logística Inteligente"

Los investigadores de Oxford construyeron un nuevo sistema llamado Asteria. Piensa en Asteria no como un nuevo profesor, sino como un brillante gerente de logística que reorganiza todo el aula para que el profesor inteligente finalmente pueda hacer su trabajo sin destruir la escuela.

Así es como Asteria resuelve los tres problemas más grandes, utilizando analogías simples:

1. El problema de "Los Muebles en una Habitación Pequeña" (Memoria)

El Problema: El profesor inteligente necesita mantener enormes y complejos gráficos (matrices) en su mente (la memoria de la GPU). En una computadora estándar, no hay suficiente espacio. Es como intentar meter una cama tamaño king, una mesa de comedor y un armario en un apartamento estudio. La habitación se llena y el profesor tiene que renunciar.
La Solución de Asteria: Asteria actúa como un experto en muebles plegables. Se da cuenta de que el profesor no necesita sostener todo en sus manos al mismo tiempo.

Mantiene las partes más activas de los gráficos en la GPU (el escritorio).
Mueve las partes pesadas y de uso menos frecuente a la CPU (el pasillo) o incluso a un disco duro (el garaje).
Crucialmente, solo trae las partes pesadas de vuelta al escritorio exactamente cuando se necesitan. Esto permite que el profesor inteligente trabaje en una laptop pequeña tan bien como en un superordenador.

2. El problema de "El Embotellamiento" (Velocidad)

El Problema: Cada pocos pasos, el profesor inteligente debe realizar un cálculo masivo y complejo (como resolver un rompecabezas gigante) para actualizar sus gráficos. Esto toma mucho tiempo y detiene a toda la clase para trabajar. Es como un camión de reparto que bloquea toda la autopista mientras descarga un solo paquete. La GPU (el cerebro de la computadora) se queda inactiva, esperando.
La Solución de Asteria: Asteria introduce una línea de montaje paralela.

En lugar de detener la clase principal para hacer las matemáticas pesadas, Asteria envía el "trabajo pesado" a un equipo de trabajadores en la oficina trasera (la CPU).
Mientras la clase principal (la GPU) sigue enseñando al robot, los trabajadores de la oficina trasera resuelven silenciosamente los rompecabezas complejos en segundo plano.
Para cuando la clase necesita los nuevos gráficos, la oficina trasera ya los ha terminado y los ha deslizado hacia allí. La clase principal nunca tiene que detenerse. El "embotellamiento" desaparece.

3. El problema de "El Chat de Grupo" (Entrenamiento Distribuido)

El Problema: Cuando se entrena con muchas computadoras a la vez, todos suelen tener que detenerse y ponerse de acuerdo en la misma información exacta antes de pasar al siguiente paso. Esto es como un chat de grupo donde todos tienen que esperar a que la persona más lenta responda antes de que cualquiera pueda escribir de nuevo.
La Solución de Asteria: Asteria utiliza una política de "Lo suficientemente bueno".

En lugar de esperar a que todos estén perfectamente sincronizados, permite que las computadoras trabajen con información ligeramente "desactualizada" (ligeramente anticuada) durante un breve período.
Solo envía actualizaciones cuando son verdaderamente necesarias.
Esto mantiene al grupo moviéndose rápido, como un equipo que confía en que los demás seguirán trabajando mientras se ponen al día con las actualizaciones más tarde, en lugar de detener todo el proyecto cada cinco minutos.

Los Resultados: ¿Qué Descubrieron?

Los investigadores probaron Asteria en hardware real, incluyendo una computadora potente individual (Nvidia DGX Spark) y un gran clúster de computadoras (Nvidia GH200).

Funciona en máquinas pequeñas: Pudieron entrenar un modelo de lenguaje grande (1 mil millones de parámetros) en una sola máquina que anteriormente no podía manejar los requisitos de memoria de este profesor inteligente.
Es más rápido en tiempo real: Al ocultar los "embotellamientos", el entrenamiento termina en menos tiempo de reloj real, incluso aunque las matemáticas sean más difíciles.
Ahorra energía: Al mantener el cerebro de la computadora (GPU) ocupado y no permitir que se quede inactiva esperando cálculos, Asteria en realidad utiliza menos energía total para lograr el mismo resultado en comparación con los antiguos y torpes métodos.
No pierde calidad: El modelo aprende tan bien como lo haría con el profesor inteligente "nativo" (sin optimizar), pero llega allí mucho más rápido y barato.

En Resumen:
Asteria no inventa una nueva fórmula matemática. En cambio, reevalúa cómo la computadora ejecuta esas matemáticas. Separa el trabajo pesado del trabajo principal, utiliza todo el espacio de almacenamiento disponible de manera inteligente y permite que las computadoras trabajen de forma asíncrona. Esto convierte un método "teóricamente excelente pero prácticamente imposible" en una herramienta práctica para entrenar la próxima generación de IA.

Resumen Técnico: Asteria – Optimización de Segundo Orden Orquestada en Tiempo de Ejecución para el Entrenamiento Escalable de LLM

1. Declaración del Problema

Los métodos de optimización de segundo orden (por ejemplo, Shampoo, SOAP, KL-Shampoo) ofrecen una eficiencia estadística superior y una convergencia más rápida para el entrenamiento de Modelos de Lenguaje Grande (LLM) en comparación con los optimizadores de primer orden como AdamW. Sin embargo, su implementación práctica a gran escala se ve obstaculizada por barreras severas a nivel del sistema:

Restricciones de Memoria: Los métodos de segundo orden requieren almacenar grandes matrices de covarianza ( $L$ y $R$ ) y realizar descomposiciones espectrales densas, lo que genera una huella de memoria de $O(N^2)$ y un costo computacional de $O(N^3)$ . En entornos con restricciones de memoria (por ejemplo, estaciones de trabajo de nodo único con Arquitecturas de Memoria Unificada), estos estados agotan la memoria física, provocando fallos por falta de memoria (OOM) o forzando reducciones en el tamaño del modelo y el tamaño del lote.
Picos de Latencia: El cálculo periódico de raíces inversas (por ejemplo, $L^{-1/4}$ ) implica operaciones matriciales de costo cúbico. En las implementaciones existentes, estas actualizaciones síncronas bloquean la tubería de entrenamiento de la GPU, causando picos de latencia periódicos severos que interrumpen la superposición de granularidad fina entre el cálculo y la comunicación, esencial para un entrenamiento distribuido eficiente (por ejemplo, FSDP).
Sobrecarga de Sincronización: Los métodos distribuidos actuales de segundo orden dependen de actualizaciones globales de estado rígidas y síncronas (AllGather) a través de todas las ranuras. Esto acopla la comunicación rápida intra-nodo con la sincronización lenta inter-nodo, impidiendo que el sistema explote límites de obsolescencia asíncronos y conscientes de la topología.

Las soluciones existentes, como la fragmentación basada en DTensor, no abordan estos problemas en configuraciones de nodo único con recursos limitados o cuando la sobrecarga de comunicación de la sincronización global supera los beneficios.

2. Metodología: El Tiempo de Ejecución Asteria

Asteria es un sistema de tiempo de ejecución diseñado conjuntamente en hardware y software construido sobre PyTorch FSDP. Desacopla la lógica de optimización de segundo orden de la ruta crítica de entrenamiento de la GPU mediante tres mecanismos centrales:

A. Jerarquización de Memoria Asimétrica Adaptativa a la Arquitectura

A diferencia de los optimizadores de primer orden que tratan el estado como vectores homogéneos, Asteria reconoce el ciclo de vida asimétrico de los estados de segundo orden:

GPU: Almacena factor_matrices (factores de Kronecker) que se actualizan con frecuencia durante el entrenamiento.
Host (CPU/UVM): Almacena inv_factor_matrices (factores inversos). Estos se mantienen en memoria accesible por la CPU para evitar la presión sobre la memoria de la GPU.
NVMe (Opcional): Proporciona un nivel de preparación para ejecuciones con capacidad limitada, permitiendo que los estados se intercambien mediante computación serializada o paginación JIT.
Mecanismo: El tiempo de ejecución captura instantáneas del lado del host de las matrices de factores, despacha trabajadores asíncronos de CPU para calcular raíces inversas y escribe los resultados de vuelta en la memoria del host. Cuando es necesario, la GPU accede a estos mediante vistas de Memoria Virtual Unificada (UVM) o prefetching explícito, evitando copias redundantes de host a dispositivo.

B. Tubería de Estado Fantasma Orquestada por Ganchos

Para eliminar los picos de latencia, Asteria elimina el cálculo de raíces inversas de $O(N^3)$ de la ruta crítica de la GPU:

Flujo Fantasma: Un flujo CUDA de baja prioridad maneja la preparación de estados del lado de la GPU (prefetching y paginación) sin bloquear el flujo principal de entrenamiento.
Ejecución Asíncrona: Las costosas actualizaciones de raíces inversas se ejecutan de forma asíncrona en hilos de trabajadores de CPU.
Integración de Ganchos: Ganchos ligeros de PyTorch (register_forward_hook, register_full_backward_pre_hook) actúan como señales de programación. Desencadenan el vaciado de transferencias de estado listas y el prefetching de factores inversos que pronto serán consumidos, asegurando que la GPU tenga precondicionadores frescos sin detener el bucle principal de entrenamiento.

C. Coherencia de Obsolescencia Acotada Consciente de la Topología

Para el entrenamiento distribuido, Asteria reemplaza la sincronización global rígida con un protocolo de coherencia selectivo:

Seguimiento de Obsolescencia: Un CoherenceRegistry rastrea la frescura de cada bloque de precondicionador. Solo los bloques que exceden un presupuesto de obsolescencia configurado se sincronizan.
Sincronización Jerárquica: La sincronización ocurre en dos etapas: primero dentro de los grupos de procesos locales al nodo, y luego a través de representantes de nodos. Esto minimiza el tráfico entre nodos y evita viajes de ida y vuelta innecesarios de CPU a GPU manteniendo la comunicación en memoria accesible por el host.
Barrera Acotada: El tiempo de ejecución impone una barrera local que asegura que, si una versión de un precondicionador se vuelve demasiado obsoleta mientras un trabajo de CPU está en vuelo, el optimizador espera a su finalización, equilibrando la frescura con el rendimiento.

3. Contribuciones Clave

Desacoplamiento en Tiempo de Ejecución: Asteria demuestra que la optimización de segundo orden puede volverse práctica no simplificando el algoritmo, sino replanteando la gestión de estados, la computación en segundo plano y la sincronización a nivel de tiempo de ejecución.
Escalabilidad de Memoria: Permite el entrenamiento de segundo orden en hardware de nodo único con restricciones de memoria (por ejemplo, una sola GPU GB10 con 128 GB de memoria unificada) donde la fragmentación tradicional es imposible.
Ocultamiento de Latencia: Al mover los cálculos de raíces inversas al host y utilizar una tubería fantasma, Asteria aplanar los picos de latencia periódicos, logrando un perfil de ejecución de hardware suave comparable a los métodos de primer orden.
Eficiencia Energética: El sistema mejora la compensación entre pérdida y energía, logrando una mayor reducción de la pérdida por unidad de energía al reducir el tiempo de inactividad de la GPU y optimizar la superposición de la computación del lado del host.

4. Resultados Experimentales

Los autores evaluaron Asteria en modelos de la familia OLMo (660M, 1B, 7B parámetros) utilizando el marco de entrenamiento OLMo en dos plataformas:

Nodo Único (DGX Spark, GPU GB10):
- Rendimiento: Asteria-SOAP y Asteria-KL-Shampoo eliminaron los picos periódicos de tiempo por paso observados en implementaciones nativas (reduciendo el tiempo máximo por paso de ~96s a ~66.5s).
- Energía: Mientras que los métodos de segundo orden nativos consumieron aproximadamente un 17-20% más de energía total del SoC que AdamW, Asteria redujo esta brecha significativamente (por ejemplo, Asteria-KL-Shampoo consumió solo ~7.9% más de energía que AdamW).
- Eficiencia: Asteria-KL-Shampoo logró la mayor eficiencia normalizada de reducción de pérdida (3.3010), superando tanto a los métodos de segundo orden nativos como a AdamW.
Multi-Nodo (Cluster GH200):
- Convergencia: Las variantes de Asteria mantuvieron trayectorias de convergencia cercanas a sus contrapartes nativas de segundo orden, superando significativamente a AdamW en tiempo de reloj para alcanzar niveles de pérdida objetivo.
- Tolerancia a la Obsolescencia: Los experimentos que variaron el presupuesto de obsolescencia ( $S$ ) mostraron que aumentar $S$ de 1 a 5 redujo sustancialmente el tiempo de entrenamiento con un impacto negligible en la calidad final del modelo.
- Escalado Fuerte: En modelos de 7B escalando de 2 a 16 nodos, Asteria logró una mayor aceleración y un menor tiempo por paso en comparación con las líneas base nativas, demostrando ser efectivo incluso a medida que aumenta el costo relativo del mantenimiento del estado.

5. Significado y Afirmaciones

El artículo argumenta que el principal obstáculo para el entrenamiento práctico de LLM de segundo orden no es el costo matemático de las actualizaciones conscientes de la curvatura, sino el desajuste de sistemas entre el mantenimiento de estados de segundo orden y las pilas de entrenamiento modernas.

Reencuadre del Problema: Los autores afirman que la optimización de segundo orden debe considerarse como un problema de diseño conjunto que involucra algoritmos, programación de tiempo de ejecución y sistemas de memoria heterogénea.
Viabilidad: Demuestran que la memoria de GPU agregada grande no es un prerrequisito fundamental para el entrenamiento de segundo orden; con una colocación consciente del ciclo de vida y acceso escalonado, estos métodos pueden ejecutarse con presupuestos de memoria ajustados.
Practicidad: Al tratar la frescura del precondicionador y la comunicación como recursos gestionados en tiempo de ejecución en lugar de estados globales rígidos, Asteria reduce la brecha entre la promesa estadística de la optimización de segundo orden y su capacidad de implementación en el mundo real.

El artículo concluye que Asteria hace de la optimización de segundo orden una opción de sistemas viable para el entrenamiento de LLM, ofreciendo mejor convergencia y eficiencia energética sin requerir simplificación algorítmica.