🤖 machine learning

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

FORGE es un método de entrenamiento de LLM eficiente en memoria que fusiona el paso del optimizador en la pasada hacia atrás para eliminar los gradientes materializados al procesarlos enteramente en registros, reduciendo así a la mitad el uso de memoria del optimizador, acelerando el entrenamiento y preservando la fidelidad de precisión completa sin alterar la lógica de actualización subyacente.

Autores originales: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Publicado 2026-06-23

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

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

El gran problema: El "taller desordenado"

Imagina que estás entrenando un modelo de IA gigante (como un cerebro robótico) en una computadora. Para enseñar a este robot, la computadora tiene que realizar una cantidad masiva de matemáticas.

En la forma estándar de hacer esto (llamada "diferenciación de modo inverso"), la computadora sigue un proceso estricto de dos pasos:

Calcular los errores: Observa los datos, determina en qué se equivocó el robot y escribe una lista gigante de "notas de corrección" (gradientes) para cada una de las partes del cerebro del robot. Escribe estas notas en una pizarra gigante (la memoria de la computadora).
Aplicar las correcciones: Solo después de que la pizarra completa está llena, la computadora toma un borrador y un marcador para realmente actualizar el cerebro del robot basándose en esas notas.

El cuello de botella: El problema es que la computadora tiene que mantener esa gigante pizarra de notas en su memoria mientras también sostiene el cerebro del robot y las notas para el siguiente paso. Esta "pizarra" ocupa tanto espacio que a menudo llena la memoria de la computadora antes de que el entrenamiento pueda siquiera comenzar. Es como intentar reparar un coche en un garaje, pero tener que mantener el plano completo del coche, las herramientas y el manual de reparación extendidos en el suelo al mismo tiempo. Si el garaje es demasiado pequeño, no puedes meter el coche.

La solución: FORGE (El método de "corrección instantánea")

Los autores de este artículo, FORGE (Fused On-Register Gradient Elimination), dicen: "¿Para qué escribir las notas en la pizarra?"

Ellos argumentan que la lista gigante de notas es solo un artefacto de cómo hacemos las matemáticas, no algo que el proceso de aprendizaje realmente necesite.

Cómo funciona FORGE:
En lugar de escribir las notas y luego volver a leerlas, FORGE realiza las matemáticas y la corrección al mismo tiempo, en la zona de almacenamiento más rápida y pequeña de la computadora (llamada "registros").

La analogía: Imagina a un maestro chef cocinando un plato complejo.
- Forma antigua: El chef pica una cebolla, escribe "cebolla picada" en una libreta, pone la cebolla en un bol y luego pasa al siguiente ingrediente. Una vez que todos los ingredientes están picados y listos en la lista, el chef lee la libreta y los mezcla todos. La libreta ocupa espacio en la encimera.
- Forma de FORGE: El chef pica una cebolla y la lanza inmediatamente a la olla. Luego pica una zanahoria y la lanza de inmediato. Nunca escribe nada. Nunca necesita una libreta. La encimera se mantiene despejada.

Por qué esto es importante

El artículo afirma tres beneficios principales de este enfoque de "sin notas":

1. Ahorra un espacio masivo (Memoria)
Debido a que la computadora nunca escribe la lista gigante de notas en la memoria principal, libera una enorme cantidad de espacio.

El resultado: En un chip de computadora estándar (H200), pudieron entrenar un modelo de 8 mil millones de parámetros utilizando un 53% menos de memoria.
La analogía: Es como poder meter una mesa de comedor para 10 personas en un estudio pequeño porque dejaste de necesitar guardar las sillas adicionales en el pasillo.

2. Es realmente más rápido
Debido a que la computadora no tiene que detenerse para escribir notas, esperar y luego leerlas de nuevo, todo el proceso se acelera.

El resultado: Los pasos de entrenamiento corren aproximadamente 1.5 veces más rápido.
La analogía: Es la diferencia entre un repartidor que deja un paquete, corre de vuelta al camión para buscar el siguiente y repite el proceso (Forma antigua), frente a un repartidor que tiene una cinta transportadora que carga el paquete directamente al camión mientras lo agarra (FORGE).

3. No pierde precisión
Normalmente, cuando intentas ahorrar espacio saltándote pasos, pierdes precisión (el robot aprende un poco peor). Los autores demuestran que, debido a que están realizando las matemáticas en los "registros" de mayor calidad de la computadora (precisión total) antes de descartar los datos, el aprendizaje es matemáticamente idéntico al método antiguo y lento.

El resultado: El robot aprende igual de bien, pero de forma mucho más eficiente.

El truco de los "Tiles" (Azulejos)

¿Cómo hacen esto sin que sea un caos? Dividen el enorme problema matemático en trozos pequeños y manejables llamados "tiles" (como cuadrados de 128x128).

