🤖 machine learning

Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention

Este artículo presenta Full-Stack FP4, un marco modular que extiende la cuantización NVFP4 más allá de las proyecciones lineales para incluir estados de optimizador estables, computaciones de Root/Muon y atención de precisión mixta, logrando un rendimiento de preentrenamiento cercano a BF16 con ganancias significativas de memoria y velocidad en una sola RTX 5090.

Autores originales: Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

Publicado 2026-08-11

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

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 cerebro digital gigante a hablar, escribir y pensar. Este "cerebro" es un Modelo de Lenguaje Extenso (LLM), y para aprender, tiene que procesar montañas de información. Pero hay un inconveniente: cuanto más aprende, más hambre tiene de memoria de computadora y electricidad. Es como intentar reproducir una película masiva de alta definición en una laptop pequeña y vieja; la pantalla se vuelve borrosa y el ventilador grita. Para resolver esto, los científicos han estado intentando encoger los números que la computadora usa para hacer sus cálculos. En lugar de usar números súper precisos y pesados (como una fotografía detallada), han estado cambiando a números más pequeños y ligeros (como un boceto pixelado). Esto se llama "entrenamiento de baja precisión". Recientemente, ha llegado un nuevo tipo de número diminuto llamado "FP4" (punto flotante de 4 bits), que promete hacer estos cerebros gigantes mucho más pequeños y rápidos. Pero este es el problema: mientras que las partes principales del cerebro pueden manejar estos números diminutos, las partes "ayudantes" —como la memoria que guarda para aprender y las matemáticas complejas que realiza para entender las oraciones— siguen fallando o confundiéndose cuando se ven obligadas a usar números tan pequeños. Es como tener el motor de un auto de carreras pero intentar correrlo con la cadena de una bicicleta; el motor está listo, pero el resto de la máquina no puede seguirle el ritmo.

Este artículo presenta un nuevo y hábil conjunto de herramientas llamado Full-Stack FP4, que actúa como un maestro mecánico para estos cerebros digitales. Los investigadores se dieron cuenta de que no puedes usar las mismas reglas de números diminutos para cada parte del cerebro. Algunas partes, como las conexiones principales, pueden manejar los números diminutos sin problemas. Pero otras partes, como la "memoria" que la computadora usa para recordar lo que aprendió ayer, son muy sensibles y necesitan un cuidado especial. El equipo creó un sistema modular donde utilizan diferentes "recetas" para diferentes partes del cerebro. Para las conexiones principales, utilizan un truco llamado LoRA-SVD, que es como mantener un boceto de alta resolución de los detalles más importantes mientras se usan bloques pixelados para el resto. Para la memoria, inventaron una forma de "suavizar" los números antes de comprimirlos en cajas diminutas para que no se pierdan. Para las matemáticas que ayudan al cerebro a enfocarse en las palabras correctas (atención), decidieron mantener las partes más sensibles en alta definición mientras dejan que el resto permanezca pixelado.

Cuando probaron este sistema completo en un modelo de 3 mil millones de parámetros (un cerebro digital de tamaño mediano) utilizando 64 mil millones de palabras de datos de entrenamiento, los resultados fueron impresionantes. El cerebro "Full-Stack FP4" aprendió casi exactamente tan bien como la versión tradicional de alta capacidad. La diferencia en sus puntuaciones de aprendizaje fue un minúsculo 0.838%, algo que apenas se nota. De hecho, cuando probaron qué tan bien el cerebro respondía preguntas que nunca había visto antes, la nueva versión de números diminutos fue en realidad ligeramente mejor adivinando las palabras correctas, a pesar de que estaba usando mucha menos memoria. En una sola tarjeta gráfica potente (una RTX 5090), este nuevo método hizo que el proceso de aprendizaje del cerebro fuera de 2.5 a 2.8 veces más rápido para ciertas tareas y utilizó entre un 38% y un 42% menos de memoria que el método estándar de alta precisión.

Los investigadores advierten cuidadosamente que, si bien esto funciona de maravilla para modelos de hasta 3 mil millones de parámetros, aún no han demostrado si funcionará para cerebros aún más grandes o durante periodos de entrenamiento mucho más largos. También descubrieron que no puedes simplemente lanzar números diminutos a todo; si intentas forzar las partes más sensibles del cerebro a usar los números más pequeños sin estas recetas especiales, el aprendizaje se rompe. Pero por ahora, este enfoque "Full-Stack" muestra que podemos construir una IA más inteligente, rápida y barata al tratar las diferentes partes del cerebro con el nivel de precisión específico que realmente necesitan, en lugar de forzar a todo el sistema a ser de un solo tamaño para todos.

