Microsoft Research lanza el primer marco de entrenamiento de modelos grandes con precisión FP4, aumentando significativamente la eficiencia

El 29 de enero, Microsoft Research anunció el lanzamiento de su primer marco de entrenamiento de modelos grandes con precisión FP4. Bajo las mismas configuraciones de hiperparámetros, este marco puede lograr resultados de entrenamiento equivalentes a FP8 y BF16, lo que significa que se pueden reducir los recursos de almacenamiento y computación necesarios.
Tamaño del modelo y efectividad del entrenamiento
Los modelos entrenados con este método pueden alcanzar un tamaño de hasta 13 mil millones de parámetros, con una cantidad de tokens de entrenamiento que llega a cientos de miles de millones. Aunque actualmente se utiliza FP4 simulado a través de FP8, los resultados mejorarían aún más si se utilizara FP4 real.
Simulación de FP4 y método de entrenamiento
Dado que no existían hardware nativo compatible con FP4 en el momento del estudio, el equipo de investigación simuló FP4 utilizando TensorCore FP8. En particular, en los modelos LLaMA de 1.3B, 7B y 13B, durante el entrenamiento de 0 a 10 millones de tokens, la curva de pérdida de FP4 era casi idéntica a la de BF16.
Para realizar el entrenamiento a precisión FP4, el equipo utilizó un núcleo CUDA personalizado de multiplicación de matrices FP4 (GeMM). En este núcleo, primero se cargan las matrices A y B de FP4 en memoria compartida usando FP16 y se transforman adecuadamente, luego se realiza la multiplicación de matrices en bloques utilizando FP4, y finalmente, se reduce el resultado intermedio con FP16 para obtener una matriz de salida en formato FP16.
Estrategias de cuantificación y adaptación al hardware
Este marco utiliza el formato FP4 E2M1, en el que 2 bits representan el exponente, 1 bit la mantisa, y 1 bit el signo, para un total de 4 bits. Este formato se seleccionó para adaptarse al diseño de las unidades de cálculo cuantificadas de los principales aceleradores de ML actuales. Las matrices de pesos W y las matrices de activación A tienen estrategias de cuantificación de diferentes granularidades, lo que maximiza el efecto de aceleración de FP4 en la multiplicación de matrices.
Innovaciones en la propagación hacia adelante y hacia atrás
Al comenzar la propagación hacia adelante del modelo, el marco realiza una cuantificación FP4 simultánea de las matrices de peso W y las matrices de activación de entrada A en cada capa lineal. Durante la cuantificación, los valores de las matrices se escalan y sesgan para mapearlos al rango que FP4 puede representar, y luego se redondean a los valores discretos FP4 más cercanos mediante una tabla de búsqueda.
En el proceso de retropropagación, el equipo propuso un novedoso método de estimación de gradientes diferenciables para garantizar la eficiencia computacional y la precisión del gradiente. Además, para abordar el problema de los "valores atípicos" en las matrices de activación, el equipo introdujo una estrategia de "recorte y compensación de picos".
Diseño y aplicación de precisión mixta
En algunas fases, el equipo utilizó un diseño de precisión mixta, como el uso de FP8 para la comunicación de gradientes y FP16 para el almacenamiento del estado del optimizador. En otras partes del sistema, como los cálculos fuera de las multiplicaciones de matrices, y el escalado de la pérdida, también se utilizó FP16. Estos diseños reducen al mínimo el gasto computacional y de almacenamiento sin comprometer la estabilidad numérica del entrenamiento.
Equipo de investigación y contexto
Este marco fue desarrollado por el Microsoft Research Asia y el equipo SIGMA, todos los investigadores son de origen chino. El primer autor, Ruizhe Wang, es estudiante de doctorado en la Universidad de Ciencia y Tecnología de China y actualmente realiza prácticas en MSRA, con investigación en cuantificación de baja precisión. El profesor Zhejun Cha, Director del Departamento de Investigación de USTC, también participó en el proyecto. Los autores de correspondencia fueron Peng Cheng, Senior Research Manager de MSRA, y Yeyun Gong, Chief Research Manager. El científico destacado de MSRA, Baining Guo, también participó en el proyecto.
Resumen
El marco de entrenamiento de modelos grandes con precisión FP4 de Microsoft Research, mediante estrategias innovadoras de cuantificación y diseño de precisión mixta, ha logrado un entrenamiento de modelos más eficiente. Este marco no solo mejora la eficiencia de entrenamiento, sino que también reduce significativamente la demanda de recursos computacionales y de almacenamiento, abriendo nuevas vías para el entrenamiento de modelos grandes.
Enlace al artículo: arXiv
Enlace de referencia: Twitter