Resumen Técnico: Replanteamiento del Sesgo de Contracción en el Preentrenamiento de LLM en FP4

Declaración del Problema

A medida que los Modelos de Lenguaje Grande (LLM) escalan, los costos computacionales y de memoria del preentrenamiento requieren una transición hacia formatos de menor precisión. Aunque el formato FP8 ha sido ampliamente adoptado, el hardware emergente (por ejemplo, NVIDIA Blackwell/Rubin, AMD MI350) soporta computación nativa en FP4. Las recetas actuales de entrenamiento en FP4, como NVFP4, dependen predominantemente del formato E2M1 (2 bits de exponente, 1 bit de mantisa) combinado con la Transformada de Hadamard Aleatoria (RHT) para mitigar la sensibilidad a los valores atípicos (outliers).

Sin embargo, estas recetas basadas en E2M1 sufren de inestabilidad en el entrenamiento y una degradación significativa de la pérdida en comparación con las bases de BF16. El artículo identifica una limitación fundamental, previamente pasada por alto: el Sesgo de Contracción (Shrinkage Bias). Este es un error de redondeo negativo sistemático inherente a los formatos de cuantización no uniformes como E2M1. La asimetría geométrica de los bins de E2M1 causa una "contracción" sistemática de las magnitudes bajo el redondeo de Redondeo al Más Cercano Par (Round-to-Nearest-Even o RTNE). Crucialmente, el artículo argumenta que, si bien la RHT está diseñada para dispersar los valores atípicos y mejorar la utilización del código, inadvertidamente exacerba este sesgo en E2M1 al desplazar la masa de los tensores hacia los bins de redondeo más asimétricos del formato, lo que provoca una degradación multiplicativa de la señal a través de redes profundas.

Metodología

1. Análisis Teórico del Sesgo de Contracción

Los autores formalizan el Sesgo de Contracción como un error esperado negativo en el espacio de magnitud normalizado causado por la asimetría geométrica de los bins de RTNE.

• Origen Geométrico: En rejillas no uniformes como E2M1, el espaciamiento entre los valores representables cambia (por ejemplo, de 0.5 a 1.0 a 1.5). Para un nivel $q_i$, los anchos de los bins izquierdo y derecho ($\ell_i$ y $r_i$) suelen ser desiguales. Si $r_i > \ell_i$, el error esperado condicional bajo una densidad local uniforme es negativo ($\frac{\ell_i - r_i}{2} < 0$).
• Impacto Sistémico: A diferencia del ruido de media cero que se cancela, este sesgo actúa como un factor de atenuación sistemático ($\alpha < 1$) en las multiplicaciones de matrices (GEMMs). A través de $K$ capas, estos factores se acumulan multiplicativamente ($\prod \alpha_k \approx \exp(-\sum \delta_k)$), causando una degradación sistemática de la señal.
• Interacción con RHT: La RHT transforma los tensores de ser "limitados por el rango dinámico" a ser "limitados por la resolución local". Si bien esto mejora la utilización de los buckets, fuerza los datos hacia las regiones de magnitud media donde la asimetría de E2M1 es más pronunciada, amplificando así el sesgo de contracción.

2. La Receta UFP4

Basándose en el hallazgo de que las rejillas uniformes evitan este error geomético, los autores proponen UFP4 (Uniform FP4), una receta de entrenamiento construida sobre una rejilla uniforme de tipo E1M2/INT4.

• Rejilla Uniforme: A diferencia de E2M1, la rejilla E1M2 mantiene un espaciamiento igual entre niveles ($\ell_i = r_i$), eliminando la fuente geométrica del Sesgo de Contracción.
• Cobertura Total de RHT: Debido a que la rejilla uniforme no tiene sesgo, UFP4 aplica RHT de forma segura a las tres rutas de GEMM de entrenamiento: Propagación hacia adelante (FPROP), Gradiente de Datos (DGRAD) y Gradiente de Pesos (WGRAD). Las recetas anteriores de E2M1 restringían la RHT a la ruta del gradiente de pesos para evitar la acumulación de errores en rutas que no son hojas (non-leaf paths).
• Redondeo Estocástico (SR): El SR se aplica únicamente al gradiente ascendente ($dY$) para preservar las expectativas del gradiente, mientras que se utiliza RTNE en otros lugares.

Contribuciones Clave