Resumen Técnico: Full-Stack FP4: Preentrenamiento de LLM con Proyecciones, Optimizadores y Atención Cuantizados

1. Declaración del Problema

Los avances recientes en las GPU NVIDIA Blackwell han convertido el soporte nativo de Tensor Core NVFP4 (4 bits) en un objetivo práctico para el preentrenamiento de Modelos de Lenguaje de Gran Escala (LLM). Si bien la investigación existente ha demostrado con éxito el preentrenamiento estable de 4 bits para proyecciones lineales de Transformer, persisten brechas significativas en otros módulos críticos de entrenamiento. La mayoría de las recetas actuales de NVFP4 mantienen una mayor precisión (por ejemplo, BF16 o FP8) para los estados del optimizador, los cálculos del optimizador y los mecanismos de atención.

El desafío principal es que estos módulos poseen estructuras numéricas distintas que una regla de cuantización uniforme no puede abordar:

Proyecciones Lineales: Las señales de gradiente se debilitan durante el entrenamiento mientras que el ruido de cuantización permanece constante, lo que reduce la eficacia del entrenamiento de baja precisión continuo.
Estados de AdamW: El momento del segundo momento es no negativo, de cola pesada, persistente y sensible a los valores del denominador, lo que hace que el almacenamiento directo en 4 bits sea propenso a la distorsión.
Optimizadores Root/Muon: Estos dependen de productos de matrices Newton–Schulz repetidos donde los errores de iteración de baja precisión pueden acumularse y desestabilizar la convergencia.
Atención: El operador contiene rutas sensibles al softmax (por ejemplo, $PV$, $dOV^\top$ ) que requieren alta precisión, mientras que otras rutas (por ejemplo, $QK^\top$ ) son más robustas a la cuantización.

Consecuentemente, existe una falta de un marco unificado y modular que permita el preentrenamiento nativo en NVFP4 de todo el stack de entrenamiento sin depender de fallos de precisión mixta para componentes críticos.

2. Metodología: Full-Stack FP4

Los autores presentan Full-Stack FP4, un marco modular compuesto por cuatro recetas independientes y combinables diseñadas para abordar las dificultades numéricas específicas de cada módulo de entrenamiento.

2.1 Proyecciones Lineales LoRA-SVD

Para mitigar la degradación de las señales de gradiente en las etapas finales del entrenamiento, los autores proponen LoRA-SVD.

Mecanismo: Este método parametriza los pesos como $W = W_{res} + \beta L_2 L_1$ . Mantiene un subespacio principal compacto ( $L_2 L_1$ ) en BF16 mientras computa el residuo dominante ( $W_{res}$ ) en NVFP4.
Paso hacia atrás (Backward Pass): Un paso hacia atrás restringido al subespacio utiliza una restricción de Cholesky-QR para asegurar que las actualizaciones de la rama de bajo rango en BF16 se alineen con el subespacio principal, mientras que el residuo actualiza el componente NVFP4.
Realineación: El subespacio se realinea periódicamente (cada 2,048 pasos) utilizando SVD aleatorizado para resetear los componentes, asegurando que el subespacio BF16 permanezca óptimo.

2.2 AdamW de Estado Cuantizado (Q-AdamW)

Para estabilizar el almacenamiento de los estados de momento persistentes y de cola pesada en NVFP4:

Primer Momento ( $m_t$ ): Aplica mezcla de Hadamard para suavizar la distribución antes del Redondeo Estocástico (SR) y la cuantización adaptativa de 4/6 bits.
Segundo Momento ( $v_t$ ): Utiliza un pipeline de preprocesamiento de tres etapas antes del almacenamiento en NVFP4:
1. Raíz Cuadrada: Acorta la cola pesada.
2. Media de Teselas (Tile Mean): Separa el componente común positivo (almacenado en BF16) de un residuo aproximadamente de media cero.
3. Transformada de Hadamard: Distribuye la energía localizada en el residuo para hacerlo más simétrico.
Reconstrucción: El estado se reconstruye invirtiendo la transformada y la separación de la media, y luego elevando al cuadrado el resultado. Este pipeline evita la necesidad de operaciones de subespacio basadas en SVD.

2.3 Optimizador Root NVFP4 (Q-Root)

Para las iteraciones de Newton–Schulz (NS) utilizadas en el optimizador Root:

Ejecución Directa en NVFP4: En lugar de una ruta de precisión mixta, los productos de matrices NS se ejecutan directamente en NVFP4.
Estabilización: La estabilidad se logra mediante coeficientes dependientes de la forma (coeficientes $a, b, c$ optimizados para relaciones de aspecto de matriz específicas) y recorte de outliers p99. Este enfoque limita la amplificación del error a través de sucesivos pasos de NS sin requerir rutas auxiliares de precisión mixta.

