Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling

Resumen Técnico: Entrenamiento de LLM Eficiente en Memoria con Esparcidad Dinámica

Planteamiento del Problema
El Entrenamiento Esparcido Dinámico (DST, por sus siglas en inglés) ofrece un paradigma para mejorar la eficiencia de las redes neuronales profundas manteniendo una topología esparcida durante todo el entrenamiento. Sin embargo, aplicar el DST a los Modelos de Lenguaje de Gran Escala (LLMs) ha resultado ser un desafío debido a la inestabilidad de la optimización. Los autores identifican que la aplicación ingenua de los optimizadores estándar basados en Adam al DST resulta en severos "picos de pérdida" (loss spikes) inmediatamente después de las actualizaciones de la topología (poda y regeneración).

La causa raíz se identifica como un efecto de "arranque en frío" (cold-start): los parámetros recién regenerados carecen de estados de optimizador acumulados (estimaciones de los primeros y segundos momentos). En Adam, los términos de corrección de sesgo para estos nuevos parámetros resultan en magnitudes de actualización desproporcionadamente grandes en relación con los pesos maduros. Esto interrumpe la trayectoria de entrenamiento, causando picos de pérdida transitorios y dificultando la convergencia, particularmente a medida que las escalas del modelo aumentan o las actualizaciones de topología se vuelven más agresivas. Además, las implementaciones convencionales de DST a menudo no logran realizar ahorros de memoria porque asignan estados de optimizador de tamaño completo para todo el espacio de parámetros, a pesar de que solo un subconjunto de parámetros está activo.

Metodología: Entrenamiento Esparcido Eficiente en Memoria (SMET)
Para abordar estos desafíos, los autores proponen SMET, un marco diseñado para estabilizar el DST y reducir el consumo de memoria. El método consta de tres componentes principales:

1. Calentamiento del Estado del Optimizador (Optimizer-State Warm-up): Para mitigar el efecto de arranque en frío, SMET emplea una estrategia de calentamiento de dos etapas para los parámetros recién regenerados:

   • Calentamiento Implícito (Reinicio de Paso de Tiempo): El contador interno de pasos de tiempo ($t_i$) para los parámetros regenerados se reinicia. Esto reactiva el mecanismo de corrección de sesgo de Adam, inflando el denominador de la regla de actualización y constriñendo naturalmente la magnitud inicial de la actualización.
   • Calentamiento Explícito (Escalamiento de la Tasa de Aprendizaje): Se aplica un calentamiento lineal local donde la tasa de aprendizaje para los parámetros regenerados aumenta linealmente desde cero hasta la tasa objetivo durante $W$ pasos.
   • Estos mecanismos previenen conjuntamente las explosiones de actualización, suavizando la trayectoria de optimización inmediatamente después de los cambios de topología.

2. Escalamiento de la Tasa de Aprendizaje Dependiente de la Densidad: Para compensar las actualizaciones conservadoras introducidas por el calentamiento y para dar cuenta del reducido fan-in efectivo de las capas esparcidas, SMET escala la tasa de aprendizaje. Motivado por la Parametrización de Actualización Máxima ($\mu$P), la tasa de aprendizaje se escala por $1/\sqrt{d}$, donde $d$ es la densidad de esparcidad actual. Esto asegura que las magnitudes de actualización permanezcan comparables a los regímenes de entrenamiento denso.

3. Optimización Eficiente en Memoria: SMET implementa una estrategia de almacenamiento esparcido donde los gradientes y los estados del optimizador (estimaciones de momentos $m$ y $v$) se almacenan solo para los parámetros activos (no nulos).

   • Los gradientes se retienen exclusivamente en los índices activos durante la retropropagación.
   • Para gestionar la sobrecarga de seguimiento de los índices activos y los pasos de tiempo, se utiliza una estrategia de almacenamiento por bloques, agrupando los parámetros en bloques de tamaño $B$ para amortizar los costos de metadatos.
   • Este enfoque preserva la esparcidad durante todo el ciclo de vida del entrenamiento, a diferencia de los métodos que comprimen los estados pero producen modelos densos.

Contribuciones Clave

1. Análisis de la Inestabilidad: El artículo identifica y analiza teóricamente el fenómeno de los picos de pérdida en el pre-entrenamiento de LLMs con DST, atribuyéndolo al efecto de arranque en frío de los pesos recién regenerados en los optimizadores Adam.
2. Marco de Estabilización: Los autores proponen SMET, que integra el calentamiento del estado del optimizador y el escalamiento de la tasa de aprendizaje dependiente de la densidad para estabilizar las actualizaciones de topología y mejorar el rendimiento del entrenamiento.
3. Implementación Eficiente en Memoria: Se desarrolla una implementación de optimización esparcida que almacena gradientes y estados del optimizador solo para los parámetros activos, soportando patrones de esparcidad no estructurados y por bloques, manteniendo la esparcidad durante todo el entrenamiento.
4. Evaluación Sistemática: El trabajo proporciona una evaluación exhaustiva de DST a través de varios niveles de escala de modelo y niveles de esparcidad, demostrando que SMET supera a los métodos de DST existentes y se acerca al rendimiento del entrenamiento denso.

Resultados Experimentales
Los experimentos se realizaron en modelos basados en LLaMA (que van desde 60M hasta 1B de parámetros) pre-entrenados en los conjuntos de datos C4 y OpenWebText.

• Rendimiento: SMET supera consistentemente a otras líneas base de eficiencia de memoria (LoRA, ReLoRA, GaLore) y a los métodos de DST existentes (SET, RigL). En escalas de modelo más grandes (por ejemplo, 1B de parámetros), SMET logra una perplejidad cercana al entrenamiento denso, incluso con niveles altos de esparcidad (por ejemplo, densidad $d=0.1$).
• Estabilidad: El método elimina eficazmente los picos de pérdida observados en el DST vanilla, resultando en curvas de entrenamiento suaves.
• Eficiencia de Memoria: Al almacenar estados solo para los parámetros activos, SMET reduce significativamente la huella de memoria de los gradientes y los estados del optimizador. Para un modelo de 1B a una densidad $d=0.4$, se observan ahorros sustanciales de memoria.
• Escalabilidad: La brecha de rendimiento entre el entrenamiento esparcido y el denso se reduce a medida que aumenta el tamaño del modelo, lo que sugiere que el entrenamiento esparcido es particularmente viable para modelos de gran escala.
• Generalización: Los beneficios de SMET se generalizan a través de diferentes estrategias de regeneración (aleatoria y basada en gradientes) y diferentes arquitecturas de modelos (LLaMA y Qwen).

Significancia y Reivindicaciones
El artículo afirma que SMET allana el camino para que el entrenamiento esparcido sea una alternativa práctica al entrenamiento denso para los LLMs. Al resolver la inestabilidad de la optimización inherente al entrenamiento esparcido dinámico y lograr ahorros reales de memoria mediante el almacenamiento de estados esparcidos, SMET permite un pre-entrenamiento estable, escalable y eficiente en memoria. Los autores señalan que, si bien SMET mejora la estabilidad del entrenamiento y reduce la memoria del optimizador, aún quedan desafíos con respecto a la memoria de activación y la necesidad de soporte de hardware eficiente para la esparcidad de grano fino para lograr plenamente las aceleraciones en la inferencia. No obstante, el trabajo demuestra que, con una estabilización y gestión de memoria adecuadas, la esparcidad dinámica puede escalarse con éxito a los modelos fundacionales.