NVIDIA acelera el preentrenamiento de LLM con NVFP4 en Blackwell, JAX y MaxText

El preentrenamiento de LLM frontier choca con el límite de rendimiento de los sistemas de cómputo. NVIDIA ha mostrado cómo la combinación de JAX, MaxText y el nuevo formato NVFP4 en chips Blackwell permite acelerar de forma notable este proceso sin perder calidad.

Por qué importa cada punto porcentual

Cuando el entrenamiento se hace con billones de tokens sobre miles de aceleradores, ahorrar incluso un uno por ciento de tiempo en cada paso se traduce en varios días de tiempo real de calendario. A escala frontier, eso equivale directamente a millones de dólares en costes de cómputo.

NVFP4, un formato de coma flotante de cuatro bits que debutó en la arquitectura Blackwell, se ha convertido en una de las herramientas clave para acelerar las operaciones matriciales. En comparación con FP8, empaqueta los números con el doble de densidad, lo que reduce la carga de memoria y aumenta el ancho de banda efectivo de los núcleos tensoriales.

La principal dificultad es que la cuadrícula numérica de cuatro bits es dispersa. Con una configuración incorrecta, los gradientes se salen fácilmente de su rango, lo que provoca la divergencia del entrenamiento. NVIDIA y el equipo de MaxText lo resolvieron con esquemas de escalado personalizados y loss scaling dinámico.

Cómo funciona la precisión mixta con NVFP4

El entrenamiento de precisión mixta no es un enfoque nuevo: FP8 y BF16 ya son estándares de la industria. NVFP4 va un paso más allá y permite usar pesos de 4 bits en las multiplicaciones matriciales más pesadas, manteniendo mayor precisión donde realmente importa.

• NVFP4 se aplica a pesos y activaciones en las operaciones GEMM
• BF16 o FP32 se mantienen para acumuladores y normalización
• MaxText enruta automáticamente las operaciones al formato adecuado
• JAX compila el grafo de cómputo mediante XLA y optimiza los kernels para Blackwell
• El resultado es un mayor throughput con un consumo energético similar o menor

El stack y qué cambiar en el código

MaxText es un framework de entrenamiento de código abierto y alto rendimiento basado en JAX y desarrollado por Google. Se creó originalmente para TPU, pero se adapta activamente a clústeres de GPU, por lo que la alianza con NVIDIA resulta lógica.

NVIDIA ha incluido kernels NVFP4 de bajo nivel en cuBLAS y cuDNN, y JAX/XLA ha incorporado soporte para estas operaciones mediante adaptadores específicos. Los desarrolladores no necesitan reescribir el código de entrenamiento a mano: basta con activar los flags adecuados en la configuración de MaxText y comprobar que el clúster tiene chips Blackwell (B100, B200, GB200).

«La precisión numérica es uno de los parámetros con más palanca, pero el preentrenamiento de precisión mixta de bajos bits es difícil de implementar bien», señala el equipo de NVIDIA

Developer Blog.

Qué significa

Para los equipos que preentrenan modelos frontier, NVFP4 en Blackwell es una aceleración prácticamente gratuita: el stack actual de JAX y MaxText requiere cambios mínimos de configuración. A escala de cientos o miles de GPU, incluso un aumento del 10-15 % del throughput reduce directamente el time-to-checkpoint y el presupuesto total de cómputo.

La carrera por la eficiencia del preentrenamiento entra en una fase de batalla por la precisión numérica.