Pretraining Large Language Models with NVFP4

Este estudio presenta un método innovador para el preentrenamiento estable y preciso de modelos de lenguaje grandes en formato NVFP4, que mediante transformaciones de Hadamard aleatorias, esquemas de cuantización bidimensionales y redondeo estocástico, logra un rendimiento comparable al de la precisión FP8 en un modelo de 12 mil millones de parámetros entrenado con 10 billones de tokens.

Lo esencial

¡Hola! Imagina que entrenar una Inteligencia Artificial (IA) gigante es como construir un rascacielos. Para hacerlo, necesitas millones de ladrillos (datos) y una grúa muy potente (computadoras).

Hasta ahora, para construir estos rascacielos digitales, usábamos "ladrillos" de un tamaño estándar (llamados FP8). Funcionaban bien, pero eran pesados y costosos de mover. NVIDIA se preguntó: "¿Y si pudiéramos usar ladrillos mucho más pequeños (4 bits) para hacer la construcción el doble de rápida y con la mitad de energía?".

El problema es que los ladrillos pequeños son frágiles. Si intentas construir un rascacielos entero con ellos, la estructura tiembla y se cae.

Este documento de NVIDIA presenta una solución genial llamada NVFP4. Es como un "kit de construcción" especial para usar esos ladrillos diminutos sin que el edificio se derrumbe. Aquí te explico cómo funciona con analogías sencillas:

1. El Problema: Los "Gigantes" y los "Enanos"

En el mundo de las matemáticas de la IA, hay números muy grandes (gigantes) y números muy pequeños (enanos).

• El problema: Cuando usas ladrillos pequeños (4 bits), si hay un "gigante" (un número muy grande) en tu grupo de datos, ocupa todo el espacio y los "enanos" se vuelven invisibles (se convierten en cero). Esto hace que la IA aprenda mal.
• La solución de NVFP4: Imagina que NVFP4 es un sistema de organización inteligente. En lugar de tener una sola regla para todo el grupo, divide los datos en grupos pequeños (bloques de 16). Si hay un gigante en un grupo, NVFP4 ajusta la escala de ese grupo específico para que el gigante quepa, sin aplastar a los enanos. Además, usa una "regla maestra" (a nivel de todo el tensor) para asegurar que nada se pierda. Es como tener un equipo de arquitectos que ajustan las escaleras de cada habitación individualmente para que todos quepan.

2. La Estrategia: No todo puede ser de plástico

Aunque queremos usar ladrillos pequeños para todo, hay partes del edificio que son tan delicadas que no pueden soportar la fragilidad de los ladrillos de 4 bits.

• La analogía: Imagina que estás construyendo un castillo de naipes. Puedes usar papel reciclado (4 bits) para las paredes, pero si intentas hacer la base o la punta con papel reciclado, se doblará y se caerá.
• La solución: NVIDIA deja las partes más críticas (las últimas capas del modelo y algunas específicas) en "papel de alta calidad" (precisión BF16 o FP8). El resto del edificio se construye con los ladrillos pequeños. Esto asegura que la estructura sea sólida.

3. El Truco de la "Bolsa de Canicas" (Transformaciones Hadamard)

A veces, los datos tienen patrones extraños que confunden a la IA.

• La analogía: Imagina que tienes una bolsa con canicas de colores. Si las sacas en el orden en que vinieron, las rojas siempre caen juntas y las azules se pierden.
• La solución: NVFP4 usa una técnica llamada Transformación Hadamard Aleatoria. Es como agitar la bolsa fuertemente antes de sacar las canicas. Esto mezcla los datos de tal forma que los "gigantes" (valores raros) se dispersan y se vuelven más fáciles de manejar para los ladrillos pequeños. Es como mezclar bien la pintura para que no queden grumos.

4. La Brújula Invertida (Redondeo Estocástico)

Cuando reduces el tamaño de los números, a veces tienes que "redondear". Si siempre redondeas hacia arriba o siempre hacia abajo, introduces un error sistemático (como una brújula que siempre apunta 5 grados al norte).

• La analogía: Imagina que tienes que adivinar si una moneda está en la mano izquierda o derecha. Si siempre adivinas "izquierda", te equivocas mucho si la moneda está a la derecha.
• La solución: NVFP4 usa Redondeo Estocástico. En lugar de decidir siempre igual, "tira un dado" para decidir hacia dónde redondear. A veces arriba, a veces abajo. A la larga, los errores se cancelan entre sí y la IA aprende con más precisión. Es como promediar muchas opiniones diferentes en lugar de seguir una sola.

