Scalify: scale propagation for efficient low-precision LLM training

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Balança, Paul, Hosegood, Sam, Luschi, Carlo, Fitzgibbon, Andrew
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866914885386895360
author Balança, Paul
Hosegood, Sam
Luschi, Carlo
Fitzgibbon, Andrew
author_facet Balança, Paul
Hosegood, Sam
Luschi, Carlo
Fitzgibbon, Andrew
contents 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
format Preprint
id arxiv_https___arxiv_org_abs_2407_17353
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Scalify: scale propagation for efficient low-precision LLM training
Balança, Paul
Hosegood, Sam
Luschi, Carlo
Fitzgibbon, Andrew
Machine Learning
68T07
I.2.7
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
title Scalify: scale propagation for efficient low-precision LLM training
topic Machine Learning
68T07
I.2.7
url https://arxiv.org/abs/2407.17353