Scalify : propagation d'échelle pour un entraînement efficace de modèles de langage en basse précision

Scalify: scale propagation for efficient low-precision LLM training

July 24, 2024

Auteurs: Paul Balança, Sam Hosegood, Carlo Luschi, Andrew Fitzgibbon

cs.AI

Résumé

Les formats de faible précision tels que le float8 ont été introduits dans les matériels accélérés pour l'apprentissage automatique afin d'améliorer l'efficacité computationnelle lors de l'entraînement et de l'inférence des grands modèles de langage. Cependant, leur adoption par la communauté du ML a été ralentie par les techniques complexes, et parfois fragiles, nécessaires pour atteindre la précision d'entraînement des formats de plus haute précision. Dans ce travail, nous présentons Scalify, un paradigme de propagation d'échelle de bout en bout pour les graphes de calcul, généralisant et formalisant les méthodes existantes de mise à l'échelle des tenseurs. Les résultats expérimentaux montrent que Scalify prend en charge nativement la multiplication matricielle en float8 et la représentation des gradients, ainsi que le stockage des états de l'optimiseur en float16. Notre implémentation de Scalify dans JAX est open-source et disponible à l'adresse suivante : https://github.com/graphcore-research/jax-scalify.

English

Low-precision formats such as float8 have been introduced in machine learning accelerated hardware to improve computational efficiency for large language models training and inference. Nevertheless, adoption by the ML community has been slowed down by the complex, and sometimes brittle, techniques required to match higher precision training accuracy. In this work, we present Scalify, a end-to-end scale propagation paradigm for computational graphs, generalizing and formalizing existing tensor scaling methods. Experiment results show that Scalify supports out-of-the-box float8 matrix multiplication and gradients representation, as well as float16 optimizer state storage. Our JAX implementation of Scalify is open-sourced at https://github.com/graphcore-research/jax-scalify