1. Identificación del Sesgo de Contracción: El artículo define y demuestra formalmente que los formatos no uniformes (E2M1) sufren inherentemente de un error de redondeo negativo sistemático debido a la asimetría geométrica, distinto de los errores de cuantización inducidos por la distribución.
2. Explicación Unificada de la Inestabilidad: Proporciona una explicación teórica y empírica para la inestabilidad del entrenamiento en las recetas existentes basadas en E2M1, mostrando cómo el Sesgo de Contracción se acumula multiplicativamente y es exacerbado por la RHT.
3. Propuesta de UFP4: Una nueva receta de entrenamiento en 4 bits que utiliza rejillas uniformes (E1M2/INT4) que permite una RHT omnipresente en todos los GEMM de entrenamiento, convirtiendo la mejora en la utilización de los buckets en una mayor calidad de cuantización.
4. Validación a Escala Industrial: El enfoque se valida en modelos Dense 1.5B, MoE 7.9B y MoE 124B, demostrando que las rejillas uniformes son formatos prácticos de primer nivel para el entrenamiento en FP4.

Resultados Experimentales

Diagnósticos de Cuantización y GEMM

• Efecto de RHT: En tensores con alta presencia de valores atípicos, la RHT mejora significamente la relación de bucket efectiva para ambos formatos. Sin embargo, la Relación Señal-Ruido de Cuantización (SQNR) resultante diverge: E1M2 muestra una SQNR mejorada post-RHT, mientras que E2M1 a menudo se degrada ( $\Delta$SQNR negativa) porque los datos desplazados caen en bins asimétricos.
• Salidas de GEMM: En pruebas de un solo GEMM, E1M2 preserva o mejora la SQNR de salida a través de las rutas de propagación hacia adelante y hacia atrás con RHT, mientras que E2M1 sufre degradación, particularmente en las rutas hacia atrás (backward paths).

Preentrenamiento de Extremo a Extremo

• Degradación de la Pérdida: En tareas de preentrenamiento de larga duración, UFP4 (E1M2) logra consistentemente una menor degradación de pérdida relativa respecto a BF16 en comparación con las bases fuertes de E2M1.
  • Dense 1.5B: El error relativo se redujo de 1.2570% (E2M1) a 0.9673% (UFP4).
  • MoE 7.9B: Se redujo de 2.3596% a 1.8469%.
  • MoE 124B: Se redujo de 1.7308% a 1.3863%.
• Leyes de Escalamiento: La ventaja de UFP4 persiste a través de diversas escalas de modelos (de 10M a 324M de parámetros). Las curvas de pérdida ajustadas muestran que la brecha de E1M2 respecto a BF16 sigue siendo menor que la brecha de E2M1, y la penalización de FP4 no crece con la escala de cómputo.
• Estudios de Ablación:
  • Alcance de RHT: La cobertura total de RHT (en los tres GEMM) produce los mejores resultados para UFP4, mientras que restringir la RHT (como se hace en las recetas de E2M1) limita el rendimiento.
  • Restricción de Rango: Los intentos de emular una rejilla uniforme restringiendo el rango dinámico de E2M1 (por ejemplo, limitando el valor máximo) no lograron igualar el rendimiento de la E1M2 nativa, sacrificando el rango dinámico sin eliminar completamente el sesgo.

Eficiencia

• Fusión de Kernels: Los autores demuestran que la RHT puede fusionarse con la cuantización FP4 de manera eficiente. En el hardware probado (SM90/SM100), la latencia de la fusión RHT+cuantización es solo ~1.06x–1.07x la de la cuantización independiente, comparado con ~1.4x–1.6x para implementaciones no fusionadas.

Significado y Reivindicaciones

El artículo argumenta que el estándar actual de la industria de tratar a E2M1 como el único formato viable de FP4 es subóptimo debido a sus limitaciones geométricas. El significado de este trabajo radica en:

1. Desafiar el Statu Quo: Sugiere que los futuros aceleradores de ML deberían soportar rejillas uniformes de tipo E1M2/INT4 como primitivas de entrenamiento de primer nivel junto con E2M1, en lugar de depender únicamente de E2M1.
2. Desajuste de Régimen de Formato: Destaca que la efectividad de las técnicas de mitigación de valores atípicos como la RHT depende de la geometría de la rejilla subyacente. La RHT es beneficiosa para las rejillas uniformes pero perjudicial para las rejillas no uniformes en regímenes de entrenamiento profundos.
3. Receta Práctica: UFP4 ofrece una receta concreta y escalable que reduce la brecha entre el entrenamiento en FP4 y BF16 sin requerir modificaciones complejas de estimadores, simplemente alineando la geometría de la rejilla con la distribución del tensor post-RHT.

Los autores concluyen que, si bien E2M1 sigue siendo adecuado para cargas de trabajo limitadas en rango e inferencia, el hardware de entrenamiento debe evolucionar para soportar rejillas uniformes para alcanzar plenamente el potencial del preentrenamiento en FP4.