5. El Resultado: ¡El edificio se mantiene en pie!

NVIDIA probó esto construyendo un "rascacielos" de 12 mil millones de parámetros (un modelo muy grande) con 10 billones de palabras de datos.

• El hallazgo: El edificio construido con ladrillos pequeños (NVFP4) quedó tan fuerte y alto como el construido con ladrillos grandes (FP8).
• La ventaja: Al usar ladrillos más pequeños, la construcción fue más rápida y consumió mucha menos energía. Además, comparado con otro método de ladrillos pequeños (MXFP4), NVFP4 construyó el edificio con menos ladrillos (menos datos necesarios) para lograr el mismo resultado.

En resumen

NVIDIA ha descubierto cómo usar una versión "miniatura" de los datos para entrenar IAs gigantes. Lo lograron no solo haciendo los ladrillos más pequeños, sino creando un sistema de construcción inteligente que:

1. Ajusta las escalas localmente.
2. Protege las partes delicadas del edificio.
3. Mezcla los datos para evitar atascos.
4. Usa el azar para corregir errores.

Esto significa que en el futuro, podríamos tener IAs más inteligentes que entrenan más rápido y consumen menos electricidad, lo cual es una gran noticia para el planeta y para la tecnología.

Resumen técnico

Aquí tienes un resumen técnico detallado del informe técnico de NVIDIA sobre el preentrenamiento de Modelos de Lenguaje Grande (LLM) utilizando el formato NVFP4.

Título: Preentrenamiento de Modelos de Lenguaje Grande con NVFP4

1. El Problema

El entrenamiento de modelos de lenguaje de vanguardia (LLMs) requiere una inversión masiva de recursos computacionales, tiempo y energía, alcanzando órdenes de magnitud de yottaFLOPS. Aunque el formato de punto flotante de 8 bits (FP8) se ha adoptado ampliamente para mejorar la eficiencia, existe una necesidad crítica de reducir aún más la precisión a 4 bits (FP4) para duplicar o triplicar el rendimiento aritmético y reducir el uso de memoria a la mitad.

Sin embargo, la cuantización a 4 bits presenta desafíos significativos:

• Inestabilidad en el entrenamiento: La precisión reducida puede causar divergencia en el entrenamiento, especialmente en modelos grandes y en horizontes de tokens largos.
• Problemas de convergencia: Los valores atípicos (outliers) y los sesgos en la estimación de gradientes pueden degradar la capacidad del modelo para aprender.
• Limitaciones de formatos existentes: Los formatos de microescalado estándar (como MXFP4) pueden perder rango dinámico o precisión debido a sus factores de escala limitados, afectando la representación de valores críticos.

2. Metodología Propuesta

NVIDIA introduce un enfoque integral para el entrenamiento estable y preciso utilizando el formato NVFP4 (un formato de 4 bits optimizado por NVIDIA que extiende el enfoque de "microescalado"). La metodología se basa en cuatro pilares técnicos clave:

• Formato NVFP4:

  • A diferencia de MXFP4 (bloques de 32 elementos), NVFP4 utiliza bloques más pequeños de 16 elementos, lo que permite un ajuste más fino del rango dinámico local.
  • Emplea una estrategia de escalado de dos niveles: un factor de escala global en FP32 a nivel de tensor y un factor de escala local en E4M3 (4 bits de exponente, 3 de mantisa) a nivel de bloque. Esto permite representar el valor máximo del bloque con mayor precisión que los formatos que usan solo potencias de dos.
  • Soporte nativo en las Tensor Cores de la arquitectura NVIDIA Blackwell.

• Precisión Mixta Selectiva:

  • Se identifican capas lineales sensibles (especialmente las últimas capas del modelo) que son críticas para la convergencia.
  • La estrategia mantiene aproximadamente el 15-16% de las capas lineales (las primeras dos y las últimas ocho bloques en el modelo de 12B) en precisión alta (BF16), mientras que el resto de las operaciones GEMM se realizan en NVFP4.

