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

El preentrenamiento de LLMs de frontera choca con los límites de throughput de los sistemas computacionales. NVIDIA ha demostrado cómo la combinación de JAX, MaxText y el nuevo formato NVFP4 en chips Blackwell permite acelerar significativamente este proceso sin pérdida de calidad.

Por Qué Cada Porcentaje Es Importante

Cuando el entrenamiento se realiza en billones de tokens a través de miles de aceleradores, ahorrar incluso un porcentaje de tiempo en cada paso se traduce en varios días de tiempo de calendario real. A la escala del preentrenamiento de frontera, esto es una conversión directa a millones de dólares en gastos de computación. NVFP4 — un formato de punto flotante de cuatro bits que debutó en la arquitectura Blackwell — se convirtió en una de las herramientas clave para acelerar operaciones matriciales.

Comparado con FP8, compacta números dos veces más densamente, lo que reduce la carga de memoria y aumenta el throughput efectivo de los núcleos de tensor. El principal desafío: la cuadrícula numérica de cuatro bits es dispersa. Con una configuración incorrecta, los gradientes fácilmente superan sus límites — esto lleva a la divergencia del entrenamiento.

NVIDIA y el equipo de MaxText abordaron esto a través de esquemas de escalado personalizados y escalado dinámico de pérdida.

Cómo Funciona Mixed-Precision con NVFP4

El entrenamiento mixed-precision no es un enfoque nuevo: FP8 y BF16 ya se han convertido en un estándar industrial. NVFP4 va un paso más allá, permitiendo pesos de 4 bits en las multiplicaciones matriciales más intensivas en computación manteniendo mayor precisión donde realmente importa.

• NVFP4 se aplica a pesos y activaciones en operaciones GEMM
• BF16 o FP32 permanecen para acumuladores y normalización
• MaxText enruta automáticamente operaciones al formato apropiado
• JAX compila el gráfico computacional a través de XLA, optimizando núcleos para Blackwell
• Resultado — crecimiento de throughput con consumo de energía comparable o menor

El Stack y Qué Cambiar en el Código

MaxText es un framework de entrenamiento de alto rendimiento de código abierto basado en JAX, desarrollado por Google. Fue creado originalmente para TPU, pero se está adaptando activamente para clusters de GPU, y la asociación con NVIDIA es un ajuste natural aquí. NVIDIA incluyó núcleos NVFP4 de bajo nivel como parte de cuBLAS y cuDNN, y JAX/XLA recibió soporte para estas operaciones a través de adaptadores especiales. Los desarrolladores no necesitan reescribir el código de entrenamiento manualmente — es suficiente habilitar los indicadores necesarios en las configuraciones de MaxText y asegurar que el clúster tenga chips Blackwell instalados (B100, B200, GB200).

"La precisión numérica es uno de los parámetros más apalancados, pero el preentrenamiento mixed-precision de bajo bit es difícil de implementar correctamente," señala el equipo del NVIDIA

Developer Blog.

Qué Significa Esto

Para equipos dedicados al preentrenamiento de modelos de frontera, NVFP4 en Blackwell es prácticamente aceleración gratuita: el stack existente en JAX y MaxText requiere cambios de configuración mínimos. A escala de cientos y miles de GPUs, incluso ganancias de throughput de 10–15% reducen directamente el tiempo hasta checkpoint y el presupuesto computacional general. La carrera por la eficiencia del preentrenamiento está entrando en la fase de batalla por la precisión numérica.