2.4 Atención de Precisión Mixta (Q-Attn)

Para equilibrar la precisión y la eficiencia en la atención:

División de Precisión: El framework asigna NVFP4 a las rutas robustas ( $QK^\top$ y $dS$) y retiene BF16 para las rutas sensibles al softmax ($PV$, $P^\top dO$ , y $dOV^\top$ ).
Consistencia: Las vistas cuantizadas de $Q, K, V$ se reutilizan entre los pasos hacia adelante (forward) y hacia atrás (backward) para asegurar la consistencia de los tensores y evitar el tráfico redundante de memoria.

3. Contribuciones Clave

Proyecciones LoRA-SVD: Una reparametrización combinable que reduce la brecha de pérdida solo lineal de 1.40% a 0.61% comparado con el NVFP4 estándar, protegiendo un subespacio compacto en BF16.
AdamW de Estado Cuantizado: La primera demostración del almacenamiento estable de los estados de momento de AdamW en NVFP4 en una ejecución de preentrenamiento integrada, utilizando un pipeline de estado transformado que reduce el error de reconstrucción del segundo momento en aproximadamente un 60%.
Root NVFP4: Una ruta directa de baja precisión para Newton–Schulz habilitada por coeficientes conscientes de la forma y recorte, reduciendo el error de iteración de ~10% a ~3%.
Atención de Precisión Mixta Entrenable: Una estrategia que protege las rutas sensibles al softmax en BF16 mientras cuantiza las rutas de forward y backward adecuadas, permitiendo un operador de preentrenamiento consistente en NVFP4.

4. Resultados Experimentales

Los autores evaluaron el framework en un Transformer de 3B de parámetros entrenado durante 64B de tokens.

Pérdida de Entrenamiento: El baseline de BF16 alcanzó una pérdida de 2.267, mientras que Full-Stack FP4 alcanzó 2.286, resultando en una brecha relativa de 0.838%.
Rendimiento en Tareas Aguas Abajo (Downstream): La evaluación zero-shot mostró que Full-Stack FP4 tuvo una perplejidad promedio de 26.665 (frente a 26.675 para BF16) y una precisión promedio 0.10 puntos porcentuales menor que el baseline de BF16.
Estudios de Ablación:
- LoRA-SVD por sí solo redujo la brecha de pérdida solo lineal de 1.40% a 0.61%.
- Las recetas modulares pueden habilitarse de forma independiente o combinarse, siendo "Full-Stack" la representación de la cobertura agregada.
Eficiencia de Hardware Nativo (RTX 5090):
- Aceleración de Root: Full-Stack FP4 logró aceleraciones de 2.50×–2.83× sobre el BF16 optimizado y 1.83×–2.17× sobre NVIDIA Transformer Engine (TE) para el paso del optimizador Root.
- Ahorro de Memoria: El framework redujo el pico de memoria de AdamW en un 37.9%–42.5% comparado con BF16.
- Ahorro de Tiempo: Los ahorros de tiempo con pérdida igualada (matched-loss) se calcularon en un 56.9%–61.9% al contabilizar la reducción del tiempo de paso y el ligero aumento de tokens requeridos para igualar la pérdida de BF16.

5. Significado y Limitaciones

Significancia:
El artículo demuestra que un enfoque de "Full-Stack" para el preentrenamiento en NVFP4 es factible, cubriendo proyecciones, estados del optimizador, computación del optimizador y atención. Establece que el preentrenamiento nativo en 4 bits puede lograr un rendimiento comparable a BF16 (dentro de una brecha de pérdida <1%) ofreciendo al mismo tiempo beneficios sustanciales de memoria y throughput en hardware de clase Blackwell. El trabajo proporciona un conjunto de recetas modulares que permiten a los practicantes adoptar la baja precisión de forma incremental.

Limitaciones:

Escala: Los experimentos están limitados a modelos de hasta 3B de parámetros y 64B de tokens. El artículo no establece el comportamiento en escalas de 7B+ o durante horizontes de entrenamiento más largos.
Medición de Eficiencia: La eficiencia de hardware nativo se midió en cuatro bloques de decodificador en una sola RTX 5090, no en un modelo distribuido completo.
Estimaciones: Los cálculos de ahorro de tiempo son estimaciones basadas en configuraciones fijas de batch size y no consideran las ganancias potenciales de throughput al aumentar el batch size usando la memoria liberada.
Fusión de Kernels: Algunos kernels (backward de LoRA, atención) aún no están totalmente fusionados, y la implementación aún no utiliza todas las optimizaciones de FlashAttention-4.
Especificidad de Formato: Las recetas y mediciones son específicas de NVFP4 y no pretenden ser transferibles a MXFP4.