🤖 AI

TetraJet-v2: Accurate NVFP4 Training for Large Language Models with Oscillation Suppression and Outlier Control

TetraJet-v2 es un método de entrenamiento totalmente cuantizado en 4 bits de extremo a extremo para Modelos de Lenguaje Grande que utiliza precisión NVFP4 en todas las capas lineales e introduce OsciReset y OutControl para suprimir eficazmente la oscilación de los pesos y controlar los valores atípicos, logrando así un rendimiento casi sin pérdidas con aceleraciones significativas en comparación con FP8.

Autores originales: Yuxiang Chen, Yifan Liu, Xiaoming Xu, Pengle Zhang, Michael Beyer, Martin Rapp, Jun Zhu, Jianfei Chen

Publicado 2026-05-06

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Yuxiang Chen, Yifan Liu, Xiaoming Xu, Pengle Zhang, Michael Beyer, Martin Rapp, Jun Zhu, Jianfei Chen

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 una cantidad masiva de información a un estudiante gigante y brillante (un Modelo de Lenguaje Grande). Por lo general, para hacerlo de manera efectiva, necesitas un aula de alta calidad y costosa, con iluminación perfecta y pizarras gigantes (computación de alta precisión). Esto cuesta una fortuna en electricidad y hardware.

El artículo presenta un nuevo método llamado TetraJet-v2 que intenta enseñar a este mismo estudiante utilizando un aula diminuta y de bajo costo, con una pequeña pizarra de tiza y luces tenues (precisión de 4 bits). El objetivo es hacer que el proceso de entrenamiento sea 1,67 veces más rápido y mucho más barato, sin que el estudiante aprenda cosas incorrectas o olvide lo que acaba de aprender.

Aquí se explica cómo resolvieron los problemas que suelen ocurrir cuando intentas enseñar en un aula tan "de baja calidad":

1. El Problema: La "Mano Temblorosa" (Oscilación de Pesos)

Imagina que el estudiante está intentando escribir un número en la pizarra. En un aula perfecta, pueden escribir "0,25" exactamente. Pero en esta aula de baja precisión, solo pueden escribir números enteros o medios números (como 0 o 0,5).

Si la respuesta real del estudiante está exactamente en el medio (0,25), su mano empieza a temblar. Un momento escriben "0", al siguiente escriben "0,5", luego vuelven a "0". Quedan atrapados en un bucle, oscilando de un lado a otro. En el artículo, llaman a esto Oscilación de Pesos. Esto desperdicia energía y evita que el estudiante se asiente nunca en la respuesta correcta.

La Solución: OsciReset (El Truco de "Centrar")
Los autores inventaron una técnica llamada OsciReset. Imagina a un profesor que vigila la mano del estudiante. En cuanto ve que el estudiante empieza a temblar entre dos opciones, el profesor toma suavemente la tiza y la coloca justo en el medio de la zona segura (el "centro del contenedor").

Por qué funciona: Detiene el temblor inmediatamente. No congela al estudiante (lo cual detendría el aprendizaje); simplemente le da un punto de partida estable para continuar aprendiendo en una nueva dirección, más suave. Esta es la primera vez que se resuelve este problema específico de "temblor" para estos modelos gigantes.

2. El Problema: Los "Gritones" (Valores Atípicos)

En un aula normal, la mayoría de los estudiantes hablan a un volumen normal. Pero en estos modelos gigantes, unos pocos "canales" específicos (piensa en ellos como micrófonos específicos) empiezan de repente a gritar a volumen máximo. Estos se llaman Valores Atípicos.

Como la pizarra es tan pequeña (baja precisión), si un micrófono grita demasiado fuerte, el profesor tiene que bajar el volumen para todos los demás para que el que grita no rompa la pizarra. Esto significa que los susurros importantes y silenciosos de los otros estudiantes se pierden en el ruido.

La Solución: OutControl (Los "Micrófonos VIP")
Los autores se dieron cuenta de que estos micrófonos gritones son predecibles; siempre son los mismos pocos canales. Así que crearon un sistema llamado OutControl.

Cómo funciona: Identifican los "micrófonos VIP" (los valores atípicos) desde el principio y les dan su propio micrófono especial y de alta calidad (manteniéndolos en un formato de mayor precisión como FP8). El resto de los estudiantes continúa usando la pequeña pizarra barata (FP4).
Resultado: El gritón no obliga a bajar el volumen para todos los demás, por lo que los susurros silenciosos siguen escuchándose claramente.

3. La Pizarra de "Doble Bloque"

Para asegurar que las matemáticas funcionen perfectamente en esta pizarra diminuta, rediseñaron cómo se escriben los números. En lugar de intentar caber un número enorme en un solo cuadrado pequeño, utilizan un sistema de Doble Bloque.

Imagina a un gran grupo de estudiantes (un bloque de 128) compartiendo una sola regla grande (una escala global), pero dentro de ese grupo, están divididos en equipos más pequeños de 16, cada uno con su propia regla diminuta. Esto les permite manejar tanto números muy pequeños como números muy grandes sin perder precisión.