Procesan un tile: Calculan el error, corrigen el cerebro y descartan el error inmediatamente.
Luego pasan al siguiente tile.
Debido a que hacen esto tile por tile, la computadora nunca tiene que retener la lista completa de errores al mismo tiempo.

Lo que esto nos permite hacer

El artículo muestra que con FORGE:

Puedes entrenar modelos más grandes en la misma computadora.
Puedes usar "lotes" (batches) más grandes (enseñarle al robot más ejemplos a la vez) sin quedarte sin memoria.
En una prueba específica, pudieron entrenar un modelo en cuatro GPUs que habría requerido ocho GPUs usando el método antiguo.

Resumen

FORGE es una nueva forma de entrenar IA que evita que la computadora llene su memoria con "notas de corrección" temporales. Al calcular los errores y corregir el modelo instantáneamente en la parte más rápida del chip, reduce el uso de memoria a la mitad y acelera el entrenamiento, todo sin hacer que la IA sea menos inteligente. Convierte un taller congestionado y lento en una línea de ensamblaje ágil y de alta velocidad.

Resumen Técnico: FORGE (Eliminación de Gradientes Fusionada en Registros)

1. El Problema: El Suelo de Gradiente (The Gradient Floor)

Los métodos actuales de entrenamiento de Grandes Modelos de Lenguaje (LLM) eficientes en memoria se centran principalmente en tres de los cuatro depósitos de memoria en una GPU: estados del optimizador (mediante compresión), activaciones (mediante recomputación) y pesos del modelo. Sin embargo, dejan sin tocar el cuarto depósito: el gradiente de peso materializado.

En la diferenciación de modo inverso estándar, la pasada hacia atrás (backward pass) calcula el tensor de gradiente completo ( $\nabla W$ ) para cada capa, lo escribe en la Memoria de Alto Ancho de Banda (HBM) y solo entonces el optimizador lo lee de nuevo para actualizar los pesos. Este esquema de dos fases crea un "suelo de gradiente" donde, en el límite entre la pasada hacia atrás y el paso del optimizador, los gradientes de todas las capas deben coexistir en memoria. Para un modelo de 8B de parámetros, este depósito de gradiente transitorio por sí solo consume aproximadamente 16 GiB, lo que a menudo dicta si una configuración de entrenamiento cabe en una sola GPU. Los métodos existentes que fusionan los pasos de backward y del optimizador (por ejemplo, LOMO, FlashOptim) operan a la granularidad de parámetros o capas completas, lo que significa que el tensor de gradiente completo para un parámetro todavía se materializa en memoria antes de que ocurra la actualización.

2. Metodología: FORGE

FORGE (Fused On-Register Gradient Elimination) elimina por completo el tensor de gradiente materializado fusionando el paso del optimizador directamente en el cálculo del gradiente de peso a una granularidad de tesela (específicamente teselas de $128 \times 128$ ).

Mecanismo Central

En lugar de calcular la matriz de gradiente completa y almacenarla, FORGE procesa la matriz de pesos tesela por tesela:

Acumulación de Gradientes por Tesela: Para una tesela de peso específica, el kernel acumula los productos parciales del gradiente ascendente ( $\Delta Y$ ) y la activación de entrada ( $X$ ) directamente en los registros FP32 de la GPU.
Optimización en Registro: Una vez que el gradiente de la tesela se ha acumulado completamente en los registros, se aplica inmediatamente la actualización del optimizador (por ejemplo, AdamW) a esa tesela, también dentro de los registros.
Descarte Inmediato: La tesela de gradiente es consumida en el instante en que se produce. Nunca se escribe en la HBM como un tensor.
Escritura de Retorno (Write-Back): Solo se escriben de vuelta en la memoria global la tesela de peso actualizada y las teselas de estado del optimizador (momentos) (opcionalmente cuantizadas).

Precisión y Fidelidad

Exactitud: Para cualquier optimizador elemento a elemento separable (donde la actualización de un peso depende solo de la coordenada de ese peso, por ejemplo, AdamW, SGD, Lion), la actualización por tesela es matemmente idéntica a un paso de gradiente completo estándar. La única desviación es el orden de la suma de punto flotante sobre los tokens, lo cual es inherente a cualquier implementación de GEMM y no constituye una aproximación de la lógica del optimizador.
Fidelidad de Baja Precisión: En regímenes de entrenamiento BF16 o INT8, FORGE preserva la fidelidad de precisión completa de dos maneras:
- El gradiente nunca se trunca a BF16 antes de la optimización; se calcula y se consume en registros FP32.
- Para evitar la "zona muerta" de BF16 (donde las actualizaciones pequeñas se redondean a cero), FORGE emplea redondeo estocástico para las escrituras de retorno de los pesos, asegurando una acumulación imparcial de las actualizaciones de sub-ULP.
- Para los estados del optimizador INT8, FORGE utiliza tamaños de bloque más pequeños (64 frente a 2048 en optimizadores de 8 bits estándar) para minimizar el error de cuantización.

