Scalify: propagação de escala para treinamento eficiente de LLM de baixa precisão

Scalify: scale propagation for efficient low-precision LLM training

July 24, 2024

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

cs.AI

Resumo

Formatos de baixa precisão, como float8, foram introduzidos em hardware acelerado de aprendizado de máquina para melhorar a eficiência computacional no treinamento e inferência de grandes modelos de linguagem. No entanto, a adoção pela comunidade de ML tem sido retardada pelas técnicas complexas e, por vezes, frágeis necessárias para equiparar a precisão de treinamento de alta precisão. Neste trabalho, apresentamos o Scalify, um paradigma de propagação de escala de ponta a ponta para grafos computacionais, generalizando e formalizando os métodos de escalonamento de tensores existentes. Os resultados dos experimentos mostram que o Scalify suporta multiplicação de matrizes float8 e representação de gradientes prontas para uso, bem como armazenamento de estado do otimizador float16. Nossa implementação do Scalify em JAX está disponível em código aberto em 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