Los Resultados

Cuando juntaron todos estos trucos (el truco de centrar, los micrófonos VIP y la pizarra de doble bloque):

Velocidad: Entrenaron los modelos 1,67 veces más rápido que el estándar actual (FP8).
Precisión: Los modelos aprendieron casi tan bien como si hubieran sido entrenados en el aula costosa y de alta precisión. Redujeron la brecha de rendimiento entre el método barato y el método costoso en más del 50%.
Escala: Probaron esto en modelos de hasta 370 millones de parámetros, entrenándolos con más de 200 mil millones de palabras.

En resumen: TetraJet-v2 es un nuevo conjunto de reglas que nos permite entrenar modelos gigantes de IA en hardware mucho más barato y rápido, deteniendo sus temblores, protegiendo las voces fuertes y organizando la pizarra de manera más eficiente.

Resumen Técnico: TetraJet-v2

1. Enunciado del Problema

El preentrenamiento de Modelos de Lenguaje Grande (LLM) es cada vez más prohibitivamente costoso, impulsando la necesidad de Entrenamiento Totalmente Cuantizado de Baja Precisión (FQT). Si bien formatos de 4 bits como NVFP4 ofrecen ganancias de eficiencia significativas frente a FP8 y BF16, lograr un entrenamiento estable y casi sin pérdidas a esta precisión sigue siendo un desafío crítico. Los intentos anteriores de entrenamiento totalmente en 4 bits han sufrido una degradación severa del rendimiento.

Los autores identifican dos cuellos de botella de optimización principales que obstaculizan el entrenamiento totalmente en FP4, distintos del simple error de cuantización:

Oscilación de Pesos: Un modo de fallo dominante donde pequeñas actualizaciones en los pesos maestros de alta precisión provocan que sus representaciones cuantizadas cambien repetidamente entre contenedores adyacentes (por ejemplo, entre 0 y 0.5). Este fenómeno interrumpe la optimización, particularmente a medida que las tasas de aprendizaje decaen, impidiendo que el modelo converja a un mínimo local estable.
Características Atípicas (Outliers): Un pequeño número de canales de activación exhibe magnitudes masivas, creando rangos dinámicos que FP4 no puede representar con precisión. Esto conduce a factores de escala altos que destruyen la precisión de los valores más pequeños dentro del mismo grupo de cuantización.

Los métodos existentes, como la receta NVFP4 de NVIDIA, a menudo dependen de precisión mixta (manteniendo algunos bloques en BF16) o no abordan completamente estas dinámicas de optimización en un escenario de preentrenamiento desde cero.

2. Metodología: TetraJet-v2

TetraJet-v2 es un método de FQT de 4 bits de extremo a extremo diseñado específicamente para NVFP4, aplicando el formato a activaciones, pesos y gradientes en todas las capas lineales. La metodología aborda los cuellos de botella identificados mediante tres componentes principales:

A. Cuantización Doble de Bloque No Sesgada para Capas Lineales NVFP4

Para maximizar la precisión dentro de las restricciones numéricas de NVFP4 (que utiliza grupos de 16 elementos con escala de 8 bits), los autores proponen un esquema de cuantización doble de bloque:

Estructura: Un bloque externo de 128 elementos se escala globalmente, mientras que los bloques internos de 16 elementos se escalan localmente. Esto evita calcular el valor absoluto máximo sobre todo el tensor, mejorando la compatibilidad con el hardware y la precisión en comparación con la escala por tensor.
Estimación de Gradientes No Sesgada: El diseño de la capa lineal garantiza una estimación de gradientes no sesgada mediante:
- Alinear los tensores en la pasada hacia atrás (estimando $\hat{X}$ en lugar de $X$ ).
- Utilizar Redondeo Estocástico en la pasada hacia atrás para asegurar que $E[\text{round}(x)] = x$ , evitando sesgos en las expectativas de los gradientes.
- Adoptar una forma de grupo de pesos de $1 \times 16$ (no cuadrada) para mayor precisión, requiriendo una alineación específica para la matriz de pesos transpuesta.

B. OsciReset: Supresión de Oscilación

Para abordar la oscilación de pesos, los autores introducen OsciReset, el primer algoritmo diseñado específicamente para suprimir este cuello de botella en el preentrenamiento de LLM.

Detección: El método rastrea la trayectoria de optimización de los pesos, calculando una métrica de "Riesgo de Oscilación" definida como la relación entre la distancia recorrida por el peso cuantizado ( $dist_Q$ ) y la distancia recorrida por el peso maestro ( $dist_M$ ). Un riesgo alto indica que un peso está cambiando de contenedor a pesar de pequeñas actualizaciones en el peso maestro.
Mecanismo: Cuando la tasa de aprendizaje es baja y el riesgo de oscilación supera un umbral, el peso maestro del elemento oscilante se restablece al centro de su contenedor de cuantización actual.
Ventaja: A diferencia de métodos anteriores que congelan pesos o utilizan regularización pesada (lo cual perjudica la optimización global), OsciReset permite que el peso continúe optimizándose en una nueva trayectoria menos propensa a la oscilación, sin degradar el rendimiento inmediatamente.