Compatibilidad con Entrenamiento Distribuido

FORGE es demostrablemente exacto bajo Paralelismo de Tensores (TP) y Paralelismo de Secuencia (SP) porque el gradiente local para un fragmento (shard) está completo al momento del backward. Sin embargo, es incompatible con el Paralelismo de Datos (DP) o el Paralelismo de Contexto en su forma actual, ya que estos requieren un all-reduce de los gradientes parciales antes de que se pueda aplicar la actualización, lo que requeriría materializar el tensor de gradiente.

3. Contribuciones Clave

Eliminación Estructural, No Compresión: FORGE identifica el gradiente materializado como un artefacto del esquema de entrenamiento en lugar de una necesidad matemática. Elimina el término de memoria de gradiente $2P$ (conteo de parámetros), reduciendo el depósito de gradiente de $O(P)$ a $O(T_R T_C)$ (tamaño de la tesela), independientemente del tamaño del modelo.
Exactitud Demostrable: El artículo demuestra que, para optimizadores elemento a elemento, el esquema fusionado por teselas produce actualizaciones idénticas a AdamW estándar sobre gradientes FP32. Esta exactitud se mantiene bajo paralelismo de tensores y de secuencia.
Fidelidad de Precisión Completa en Baja Precisión: Al mantener los gradientes en registros FP32 hasta la optimización, FORGE evita el sesgo introducido al truncar los gradientes a BF16 antes de la optimización. Combinado con el redondeo estocástico, iguala la convergencia del maestro FP32 incluso cuando se utilizan estados BF16 o INT8.
Generalidad: El método es agnóstico a la arquitectura para capas lineales y soporta ocho optimizadores elemento a elemento totalmente fusionados (AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum).

4. Resultados Experimentales

Evaluado en GPUs NVIDIA H200 y B200 en modelos que van desde 0.6B hasta 14B de parámetros (Llama-3.1-8B, familia Qwen3):

Reducción de Memoria: FORGE reduce el pico de memoria entre un 38–53% comparado con AdamW fusionado. Para Llama-3.1-8B, el pico de memoria cae de ~75 GiB a 35.4 GiB (una reducción del 53%).
Rendimiento (Throughput): FORGE logra una aceleración de 1.52× en el tiempo de paso (99.1 ms frente a 150.9 ms) al eliminar el lanzamiento del kernel del optimizador por separado y el viaje de ida y vuelta (round-trip) de la HBM para los gradientes.
Escalabilidad:
- En una sola H200, FORGE entrena Qwen3-14B (que causa errores de Fuera de Memoria/OOM con AdamW fusionado estándar) a 63.2 GiB.
- En una configuración de 4 GPUs A100-40GB, FORGE permite un tamaño de micro-lote (micro-batch) de 4 (ajustándose a 24.6 GiB), mientras que el AdamW fusionado estándar solo permite un micro-lote de 1. Esto permite que FORGE logre un rendimiento 2.42× mayor que la línea base estándar en el mismo hardware.
Convergencia: FORGE sigue la convergencia de AdamW fusionado dentro de 0.001 nats (desviación promedio) en tareas de pre-entrenamiento continuo, confirmando que las desviaciones de baja precisión están dentro del ruido de ejecución.
Componibilidad: FORGE se compone limpiamente con la cuantización FP8 y el checkpointing de activaciones, absorbiendo los costos de descuantización que típicamente penalizan a los optimizadores estándar.

5. Significado y Reivindicaciones

El artículo afirma que FORGE representa un cambio estructural en lugar de una optimización incremental. Al eliminar el suelo de gradiente, convierte el exceso de memoria del optimizador en capacidad útil para lotes más grandes, secuencias más largas o modelos más grandes en el hardware existente.

Límite de Capacidad: Los autores afirman que la eliminación del suelo de gradiente no es solo una mejora de un factor constante, sino un límite de capacidad. Por ejemplo, las configuraciones de entrenamiento estándar para modelos de 14B en H200 individuales son imposibles debido al suelo de gradiente, mientras que FORGE las hace factibles.
Ortogonalidad: El método es ortogonal a otras técnicas de ahorro de memoria (como la cuantificación de estados o la recomputación de activaciones) y puede combinarse con ellas para reducir aún más la huella de memoria.
Limitaciones: El método está actualmente limitado a optimizadores separables elemento a elemento y no soporta paralelismo de datos sin modificaciones. También requiere que el gradiente de entrada sea leído antes de la actualización de los pesos, lo que excluye embeddings con lazos (tied embeddings) o expertos de MoE con múltiples consumidores de gradiente en un solo paso.

En resumen, FORGE demuestra que el "suelo de gradiente" es un artefacto innecesario de los bucles de entrenamiento actuales. Al fusionar el optimizador en el proceso de backward a nivel de tesela, logra ahorros de memoria y aceleraciones significativas manteniendo la exactitud matemática para una amplia clase de optimizadores.