FP8 y NVFP4 en Transformers: hasta 3,48x con Blackwell B300

Las arquitecturas Transformer son la columna vertebral de la mayoría de los modelos de lenguaje y de IA generativa actuales. A medida que crecen, los entrenamientos consumen más horas de GPU y más tiempo de iteración para los equipos de ingeniería. Acelerarlos no es solo una optimización: define cuán rápido se puede experimentar y qué tan grande es el modelo que se puede entrenar dentro del presupuesto. Las GPUs NVIDIA Hopper y NVIDIA Blackwell atacan ese problema con soporte para operadores de baja precisión como FP8 y NVFP4.

El grueso del entrenamiento se va en operaciones GEMM (multiplicaciones de matrices), y las precisiones bajas aceleran esas multiplicaciones haciéndolas más rápidas y baratas. El detalle: la configuración de tu Transformer no te dice qué GEMMs corren realmente en tu modelo. Para entender dónde se va el tiempo hay que traducir la config y el batch size a las formas exactas M×K×N que ejecuta el modelo, y luego medirlas en cada precisión. Solo así se puede decidir la precisión óptima antes de comprometerse con un entrenamiento caro.

NVIDIA Transformer Engine (TE) se encarga de la cuantización y el despacho de kernels para abrir esos formatos de baja precisión. La guía publicada por NVIDIA muestra cómo pasar de la configuración del modelo a workloads concretos de GEMM, medirlos con un microbenchmark y estimar dónde la precisión menor se traduce en aceleración real. El caso de uso es CodonFM 5B, un modelo de lenguaje para biología enfocado en ARN.

¿Cuál es la configuración base del benchmark?

Tomemos un modelo de 5.000 millones de parámetros como CodonFM 5B, con la siguiente configuración:

hidden_size: 4096
intermediate_size: 16384
num_attention_heads: 32
num_hidden_layers: 24

Y un setup de entrenamiento:

micro_batch_size: 31
sequence_length: 512

El benchmark toma esos hiperparámetros y, con un solo comando, deriva las formas GEMM, las mide en cada precisión y calcula la aceleración total:

python benchmark.py \
  --hidden_size 4096 \
  --intermediate_size 16384 \
  --num_attention_heads 32 \
  --num_hidden_layers 24 \
  --micro_batch_size 31 \
  --sequence_length 512 \
  -o ./images/b300_model_config_speedup.png

Si se quiere desactivar los flags específicos de Blackwell, los flags --no-fp8 --no-fp4 dejan BF16 más las tres recetas FP8 tensor-wise que funcionan en Hopper. --no-fp8 desactiva MXFP8 y --no-fp4 desactiva NVFP4.

Autocast versus prequantize: dos formas de medir

Por defecto la herramienta corre en modo autocast, que es lo que hace TE durante el entrenamiento: los inputs se cuantizan dinámicamente a la precisión objetivo antes de cada GEMM, por lo que el tiempo medido incluye tanto el costo de cuantización como el kernel GEMM en sí. Esto entrega la foto realista por GEMM dentro de un paso de entrenamiento.

La herramienta calcula M = 31 × 512 = 15.872 tokens, deriva las 12 formas GEMM, mide cada una en las precisiones habilitadas e imprime los resultados completos. Fprop, Dgrad y Wgrad se miden por separado para capturar el impacto de los distintos aspect ratios sobre la selección de kernel.

Para aislar el rendimiento puro del kernel GEMM, hay que agregar --pre-quantize. Esto cuantiza todos los inputs una sola vez antes del loop medido, así el tiempo refleja únicamente la ejecución del kernel: sin cuantización dinámica, sin cálculo de block scaling, sin conversión de formato durante la región medida.

Atención: FP8 DelayedScaling siempre corre en autocast, incluso con --pre-quantize, porque depende de un historial amax que requiere cuantización dinámica. Sus tiempos no son directamente comparables con los de las otras precisiones en modo prequantized.

python benchmark.py \
  --hidden_size 4096 \
  --intermediate_size 16384 \
  --num_attention_heads 32 \
  --num_hidden_layers 24 \
  --micro_batch_size 31 \
  --sequence_length 512 \
  --pre-quantize \
  -o ./images/b300_model_config_speedup_prequant.png

Comparar las aceleraciones de autocast y prequantized dice exactamente cuánto cuesta el overhead de cuantización: NVFP4 frente a BF16 pasa de 1,98x (autocast) a 3,48x (solo kernel). La brecha entre los dos números es el costo de la cuantización dinámica, los transformes Hadamard y el block scaling que ocurren en cada paso de entrenamiento.

Los resultados de autocast sirven para predecir aceleraciones reales de entrenamiento (es lo que hace TE en producción). Los resultados prequantized sirven para entender si el cuello de botella es el overhead de cuantización, o para comparar el throughput puro de los tensor cores entre precisiones, independiente de la implementación de la cuantización.

¿Qué se ve al aterrizarlo sobre un modelo real?

Sobre la misma configuración de CodonFM 5B, NVIDIA corrió el benchmark completo en una B300. Las aceleraciones por forma de NVFP4 frente a MXFP8 en los resultados de Fprop son:

QKV proj:   0,579 / 0,392  =  1,48x
Attn out:   0,269 / 0,256  =  1,05x  (apenas más rápido — el overhead casi iguala la ganancia GEMM)
MLP up:     0,924 / 0,635  =  1,46x
MLP down:   1,076 / 0,649  =  1,66x