• Transformadas de Hadamard Aleatorias (RHT):

  • Se aplican transformadas ortogonales aleatorias a las entradas de los gradientes de los pesos (Wgrad) para dispersar los valores atípicos (outliers) y distribuirlos en una aproximación de distribución gaussiana. Esto facilita su representación en el rango limitado de FP4.
  • Se utiliza un tamaño de matriz de Hadamard de 16x16 y un vector de signos aleatorio fijo compartido en todo el entrenamiento.

• Escalado Bidimensional (2D) y Redondeo Estocástico:

  • Escalado 2D: Para garantizar la consistencia entre el paso forward y backward (evitando violaciones de la regla de la cadena debido a diferentes representaciones cuantizadas), los pesos se escalan en bloques 16x16. Esto asegura que la misma representación cuantizada se use en ambas direcciones.
  • Redondeo Estocástico (Stochastic Rounding): Se aplica exclusivamente a los gradientes (no a pesos ni activaciones) para eliminar el sesgo sistemático introducido por el redondeo determinista, mejorando la estimación no sesgada de los gradientes.

3. Contribuciones Clave

• Validación a Escala Masiva: Se logra el primer entrenamiento documentado públicamente de un modelo de 12 mil millones de parámetros durante 10 billones de tokens utilizando precisión de 4 bits.
• Nuevo Formato NVFP4: Demostración de que el formato NVFP4, con su escalado E4M3 y bloques de 16 elementos, supera a los formatos MXFP4 existentes en términos de estabilidad y precisión.
• Metodología Robusta: Definición de un conjunto de técnicas (RHT, escalado 2D, redondeo estocástico selectivo) que son esenciales para la convergencia en modelos grandes y horizontes de entrenamiento largos.
• Soporte de Hardware: Integración completa en el motor Transformer Engine de NVIDIA para GPUs Blackwell.

4. Resultados

El modelo de 12B entrenado con NVFP4 se comparó contra un modelo de referencia entrenado en FP8:

• Pérdida de Entrenamiento (Loss): La curva de pérdida de validación del modelo NVFP4 sigue muy de cerca a la del modelo FP8. Durante la fase estable, el error relativo de pérdida se mantiene por debajo del 1%, ampliándose ligeramente a 1.5% al final del entrenamiento (durante la decadencia de la tasa de aprendizaje).
• Rendimiento en Tareas (Downstream Tasks):
  • MMLU-Pro: NVFP4 alcanzó una precisión de 62.58%, casi idéntico al 62.62% del modelo FP8.
  • MMLU, Matemáticas y Razonamiento: El rendimiento fue comparable o incluso superior en algunas métricas (ej. GSM8k CoT: 92.27% vs 89.08% en FP8).
  • Codificación: Hubo una ligera caída en tareas de codificación (MBPP+), atribuida posiblemente a evaluaciones ruidosas o a la necesidad de ajustar el checkpoint final.
• Comparación NVFP4 vs. MXFP4:
  • NVFP4 converge a una pérdida menor que MXFP4 con el mismo número de tokens.
  • Para igualar el rendimiento de NVFP4, MXFP4 requeriría un 36% más de tokens (1.36T vs 1T), lo que implica un aumento significativo en el tiempo de entrenamiento y coste computacional.

5. Significado e Impacto

Este trabajo representa un hito fundamental en la eficiencia del entrenamiento de IA:

1. Viabilidad del FP4: Demuestra que el entrenamiento de modelos de frontera con precisión de 4 bits es técnicamente viable y estable a escalas de billones de tokens, algo que antes se consideraba arriesgado o imposible.
2. Eficiencia de Recursos: Al reducir la precisión a 4 bits, se reduce el ancho de banda de memoria y el uso de memoria en un 50% respecto a FP8, y se aumenta el rendimiento matemático en 2x a 3x en hardware Blackwell.
3. Futuro de los LLMs: Abre la puerta a entrenar modelos más grandes o a reducir drásticamente el coste energético y temporal de los modelos existentes, facilitando la próxima generación de modelos de inteligencia artificial más capaces y accesibles.
4. Estandarización: Establece un nuevo estándar de mejores prácticas (mezcla de precisiones, transformadas de Hadamard, escalado 2D) para la comunidad de investigación en cuantización de LLMs.

En resumen, NVIDIA ha demostrado que, con las técnicas algorítmicas adecuadas, el formato NVFP4 puede ofrecer un rendimiento de entrenamiento casi equivalente al FP8, pero con una eficiencia de hardware muy superior, marcando un paso decisivo hacia el entrenamiento ultra-eficiente de modelos de inteligencia artificial.