Scalify: scale propagation for efficient low-precision LLM training
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
66a1b2c604ab5290ec0d9625
Paul Balança
admin_assigned
2024-07-25T09:08:50.489Z
60d38165d7b174177faabd67
/avatars/809b2a3fc3502eefcd48aa70559f8275.svg
Paul Balanca
user
balancap
66a1b2c604ab5290ec0d9626
Sam Hosegood
admin_assigned
2024-07-25T09:08:56.499Z
65776369bc075aaecd83c517
/avatars/ff9e69a9a8a4043fd761ee69da80def9.svg
Sam Hosegood
user
slhosegood
66a1b2c604ab5290ec0d9627
Carlo Luschi
66a1b2c604ab5290ec0d9628
Andrew Fitzgibbon
admin_assigned
2024-07-25T09:09:09.934Z
6360088814657fb8cff4a88a
/avatars/abc7f05b6c2d1bd97a1d6d2390209601.svg
Andrew Fitzgibbon
user
awf
Scalify: scale propagation for efficient low-precision LLM training
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
Scalify is an end-to-end scale propagation paradigm that supports float8 matrix multiplication and float16 optimizer state storage, improving low-precision training in large language models.
float8 low-precision formats computational efficiency large language models tensor scaling methods Scalify matrix multiplication gradients float16 optimizer state storage JAX