Hay varios puntos relevantes:

• El GEMM de salida de atención casi no se beneficia de la precisión baja. Frente al baseline MXFP8 hay solo un 1,05x. Es la matriz de pesos más pequeña de la capa (4096×4096), apenas lo suficientemente grande como para que la precisión baja compense el overhead. En cambio, el GEMM mucho más grande de MLP Down rinde 1,66x de NVFP4 sobre MXFP8 en el mismo hardware: el GEMM de MLP down es lo bastante grande como para amortizar el costo de cuantización, el de attention output no.
• Los GEMMs grandes muestran ganancias reales pero subteóricas. Los tensor cores FP4 entregan entre 1,46x y 1,66x sobre MXFP8 en los GEMMs grandes. Está bastante por debajo del 2x a 3x teórico del spec de hardware. Si se incluye el GEMM de attention output, la aceleración Fprop combinada baja a 1,47x. Al sumar los tiempos de Wgrad, el overhead no-GEMM y los costos de cuantización propios de NVFP4, la diferencia end-to-end entre NVFP4 y MXFP8 en entrenamiento es consistente con estos números a nivel de kernel.
• FP8 DelayedScaling es sorprendentemente competitivo en Blackwell. A 7,80 ms/capa en autocast supera tanto a FP8 CurrentScaling (9,15 ms) como a MXFP8 (8,98 ms). En modo prequantized FP8 CurrentScaling pasa adelante (6,81 ms frente a 8,12 ms), lo que sugiere que el enfoque de historial amax de DelayedScaling tiene menor overhead de cuantización pero throughput puro de kernel similar. Es un buen ejemplo de cómo la comparación entre autocast y prequantized hace aparecer ganadores distintos según se mida con o sin el impuesto de la cuantización.
• Los resultados prequantized revelan el verdadero potencial del kernel. Correr con --pre-quantize elimina por completo el overhead de cuantización: NVFP4 frente a BF16 salta de 1,98x (autocast) a 3,48x (solo kernel). Eso muestra que los tensor cores FP4 sí están entregando aceleraciones reales: es el overhead de cuantización en autocast el que cierra la brecha.
• La comparación Fprop versus Dgrad muestra que la aproximación del 2x es imprecisa para los formatos cuantizados. Mientras BF16 Dgrad queda dentro del 2% de Fprop, los formatos cuantizados muestran sumas Dgrad 5-13% más lentas. El Dgrad de QKV Proj es especialmente asimétrico (33-51% más lento que Fprop para FP8/FP4) porque intercambiar K (4096) y N (12288) cambia dramáticamente el aspect ratio de la matriz y la selección de kernel. Por eso la herramienta mide Fprop y Dgrad por separado en vez de contar el tiempo de Fprop dos veces.

¿Cómo se compara contra el speedup real del entrenamiento?

Una vez que se tiene la aceleración estimada de los GEMMs, hay que compararla con la aceleración end-to-end observada:

• GEMM speedup ≈ training speedup: los GEMMs dominan el paso, todo funciona como se espera.
• GEMM speedup >> training speedup: el overhead fuera de los GEMMs se está comiendo las ganancias. Para NVFP4 en particular ese overhead incluye los transformes Random Hadamard en los inputs de Wgrad, stochastic rounding en gradientes, block scaling 2D para los pesos y la pasada extra de memoria para calcular amax per-tensor. Son operaciones adicionales que MXFP8 no necesita y pueden cerrar significativamente la brecha aunque los GEMMs FP4 puros sean mucho más rápidos.
• GEMM speedup ≈ 1,0 incluso en el microbenchmark: los kernels FP4 no son más rápidos en esas formas, o están cayendo silenciosamente a FP8.

El último caso conviene chequearlo de verdad. Se puede setear NVTE_LOG_LEVEL=1 o inspeccionar con NVIDIA Nsight Systems para confirmar que TE está despachando kernels FP4. TE puede caer silenciosamente a FP8 o BF16 para capas u operaciones que aún no soportan FP4, lo que explicaría rendimientos idénticos sin otros síntomas. Otra opción es comparar el uso de memoria GPU entre corridas MXFP8 y NVFP4: si la memoria es casi igual, es una señal fuerte de que los pesos FP4 no se están almacenando realmente.

Comparativa rápida: precisiones en entrenamiento de LLM (2026)

Implicancias para equipos en Chile y LatAm

Para los equipos de la región que arriendan tiempo de GPU en proveedores hyperscale o entrenan en clusters chicos, el dato práctico es claro: subirse a Blackwell no garantiza el 3x del marketing. La diferencia entre el speedup de marketing (kernel-only) y el observado en producción (autocast) puede ser de hasta 1,5x. Hay que medir antes de comprometer presupuesto en arriendo o en compra. Una hora de B300 en hyperscaler suele rondar los USD 30-60/h, así que una mala estimación de aceleración convierte un entrenamiento previsto en USD 10.000 en uno de USD 17.000.

La buena noticia: el benchmark de NVIDIA es open source, se corre con un comando y entrega los números reales antes de meter el modelo en el cluster. Para quienes entrenan modelos pequeños (4096×4096 y menos) la conclusión es más dura todavía: el overhead de NVFP4 puede neutralizar la ganancia del tensor core, y FP8 sigue siendo la opción más segura.