C. OutControl: Control de Atípicos

Para gestionar los valores atípicos, TetraJet-v2 emplea una estrategia híbrida:

Transformación Hadamard Aleatoria (RHT): Aplicada a los cálculos de la pasada hacia atrás ($dX $y$ dW$) para distribuir la energía de los valores atípicos de manera más uniforme entre los canales, mitigando la varianza de cuantización.
Retención Estática de Atípicos: Reconociendo que los canales atípicos son estructuralmente consistentes (canales específicos siempre tienen normas altas), el método identifica el top $p\%$ (por ejemplo, 10%) de canales atípicos durante una breve ventana de calibración. Estos canales se retienen en mayor precisión (por ejemplo, FP8) tanto para la pasada hacia adelante como para la hacia atrás, mientras que el resto de la red permanece en FP4. Esta selección estática evita la sobrecarga computacional de ajustes dinámicos por lote.

3. Contribuciones Clave

Cuantización Doble de Bloque No Sesgada: Un diseño novedoso de capa lineal NVFP4 con estimación de gradientes no sesgada y formas de bloque óptimas ( $1 \times 16$ ) para LLM, superando las recetas NVFP4 existentes.
OsciReset: El primer algoritmo efectivo para suprimir la oscilación de pesos en LLM durante el preentrenamiento, resolviendo una fuente principal de inestabilidad en el entrenamiento de baja precisión.
OutControl: Un algoritmo de precisión mixta que combina RHT para la pasada hacia atrás con retención estática de atípicos para ambas pasadas (hacia adelante y hacia atrás) para preservar la precisión.
Evaluación Exhaustiva: Experimentos extensos de preentrenamiento en modelos OLMo-2 (hasta 370M de parámetros, 212B de tokens) que demuestran la eficacia de la pila completa.

4. Resultados Experimentales

Los autores evaluaron TetraJet-v2 en modelos OLMo-2 (70M, 150M y 370M de parámetros) entrenados desde cero.

Reducción de la Brecha de Rendimiento: TetraJet-v2-full (incorporando todos los métodos propuestos) redujo la brecha de rendimiento frente al entrenamiento de precisión completa (BF16) en un promedio del 51.3% a través de los tamaños de modelo.
Perplejidad (PPL): En el modelo de 370M entrenado con 212B de tokens, TetraJet-v2-full logró una PPL de validación de 25.11, en comparación con 23.70 para BF16 y 26.20 para la línea base NVFP4 de NVIDIA.
Tareas de Descenso: El método logró el rendimiento promedio más alto entre todos los métodos FP4 en una suite de puntos de referencia (ARC, BoolQ, MMLU, etc.), con un modelo de 370M alcanzando una precisión promedio del 43.60%.
Eficiencia: Las implementaciones de kernels CUDA lograron una aceleración de extremo a extremo de 1.67× sobre FP8 (TransformerEngine) en hardware RTX 5090, verificando ganancias de rendimiento prácticas a pesar de la complejidad añadida de la supresión de oscilación y el control de atípicos.
Estudios de Ablación:
- Eliminar OsciReset condujo a una degradación severa del rendimiento y un rebote en la pérdida.
- OutControl mejoró la precisión tanto en la pasada hacia adelante como en la hacia atrás.
- Se demostró que el método es robusto a las elecciones de hiperparámetros (por ejemplo, umbral de oscilación, tiempo de inicio).

5. Significado y Afirmaciones

El artículo afirma que TetraJet-v2 representa un paso significativo hacia la viabilidad y competitividad del preentrenamiento totalmente de baja precisión (4 bits) de LLM frente al entrenamiento de alta precisión.

Optimización vs. Representación: Los autores enfatizan que la brecha de rendimiento en el entrenamiento de baja precisión no se debe únicamente al error de cuantización (representación), sino que está impulsada críticamente por problemas de optimización (oscilación y atípicos) que los métodos anteriores no abordaron sistemáticamente.
Generalizabilidad: A diferencia de las técnicas anteriores de mitigación de oscilación diseñadas para Vision Transformers (que a menudo dependen de ajuste fino o sintonización específica de hiperparámetros), OsciReset se presenta como un método simple y efectivo aplicable al preentrenamiento de LLM desde cero con un ajuste mínimo.
Practicidad: Al lograr una aceleración de 1.67× sobre FP8 mientras reduce la brecha de precisión con BF16 en más del 50%, el trabajo demuestra que el entrenamiento totalmente cuantizado puede ofrecer tanto eficiencia computacional como calidad de modelo, reduciendo potencialmente los costos de hardware y energía de la investigación de modelos fundamentales.

Los autores reconocen limitaciones relacionadas con la escala de sus experimentos (hasta 370M de parámetros) y la falta actual de soporte de hardware mainstream para estrategias de precisión mixta como FP6 $\times$ FP4, que exploraron como una posible dirección futura para obtener mayores ganancias de precisión.