COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training
Extracted source text · arXiv 2410.19313
\documentclass{article} % For LaTeX2e
\usepackage{iclr2025_conference,times}
\newcommand{\SH}[1]{{\bf\color{red} [Song: #1]}}
\newcommand{\JS}[1]{{\bf\color{blue} [Jason: #1]}}
% Optional math commands from https://github.com/goodfeli/dlbook_notation.
\input{math_commands.tex}
\usepackage{url}
\usepackage{graphicx} % Required for inserting images
\newtheorem{definition}{Definition}
\usepackage{amsmath}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{subfigure}
\usepackage{booktabs}
\usepackage{caption}
\usepackage{multirow}
\usepackage{arydshln}
\usepackage{xcolor}
\usepackage{colortbl}
\usepackage{makecell}
\usepackage{bigstrut}
\usepackage{pifont}
\definecolor{darkgreen}{rgb}{0.15, 0.75, 0.15}
\definecolor{cvprblue}{rgb}{0.21,0.49,0.74}
\definecolor{lightblue}{rgb}{0.90, 0.95, 0.99}
\colorlet{transparentlightblue50}{lightblue!50}
% \definespotcolor[lightblueLightestSpot][lightblue][a=1,t=.3]
% \usepackage[pagebackref,breaklinks,colorlinks,citecolor=cvprblue]{hyperref}
\usepackage{hyperref}
\newcommand{\jianfei}[1]{\textcolor{red}{[Jianfei: #1]}}
\title{COAT: \underline{C}ompressing \underline{O}ptimizer states and \underline{A}ctivation for Memory-Efficient FP8 \underline{T}raining}
% Authors must not appear in the submitted version. They should be hidden
% as long as the \iclrfinalcopy macro remains commented out below.
% Non-anonymous submissions will be rejected without review.
\author{
\hspace{-10pt}
Haocheng Xi$^{1}$\thanks{Part of the work done during an internship at NVIDIA.}, Han Cai$^{2}$, Ligeng Zhu$^{2}$, Yao Lu$^{2}$, Kurt Keutzer$^{1}$, Jianfei Chen$^{4}$, Song Han$^{2,3}$ \\
\\ $^1$ University of California, Berkeley \quad $^2$ NVIDIA \quad $^3$ MIT \quad $^4$ Tsinghua University \\ \texttt{\url{https://github.com/NVlabs/COAT} \& \url{https://nvlabs.github.io/COAT/}}
}
% \author{Antiquus S.~Hippocampus, Natalia Cerebro \& Amelie P. Amygdale \\
% Department of Computer Science\\
% Cranberry-Lemon University\\
% Pittsburgh, PA 15213, USA \\
% \texttt{\{hippo,brain,jen\}@cs.cranberry-lemon.edu} \\
% \And
% Ji Q. Ren \& Yevgeny LeNet \\
% Department of Computational Neuroscience \\
% University of the Witwatersrand \\
% Joburg, South Africa \\
% \texttt{\{robot,net\}@wits.ac.za} \\
% \AND
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email}
% }
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to \LaTeX{} to determine where to break
% the lines. Using \AND forces a linebreak at that point. So, if \LaTeX{}
% puts 3 of 4 authors names on the first line, and the last on the second
% line, try using \AND instead of \And before the third author name.
\newcommand{\fix}{\marginpar{FIX}}
\newcommand{\new}{\marginpar{NEW}}
\iclrfinalcopy % Uncomment for camera-ready version, but NOT for submission.
\begin{document}
\maketitle
\input{text/0_abstract}
\vspace{-15pt}
\input{text/1_introduction}
\input{text/2_related_work}
\input{text/3_optimizer_state}
\input{text/4_activation}
\input{text/5_experiment}
\input{text/100_conclusion}
\newpage
\bibliography{iclr2025_conference}
\bibliographystyle{iclr2025_conference}
\newpage
\appendix
\input{text/10_Appendix}
\end{document}
%%%%% NEW MATH DEFINITIONS %%%%%
\usepackage{amsmath,amsfonts,bm}
% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}
% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}
% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}
\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}
\def\eps{{\epsilon}}
% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}
% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}
% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}
% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}
% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}
% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}
% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}
% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}
% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}
% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}
% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}
% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}
% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}
% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}
\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution
\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}
\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak
\begin{abstract}
FP8 training has emerged as a promising method for improving training efficiency. Existing frameworks accelerate training by applying FP8 computation to linear layers while leaving optimizer states and activations in higher precision, which fails to fully optimize memory usage. This paper introduces COAT (\textbf{C}ompressing \textbf{O}ptimizer States and \textbf{A}ctivations for FP8 \textbf{T}raining), a novel FP8 training framework designed to significantly reduce memory footprint when training large models. COAT addresses current limitations through two key innovations: (1) \textbf{Dynamic Range Expansion}, which aligns optimizer state distributions more closely with the FP8 representation range, thereby reducing quantization error, and (2) \textbf{Mixed-Granularity Activation Quantization}, which optimizes activation memory using a combination of per-tensor and per-group quantization strategies. Experiments demonstrate that COAT effectively reduces end-to-end training memory footprint by 1.54× compared to BF16 while achieving nearly lossless performance across various tasks, such as Large Language Model pretraining and fine-tuning and Vision Language Model training. COAT also achieves a 1.43× end-to-end training speedup compared to BF16, performing on par with or surpassing TransformerEngine's speedup. COAT enables efficient full-parameter training of large models on fewer GPUs, and facilitates doubling the batch size in distributed training settings, providing a practical solution for scaling large-scale model training. The code is available at \href{https://github.com/NVlabs/COAT}{\text{https://github.com/NVlabs/COAT}}.
\end{abstract}
\begin{figure}[h!]
\begin{minipage}{\textwidth}
\centering
\includegraphics[width=0.95\textwidth]{figure/FP8PrecisionFlow8.pdf}
\captionof{figure}{(a,b) Comparing the quantization flow of Transformer Engine and COAT. Both the optimizer states and activations are quantized to FP8 in COAT. (c) End-to-end per-GPU memory comparison when training Llama-2-13B on 8$\times80$G H100 using FSDP.}\label{fig:overview_of_flow_te_and_bar}
\end{minipage}
\end{figure}
\section{Introduction}
Foundation Models (FMs), such as Large Language Models (LLM) and Vision Language Models (VLM), have made significant breakthroughs in various tasks such as reasoning, understanding, and summarization~\citep{dubey2024llama3, adler2024nemotron, team2024gemma,lin2024vila}. However, the training of such models, which often comprise billions of parameters, demands substantial computational resources and memory. This presents substantial challenges, making the training of these foundation models very challenging~\citep{smith2022nlg530b,hoffmann2022chinchillia}.
Low-precision training has emerged as a promising approach to make FMs training more efficient~\citep{micikevicius2017mixedprecision,wang2018trainingfp8,zhu2020towardsint8,xi2023int4training,wortsman2023int8clip,xi2024jetfire}. By quantizing tensors used in deep neural networks into lower precision, low-precision training effectively speed up the training process and reduce the memory footprint. Currently, BF16 training ~\citep{kalamkar2019bf16study,micikevicius2017mixedprecision} is the most prevalent low-precision method, and is widely adopted in large-scale training frameworks like DeepSpeed ~\citep{rasley2020deepspeed} and Megatron-LM~\citep{shoeybi2019megatron}.
With the advent of Nvidia's H100 GPU~\citep{nvidia_h100}, FP8 training~\citet{micikevicius2022fp8format} is emerging as the next-generation low-precision technique. Compared to BF16, FP8 training has the potential to (1) double the speed and (2) halve the memory footprint. To achieve practical speedup,
Transformer Engine~\citep{nvidia_transformerengine} performs matrix multiplications in FP8 precision, leading to faster training. Transformer Engine’s memory footprint can be further improved by reducing optimizer states, gradients, weights, and activations to lower precision. As illustrated in Figure~\ref{fig:overview_of_flow_te_and_bar}, FP8-LM~\citep{peng2023fp8lm} advances this by further quantizing the gradients, weight master copy, and first-order momentum into FP8. This reduces memory and communication overhead, partially improving memory efficiency. However, they do not tackle the memory consumption of \textit{activations} and still leave the optimizer's second-order momentum in higher precision. The memory problem of activations becomes even more critical when optimizer, gradient, and weights are sharded across multiple GPUs using ZeRO or FSDP. Besides, second-order momentum is more sensitive to quantization than first-order momentum~\citep{fishman2024fp8trillion}, and activations' large spikes also make them hard to quantize to FP8~\citep{yang2024gluspike}. This potential accuracy degradation makes them missing a crucial opportunity to optimize memory further.
% \jianfei{maybe emphasize that there are two steps: 1. to compress computation to 8-bit; 2. to compress memory address/storage to 8-bit, and point out that TE does not save memory. Then, we should explain why memory efficiency is still important, particularly given the large volume of zero and 3d parallel systems.
% The positioning of our paper might be: (1) TE achieves the first step; (2) FP8-LM partially achieves the second step; (3) We further achieve the second step. }
% Despite these advancements, most existing methods overlook the memory consumption caused by activations, and the quantization of the optimizer's second-order momentum remains a challenge.
In this work, we propose COAT: \textbf{C}ompressing \textbf{O}ptimizer states and \textbf{A}ctivations for memory-efficient FP8 \textbf{T}raining to address the aforementioned issue. \emph{COAT significantly reduces the overall memory footprint by quantizing optimizer states and activations into FP8.} For optimizer states, we observe that FP8 format's \emph{representation range is under-utilized} when quantizing them, as illustrated in Figure~\ref{fig:optimizer_range_ratio}(a). To address this, we introduce a novel \emph{Dynamic Range Expansion} method which adjusts the distribution of optimizer states to better fit within the FP8 range, thereby minimizing quantization error. For activations, we propose \emph{Mixed-Granularity Activation Quantization} to achieve efficient and accurate quantization. We apply fine-grained quantization to non-linear layers and apply per-tensor quantization to linear layers. Per-tensor quantization for matrix multiplications is more efficient and better suited for TensorCores, while fine-grained quantization helps maintain accuracy. These two approaches tackle high memory consumption while ensuring minimal performance degradation. We provide an overview of COAT in Figure~\ref{fig:overview_of_flow_te_and_bar}(b) for demonstration.
We demonstrate the accurate performance of COAT on a wide range of tasks, including LLM pretraining, LLM fine-tuning, and VLM training. COAT achieves nearly lossless performance on all of these tasks. For efficiency results, COAT achieves \boldsymbol{$1.54\times$} end-to-end memory reduction compared with BF16, and \boldsymbol{$1.43\times$} end-to-end training speed up on Llama 7B, 13B, and 30B models compared to BF16. COAT also \emph{doubles the batch size} in all realistic distributed training settings, which is crucial for higher speedup and support for longer context length, leading to a more efficient training process for large-scale models.
\section{Related Work}
\paragraph{Low-precision Training}
Low precision training ~\citep{wang2018trainingfp8,chen2020statistical,lin2022device256kb,wortsman2023int8clip,xi2024jetfire} has become a prominent technique in modern deep learning, offering reductions in both computational costs and memory requirements.
FP16 (half-precision) training ~\citep{micikevicius2017mixedprecision} is the most prevalent low-precision method nowadays. It introduces loss scaling to address FP16’s narrower representation range problem. BF16 training ~\citep{kalamkar2019bf16study} refines this approach, as BF16 has a larger representation range and is more stable for large-scale training. In these approaches, forward and backward passes are computed in FP16 or BF16 precision, while master weights, gradients, and optimizers are stored in FP32.
FP8 training ~\citep{fishman2024fp8trillion,micikevicius2022fp8format} aims to push these efficiency gains further. With the introduction of Nvidia’s Hopper GPU architecture, FP8 is emerging as a practical datatype for next-generation low-precision training. Nvidia's Transformer Engine (TE) ~\citep{nvidia_transformerengine} is the first framework designed for FP8 mixed-precision training that employs FP8 Tensorcore for linear layer calculation. FP8-LM ~\citep{peng2023fp8lm} extends FP8 quantization to gradients and optimizer states, further improving the training throughput. However, they fail to reduce the memory usage of activations stored for the backward pass using FP8, and leave second-order momentum in FP16, limiting the full potential of FP8's memory advantages.
\paragraph{Memory Efficient Optimizers}
While 32-bit optimizer ~\citep{kingma2014adam,loshchilov2017adamw} states are widely adopted, several research efforts have been made to reduce the memory footprint of optimizer states through quantization. The 8-bit Adam ~\citep{dettmers2021adam8bit} introduces a novel data format called dynamic exponent (DE) for quantization, but the adoption of this new data format limits its flexibility. 4-bit Optimizer ~\citep{li2024optimizer4bit} further pushes the limit of optimizer quantization to 4-bit by addressing the zero-point problem, but is restricted to fine-tuning tasks. FP8-LM ~\citep{peng2023fp8lm} quantizes the first-order momentum to FP8 while leaving second-order momentum in FP16, which limits the overall memory savings. ~\citep{fishman2024fp8trillion} finds that second-order momentum is more sensitive to quantization, and proposes to quantize it using E5M2 format.
In addition to quantization, there are other approaches that aim to reduce the memory footprint of the optimizer states ~\citep{shazeer2018adafactor, anil2019sm3,chen2024lion,zhao2024galore}, such as low-rank decomposition and optimizer simplification that only store the first-order momentum. These methods are orthogonal to our approach.
\paragraph{Activation Quantization}
Recent research has focused on reducing memory footprint~\cite{cai2020tinytl} during neural network training through activation quantization. ActNN~\citep{chen2021actnn} introduced a 2-bit activation compressed training framework using random quantization, achieving 12x reduction in activation memory footprint1. GACT~\citep{liu2022gact} extended this concept to support various machine learning tasks and architectures, providing up to 8.1x reduction in activation memory for CNNs, transformers, and GNNs. However, these methods mostly focus on convolutional networks and are not applicable to LLMs. AQ-SGD~\citep{wang2022aqsgd} proposed compressing activation changes rather than direct values for pipeline-parallel training to reduce the communication overhead in the pipeline-parallel setting. Few-Bit Backward\citep{novikov2023fewbitbackward} quantizes the non-linear activation functions using optimal piecewise-constant approximations to reduce the memory footprint while maintaining convergence. Jetfire~\citep{xi2024jetfire} proposes INT8 data flow to quantize the activation of both linear and non-linear layers to reduce memory footprint and is applicable to language model pretraining.
\paragraph{Bit-Width Under-Utilization}
Balance Quantization~\cite{zhou2017balancedquant} recursively partitions the parameters by percentiles to reduce quantization error. DDSQ~\cite{wang2023DDQgradient} dynamically changes the quantization parameters according to different gradient distributions. N2UQ~\cite{liu2022NU2Qnonuniform} improves the quantization precision via learning input thresholds and nonuniform mapping. These approaches are aware of the dynamic range problem in quantization and propose more adaptive, distribution-aware techniques that can preserve the nuanced information contained in full precision. They share a fundamental goal of expanding the effective dynamic range of low-precision neural network representations. These methods adopt a learnable non-uniform quantization lookup table, while our work relies on the natural FP8 data format.
\section{Preliminaries}
\paragraph{FP8 Quantization}
Quantization compresses a high-precision tensor to low-precision to achieve speedup and save memory footprint, at the cost of lower precision and lower representation range.
\emph{FP8 format} consists of two encodings - E4M3 and E5M2~\citep{OCP8bit2023}. E4M3 has higher precision, while E5M2 has a larger representation range. We define E4M3's min and max values as $\Delta_{\min}^{\text{E4M3}} = 2^{-9}$ and $\Delta_{\max}^{\text{E4M3}} = 448$, while E5M2's min and max values are $\Delta_{\min}^{\text{E5M2}}
= 2^{-16}$ and $\Delta_{\min}^{\text{E5M2}} = 57344$. To \emph{quantize} an FP32 tensor $X$ into E4M3 precision, we use a quantizer $Q(\cdot)$ to map the tensor into FP8's representation range. This process can be formulated as
\[X_{\text{FP8}}, S_X = Q(X_{\text{FP32}}),\:\text{where}\: X_{\text{FP8}} = \left\lceil \frac{X_{\text{FP32}}}{S_X} \right\rfloor, \: S_X = \frac{\max\left(\lvert X_{\text{FP32}} \rvert\right)}{\Delta_{\max}^{\text{E4M3}}},
\]
% \jianfei{this is extremely tricky: when you use FP, it should not be ``round''. See my comment in the review mode}
where $S_X$ is the scaling factor, $\lceil \cdot \rfloor$ refers to round-to-nearest. Quantize into E5M2 follows a similar procedure, where we only replace $\Delta_{\max}^{\text{E4M3}}$ with $\Delta_{\max}^{\text{E5M2}}$. To map the quantized tensor back to FP32 precision, the \emph{dequantize} operation $DQ(\cdot)$ can be formulated as
$X_{\text{FP32}} = DQ(X_{\text{FP8}}, S_X) = S_X X_{\text{FP8}}.$
\paragraph{Optimizer Update Rule}
Optimizers are widely used in deep learning to update parameters. The most common gradient-based optimizer is Adam/AdamW~\citep{kingma2014adam,loshchilov2017adamw}, which uses first-order $m$ and second-order momentum $v$ to achieve better convergence. The update rule of AdamW at time step $t$ can be formulated as:
\[
\begin{aligned}
m_t &= \beta_1 m_{t-1} + (1 - \beta_1)g_{t-1} \\
\hat{m}_t &= \frac{m_t}{1 - \beta_1^t} \\
\end{aligned}
\quad \quad \quad
\begin{aligned}
v_t &= \beta_2 v_{t-1} + (1 - \beta_2)g_{t-1}^2 \\
\hat{v}_t &= \frac{v_t}{1 - \beta_2^t}
\end{aligned}
\]
\vspace{-12pt}
\begin{align}\label{eq:adam_update}
w_{t+1} = w_t - \eta \left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda w_t\right)
\end{align}
\vspace{-12pt}
% \begin{align*}
% m_t & = \beta_1 m_{t-1} + (1 - \beta_1) g_{t-1} \\
% v_t & = \beta_2 v_{t-1} + (1 - \beta_2) g_{t-1}^2 \\
% \hat{m}_t & = \frac{m_t}{1 - \beta_1^t} \\
% \hat{v}_t & = \frac{v_t}{1 - \beta_2^t} \\
% w_{t+1} & = w_t - \eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda w_t \right)
% \end{align*}
% This can be formulated as
% \[
% \begin{aligned}
% m_{t-1} &= DQ(m_{t-1}^q, S_{m_{t-1}}) \\
% m_t &= \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t \\
% \hat{m}_t &= \frac{m_t}{1 - \beta_1^t} \\
% m_t^q, S_{m_t} &= Q(m_t) \\
% \end{aligned}
% \quad \quad \quad
% \begin{aligned}
% v_{t-1} &= DQ(v_{t-1}^q, S_{v_{t-1}}) \\
% v_t &= \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2 \\
% \hat{v}_t &= \frac{v_t}{1 - \beta_2^t} \\
% v_t^q, S_{v_t} &= Q(v_t) \\
% \end{aligned}
% \]
% \begin{align}
% w_{t+1} = w_t - \eta \left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda w_t\right), \label{eq:optimizer_step}
% \end{align}
% where $m_{t-1}^q$ and $v_{t-1}^q$ refers to the FP8 optimizer states.\jianfei{these equations looks too messy, think of a way to rewrite. For example, define a function $Adam(m, v, w)$ and say that the new algorithm is $Adam(Q(m), Q(v), w)$}
% \begin{align*}
% m_{t-1} &= DQ(m_{t-1}^q, S_{m_{t-1}}) &(\text{Dequantize to FP32}) \\
% v_{t-1} &= DQ(v_{t-1}^q, S_{v_{t-1}}) &(\text{Dequantize to FP32}) \\
% m_t &= \beta_1 \ast m_{t-1} + (1 - \beta_1) \ast g_t \\
% v_t &= \beta_2 \ast v_{t-1} + (1 - \beta_2) \ast g_t^2 \\
% \hat{m}_t & = \frac{m_t}{1 - \beta_1^t} \\
% \hat{v}_t & = \frac{v_t}{1 - \beta_2^t} \\
% w_{t+1} & = w_t - \eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda w_t \right) \\
% m_t^q, S_{m_t} &= Q(m_t) &(\text{Quantize to FP8}) \\
% v_t^q, S_{v_t} &= Q(v_t) &(\text{Quantize to FP8})
% \end{align*}
where $m_t$ is the first-order momentum, $v_t$ is the second-order momentum, $g_t$ is the gradient, $\beta_1$ and $\beta_2$ are betas of AdamW, $\eta$ is learning rate, $\lambda$ is weight decay, and $\epsilon$ is used to prevent NaN or Inf.
To perform \textbf{optimizer state quantization}, we adopt per-group quantization for both first-order and second-order momentum, following previous works~\citep{dettmers2021adam8bit,li2024optimizer4bit}.
Every consecutive $G$ element forms a group ($G$ is defined as the group size), and each group is quantized independently with its own statistics. Optimizer states are stored in FP8 precision, while its scaling factor is stored in BF16. In FSDP or ZeRO-3, the statistics related to optimizer states (e.g. scaling factor) are not synchronized across GPUs, since each GPU maintains its own shard of optimizer states. More details can be found in Appendix~\ref{appendix:details_optimizer_states}.
% \begin{figure}[t!]
% \begin{minipage}[t]{0.75\textwidth}
% \centering
% \includegraphics[width=\textwidth]{figure/r_11_E4M3_E4M3_final.pdf}
% \captionof{figure}{Optimizer states has a smaller dynamic range than FP8's, especially for second order momentum.}
% % TODO: in figure, fp8 representation range, heavily under utilized, especially for second order }
% % TODO: {add to caption, use per-group quant = 128}
% \label{fig:optimizer_range_ratio}
% \end{minipage}
% \hspace{0.01\textwidth}
% \begin{minipage}[t]{0.23\textwidth}
% \centering
% \includegraphics[width=\textwidth]{figure/k_11_E4M3_E4M3_final.pdf}
% \captionof{figure}{The distribution of the exponential $k$ in the expand function. Second order requires a larger $k$ compared with first order momentum.}
% \label{fig:k_distribution}
% \end{minipage}
% \end{figure}
\begin{figure}[t!]
\centering
\includegraphics[width=\textwidth]{figure/rk_11_E4M3_E4M3_final.pdf}
\captionof{figure}{(a) Visualization of optimizer states' dynamic range under per-group quantization. FP8 E4M3's representation range is under-utilized in this case. (b) After dynamic range expansion, FP8's representation range is well utilized. (c) Distribution of $k$ for optimizer states. The second order's $k$ is larger than the first order's $k$, since the second-order momentum's dynamic range is smaller.}
% TODO: in figure, fp8 representation range, heavily under utilized, especially for second order }
% TODO: {add to caption, use per-group quant = 128}
\label{fig:optimizer_range_ratio}
\vspace{-10pt}
\end{figure}
\section{Dynamic range expansion for accurate optimizer quantization}\label{sec:optimizer_fp8}
% \jianfei{first, present an overview of our approach, then descent into the two main techniques }
In subsequent sections, we explain how COAT utilizes FP8 quantization to achieve memory-efficient FP8 training without compromising accuracy.
Section~\ref{sec:optimizer_fp8} focuses on optimizer states quantization,
while Section~\ref{sec:activation_fp8} discusses activation quantization.
\subsection{Understanding the issue of current optimizer states quantization method}
% TODO:{under per-group quantization},
Under per-group quantization, we find that one significant drawback of current quantization methods is that, they can not \emph{fully utilize the representation range} of FP8 and therefore lead to a large quantization error. Take the E4M3 data format as an example, the ratio between E4M3's maximum representable value and minimum representable value is $\Delta_{\max}^{\text{E4M3}} / \Delta_{\min}^{\text{E4M3}} = 448 \div \frac{1}{512} = 229376 \approx 2\times 10^5.$ Therefore, for a quantization group $X$, if we want to fully utilize the 256 representable value\footnote{Actually among them a small amount of values represents NaN and Inf.} of FP8, we hope the dynamic range of the quantization group $X$ should cover the entire span between $\Delta_{\min}^{\text{E4M3}}$ and $\Delta_{\max}^{\text{E4M3}}$.
To make it more formally, we define \emph{dynamic range} as the ratio between the maximum absolute value and the minimum absolute value within a quantization group $X$:
\begin{definition}[dynamic range]
Given a set that consists of \( G \) real numbers \( X = \{x_1, x_2, \dots, x_G\} \), the \textit{dynamic range} \( \gR \) is defined as
\[
\gR_{X} = \frac{\max(\lvert x_1 \rvert, \lvert x_2 \rvert, \dots, \lvert x_G \rvert)}{\min(\lvert x_1 \rvert, \lvert x_2 \rvert, \dots, \lvert x_G \rvert)},
\]
where $\lvert \cdot \rvert$ denotes the absolute value.
\end{definition}
That is to say, E4M3's dynamic range is $\gR_{\text{E4M3}} = 448 \times 512 = 229376 \approx 2\times 10^5$. However, in practice, many quantization groups within the optimizer states fail to effectively map values across this wide range. We observe that optimizer states are highly sparse, with fewer than 1\% of values having large magnitudes, while the majority are relatively small and closely clustered. Most groups exhibit low dynamic ranges since large values are so few. As visualized in Figure~\ref{fig:optimizer_range_ratio}(a), the dynamic range for first-order momentum is typically less than 1e4, and for second-order momentum, it is usually less than 1e1—both far below the available range of FP8. As a result, a substantial portion of FP8's representational capacity is wasted, leading to a large quantization error.
\subsection{Dynamic Range Expansion}
To address this problem, we introduce a \emph{expand function} $f(\cdot)$ before quantization to expand the dynamic range of the quantization group and align it with E4M3, which can be formalized as $X_{\text{FP8}}, S_X = Q(f(X_{\text{FP32}}))$. The expand function we use is defined as \[f(x) = \operatorname{sign}(x)\lvert x \rvert^k,\] where $k$ is used to control the strength of the expansion. In Appendix.~\ref{appendix:proof_expand_function_optimal} we prove that it is optimal.
For a quantization group $X$, after applying the expand function $f$ to $X$, the dynamic range becomes \[\gR_{f(X)} = \frac{\max(\lvert f(X) \rvert)}{\min(\lvert f(X) \rvert)} = \frac{\max(\lvert \operatorname{sign}(X) X^k \rvert)}{\min(\lvert \operatorname{sign}(X) X^k \rvert)} = \Big(\frac{\max(\lvert X \rvert)}{\min(\lvert X \rvert)}\Big)^k = (\gR_X)^k.\]
Therefore, when $k>1$, $\gR_{X}$ will be enlarged and become closer to the ideal $\gR_{\text{E4M3}}$. The optimal $k$ satisfy that $(\gR_X)^k = \gR_{\text{E4M3}}$, which means that $k = \log_{\gR_X}(\gR_{\text{E4M3}})$. With this optimal $k$, $f(X)$ can fully utilize the representation range of E4M3, while the original $X$ can only utilize a small portion of it. As shown in Figure~\ref{fig:optimizer_range_ratio}(c), the second-order momentum typically has a larger $k$ value ($5\sim 15$) compared to the first-order momentum's $k$ ($1\sim 3$). This corresponds to our previous observation, that second-order momentum usually has a smaller dynamic range compared with first-order momentum, so it requires a larger $k$ to align well with E4M3's dynamic range.
We calculate $k$ on-the-fly for every optimizer step and for every quantization group for accuracy consideration. When dequantize, we apply the inverse of expand function $f^{-1}(x) = x^{\frac{1}{k}}$ after dequantization to recover its original value, which can be expressed as $X^{\text{FP32}} = f^{-1}(DQ(X_{\text{FP8}}, S_X))$.
We apply regular quantizer and our dynamic range expansion method to both the first-order momentum $m$ and second-order momentum $v$. As visualized in Figure~\ref{fig:Optimizer_compander_distribution}(b),
the distribution after expansion can fully utilize the FP8 (E4M3) representation range, which proves the effectiveness of our method. We further quantify the effectiveness of our method in Table~\ref{table:optimizer_quantization_error}. In AdamW optimizer step, as stated in Eq.~\ref{eq:adam_update}, $\frac{m}{\sqrt{v} + \epsilon}$ is the actual effective term for weight update, so we report the MSE of $\frac{m}{\sqrt{v} + \epsilon}$ to quantify the performance of a quantization method. We find that E4M3 is more suitable for first-order momentum than E5M2. For second order momentum, although E4M3 better than E5M2, their quantization error is nearly the same after applying our expand function. Our Dynamic Range Expansion can effectively reduce the MSE by $1.63\times$. Appendix~\ref{appendix:details_optimizer_states} provide more results.
\begin{figure}[t!]
\begin{minipage}{0.54\textwidth}
\centering
\captionof{table}{Quantization error of $\frac{m}{\sqrt{v}}$ under different quantization settings. +Expand means applying our Dynamic Range Expansion method.}
\label{table:optimizer_quantization_error}
\resizebox{\textwidth}{!}{
\begin{tabular}{c|cccc}
\toprule
\textbf{MSE of \boldsymbol{$\frac{m}{\sqrt{v}}$}} & \multicolumn{4}{c}{\textbf{Second Order}} \\
\cmidrule(lr){1-1} \cmidrule(lr){2-5}
\textbf{First Order} & \textbf{E4M3} & \textbf{E4M3+Expand} & \textbf{E5M2} & \textbf{E5M2+Expand} \\
\midrule
\textbf{E4M3} & \textbf{\textcolor{red}{20.10}} & 18.08 & \textbf{\textcolor{red}{25.65}} & 18.16 \\
\textbf{E4M3+Expand} & 15.13 & \textbf{\textcolor{darkgreen}{12.31}} & 21.96 & \textbf{\textcolor{darkgreen}{12.43}} \\
\midrule
\textbf{E5M2} & \textbf{\textcolor{red}{37.02}} & 35.96 & \textbf{\textcolor{red}{40.30}} & 36.00 \\
\textbf{E5M2+Expand} & 17.79 & 15.48 & 23.84 & 15.57 \\
\bottomrule
\end{tabular}
}
\end{minipage}
\begin{minipage}{0.44\textwidth}
\centering
\includegraphics[width=\textwidth]{figure/ExpandQuantizeComparison2_base2_small.pdf}
\captionof{figure}{Dynamic Range Expansion can better utilize E4M3 representation range.}
\label{fig:Optimizer_compander_distribution}
\end{minipage}
\hspace{0.01\textwidth}
\end{figure}
\section{Mixed-Granularity Activation Quantization}\label{sec:activation_fp8}
\subsection{Decompose the activation memory footprint}
In the forward pass of neural networks, activations must be preserved for the backward pass to calculate gradients. As illustrated in Table~\ref{table:activation_memory_decompose}, non-linear layers such as LayerNorm/RMSNorm ~\citep{ba2016layer,zhang2019root} and Activation Functions ~\citep{hendrycks2016gelu,shazeer2020glu} typically account for approximately 50\% of the memory footprint in the Llama model series~\citet{touvron2023llama}. In contrast, linear layers contribute less than 25\%. Therefore, it is essential to optimize both linear and non-linear layers to reduce activation memory footprint.\footnote{RoPEs and FlashAttentions are more as a distinct module and should be handled separately.}
One straightforward approach to achieve this is to quantize the activations to FP8 format prior to each non-linear and linear layer and only save the quantized tensors for backward. However, this introduces significant overhead due to the additional quantization step. Furthermore, non-linear layers can be sensitive to quantization. For example, GLU activation will amplify the spike in activations, resulting in significant quantization errors if not carefully handled~\citep{yang2024gluspike,fishman2024fp8trillion}. This necessitates us to develop efficient and accurate quantization techniques.
\subsection{Mixed granularity FP8 precision flow}
\begin{table}[t]
\centering
% For GPT2-stype model, Act Func refers to GELU, and Linear refers to the summation of QKV/Attn/FC1/FC2 projection.
\caption{Activation memory footprint of different operators. U is a unit to measure memory usage, where $\text{1U}=\text{Batch Size} \times \text{Sequence Length} \times \text{Hidden Size} \times \text{2 bytes (for BF16).}$ For Llama-style model, Act Func refers to SiLU \& Multiply, and Linear refers to the summation of QKV/Attn/Up/Gate/Down projection. RMSNorm is upcast to float32 in \texttt{transformers} \href{https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py\#L70}{implementation}, so the memory usage of LayerNorm in BF16 is 4U. Our method reduces activation memory by quantizing them to FP8. More details about FlashAttention in Appendix~\ref{appendix:details_activation_linear}.}\label{table:activation_memory_decompose}
% TODO: add explanation
% to appendix}
\resizebox{0.9\textwidth}{!}{%
\begin{tabular}{cccccccccccc}
\toprule
& & \multicolumn{2}{c}{\textbf{Non-Linear}} & \multicolumn{2}{c}{\textbf{Attention}} & & & \multicolumn{2}{c}{\textbf{Reduction Ratio}} \\
\cmidrule(r){3-4} \cmidrule(r){5-6} \cmidrule(r){9-10}
\textbf{} & \textbf{} & \textbf{RMSNorm} & \textbf{Act Func} & \textbf{RoPE} & \textbf{FlashAttn} & \textbf{Linear} & \textbf{Total} & \textbf{Ideal} & \textbf{Achieved} \\ \midrule
\multirow{3}{*}{\textbf{Llama-style}} & \textbf{BF16} & 4U & 8U & 2U & 3U & 5.66U & 22.66U & 1.00$\times$ & 1.00$\times$ \\
\textbf{} & \textbf{TE} & 2U & 8U & 2U & 3U & \textcolor{darkgreen}{3.33U} & 18.33U & 1.23$\times$ & 1.20$\times$ \\
\textbf{} & \textbf{COAT} & \textcolor{darkgreen}{1U} & \textcolor{darkgreen}{4U} & 2U & 3U & \textcolor{darkgreen}{3.33U} & \textcolor{darkgreen}{13.33U} & 1.69$\times$ & 1.65$\times$\\
% \midrule
% \multirow{3}{*}{\textbf{GPT2-style}} & \textbf{BF16} & 4U & 4U & - & 3U & 7U & 18U & - \\
% \textbf{} & \textbf{TE} & 4U & 4U & - & 3U & \textcolor{darkgreen}{4U} & 14U & 1.28$\times$ \\
% \textbf{} & \textbf{COAT} & \textcolor{darkgreen}{1U} & \textcolor{darkgreen}{2U} & - & 3U & \textcolor{darkgreen}{4U} & \textcolor{darkgreen}{10U} & 1.8$\times$ \\
\bottomrule
\end{tabular}
}
\end{table}
To address the inefficiency and inaccurate problem,
we propose to use \textbf{mixed granularity FP8 precision flow} to improve the accuracy without introducing too much overhead.
FP8 precision flow requires the input and output of all linear and non-linear layers in FP8. By directly saving the input tensor in FP8 format for the backward pass, we eliminate the need for an extra quantization operation, which reduces the associated overhead. However, this method still suffers from accuracy degradation and necessitates further refinement.
% However, this method still suffers from accuracy degradation and necessitates further refinement to improve precision while keeping overhead minimal.
% TODO: VS-Quant cite (inference)
% TODO: Tensorcore does not like per-group quant, so we only use fine-grained for non-tensorcore part
We propose to vary the quantization granularity across different layers to balance precision and efficiency in a \emph{mixed-granularity} manner. For non-linear layers, VS-Quant~\citep{dai2021vsquant} or Per-Block Quant~\citep{xi2024jetfire} methods are well-suited due to their fine-grained and precise nature. For linear layers, we apply per-tensor quantization to maximize the performance of Tensor Cores.\footnote{Finer-grained methods can also be applied since they are generally compatible with FP8 precision flow.}
We observe that quantizing the input of layernorm across multiple token axes is detrimental to accuracy. As illustrated in Figure~\ref{fig:quantization_error_and_time_scaling}(a), when the number of elements that share a scaling factor is fixed, the quantization error increases significantly when quantization is performed across the token axis. Therefore instead of using per-block quantization with block size $B \times B$ as proposed in ~\citep{xi2024jetfire}, we propose to use per-group quantization with group size $1\times G$, where $G=B^2$ to keep the granularity the same. This approach enhances the accuracy of non-linear layers while maintaining efficiency.
Our precise FP8 precision flow is visualized in Figure~\ref{fig:overview_of_flow_te_and_bar}(a), where we display the full precision flow for a Llama-style decoder layer, both forward and backward pass.
\begin{figure}[t]
\begin{minipage}{\textwidth}
\centering
\includegraphics[width=\textwidth]{figure/LNQuantErrorAndReduction3.eps}
\captionof{figure}{(a) Quantization Error in forward pass. (b) Time comparison of various scaling methods.}\label{fig:quantization_error_and_time_scaling}
\end{minipage}
\vspace{-10pt}
\end{figure}
We also propose Group Scaling, an efficient per-tensor scaling method that balances the performance and precision. To perform per-tensor quantization, the maximum absolute value of the tensor needs to be calculated through max reduction, adding a lot of overhead.
In our Group Scaling, we address these problems by splitting the max reduction into two stages: (1) performing max reduction on each $1 \times G$ element and storing the results as intermediate values; (2) applying max reduction on the intermediate tensor to obtain the per-tensor max value. The first stage can be seamlessly fused with the previous operation, adding minimal overhead, while the second stage is more efficient than doing max reduction on the entire tensor, as the intermediate result is $G\times$ smaller than the original tensor. As illustrated in Figure~\ref{fig:quantization_error_and_time_scaling}(b), Group Scaling successfully reduces the max reduction overhead compared with just-in-time scaling.
In comparison, TransformerEngine proposes delayed scaling to avoid the on-the-fly max reduction required for per-tensor quantization, and can also fuse the division process into the previous operator to optimizes memory accesses.
However, we find that using the current tensor’s statistics to compute the scaling factor is no worse, and could be even better in precision, than using the delayed scaling heuristic. Therefore, we advocate for Group Scaling as a simpler and more flexible alternative that can potentially offer improved numerical stability and precision while not being significantly slower than Delayed Scaling.
% However, their delayed scaling method adds implementation complexity and introduces instability by requiring a buffer to track the maximum values in history for every per-tensor quantization operation.
% Just-in-time scaling uses the maximum absolute value of the tensor being quantized to perform per-tensor quantization. It consists of two parts (1) Reduction of the maximum absolute value, and (2) performing element-wise division to quantize to FP8, and this results in a lot of overhead. To avoid the overhead of (1), TransformerEngine proposes to use delayed scaling. Note that the overhead of (2) can not be avoided.
\section{Experiments}
\subsection{Accuracy Experiments}
\paragraph{Setups}
We compare COAT with BF16 training and TransformerEngine ~\citep{nvidia_transformerengine} baselines. For all experiments, we adopt the default hyperparameters in the official training recipe. We validate the effectiveness of our method on multiple tasks, including Large Language Model (LLM) pretraining and fine-tuning, and Vision Language model (VLM) training. For LLM pertaining, we report the perplexity on Wikitext
103~\citep{merity2016wikitext}, C4~\citep{raffel2020c4}, and Pile~\citep{gao2020pile}, and the accuracy on COPA~\citep{gordon-etal-2012-copa}, ARC~\citep{clark2018arc}, SciQ~\citep{welbl2017sciq}, and HellaSwag~\citep{zellers2019hellaswag}. For LLM fine-tuning, we conduct experiments in math corpus, and evaluate on Mathmeticas~\citep{davies2021deepmind}, SVAMP~\citep{patel2021svamp}, NumGLUE~\citep{mishra2022numglue},
and GSM8K~\citep{cobbe2021gsm8k}.
For VLM training, we report the score on VideoMME~\citep{fu2024videomme}, POPE~\citep{li2023pope}, VizWiz~\citep{gurari2018vizwiz}, GQA~\citep{hudson2019gqa}, VQAv2~\citep{goyal2017vqav2}, TextVQA~\citep{singh2019textvqa}, SEED~\citep{li2023seed}, and MMMU Validation Set~\citep{yue2024mmmu}. We use $1\times 128$ per-group quantization for optimizer states and $1\times 16$ per-group quantization for non-linear layer activations.
\subsubsection{LLM pretraining}
To evaluate our method when pretraining LLMs, We train OLMo-1B~\citep{groeneveld2024olmo} and OLMo-7B on Dolma~\citep{soldaini2024dolma}. Following the official report, we use a global batch size of 4M tokens (2048 macro batch size, with a sequence length of 2048 tokens). We use PyTorch FSDP in our experiments.
For OLMo-1B, we conduct pretraining for 300B tokens, which corresponds to 75k training steps.
We report the training curve in Figure ~\ref{fig:olmo-1b}, and report the perplexity and accuracy result in Table~\ref{table:olmo-1b}. The training curve
and downstream task performance were consistent with BF16 training baseline, validating the effectiveness of COAT.
For OLMo-7B experiment, we pretrain for 250B tokens, which corresponds to 40k training steps.we report the training curve in Figure~\ref{fig:olmo-7b} and downstream task performance in Table~\ref{table:olmo-7b}, where COAT also aligns well with the baseline and is nearly lossless.
% \begin{table}[t]
% \begin{minipage}{0.66\textwidth}
% \centering
% \captionof{table}{OLMo-1B pretraining performance on downstream tasks. TE refers to TransformerEngine.}
% \label{table:olmo-1b}
% \resizebox{\textwidth}{!}{
% \begin{tabular}{lccccc}
% \toprule
% & \textbf{Train Loss} & \textbf{WikiText} & \textbf{C4} & \textbf{Pile} & \textbf{Avg ppl} \\
% \midrule
% \textbf{BF16} & 2.995 & 30.499 & 27.966 & 17.405 & 25.290 \\
% \textbf{TE} & 3.001 & 30.736 & 28.064 & 17.434 & 25.411 \\
% \textbf{COAT} & 3.008 & 30.619 & 28.099 & 17.453 & 25.391 \\
% \midrule
% & \textbf{COPA} & \textbf{ARC(Easy)} & \textbf{SciQ} & \textbf{HellaSwag} & \textbf{Avg Acc} \\
% \midrule
% \textbf{BF16} & 60.0\% & 45.6\% & 67.3\% & 33.7\% & 51.6\% \\
% \textbf{TE} & 62.0\% & 45.4\% & 63.9\% & 33.8\% & 51.3\% \\
% \textbf{COAT} & 61.0\% & 44.2\% & 67.6\% & 33.7\% & 51.5\% \\
% \bottomrule
% \end{tabular}
% }
% \end{minipage}%
% \hspace{0.02\textwidth}
% \begin{minipage}{0.32\textwidth}
% \centering
% \includegraphics[width=\textwidth]{figure/OLMo_1B_curve.pdf}
% \captionof{figure}{OLMo-1B pretraining loss curve.}
% \label{fig:olmo-1b}
% \end{minipage}
% \end{table}
% \begin{table}[t]
% \begin{minipage}{0.70\textwidth}
% \centering
% \captionof{table}{OLMo-1B pretraining performance on downstream tasks. We report the performance after training for 300B tokens.}
% \label{table:olmo-1b}
% \resizebox{\textwidth}{!}{
% \begin{tabular}{lccccc}
% \toprule
% & \textbf{Train Loss} & \textbf{WikiText} & \textbf{C4} & \textbf{Pile} & \textbf{Avg ppl} \\
% \midrule
% \textbf{BF16} & 2.551 & 15.234 & 15.538 & 10.083 & 13.618 \\
% \textbf{COAT} & 2.568 & 15.384 & 15.695 & 10.176 & 13.751 \\
% \midrule
% & \textbf{COPA} & \textbf{ARC(Easy)} & \textbf{SciQ} & \textbf{HellaSwag} & \textbf{Avg Acc} \\
% \midrule
% \textbf{BF16} & 77.0\% & 57.3\% & 84.0\% & 54.5\% & 68.2 \% \\
% \textbf{COAT} & 75.0\% & 58.1\% & 83.9\% & 54.4\% & 67.9 \% \\
% \bottomrule
% \end{tabular}
% }
% \end{minipage}%
% \hspace{0.02\textwidth}
% \begin{minipage}{0.27\textwidth}
% \centering
% \includegraphics[width=\textwidth]{figure/OLMo_1B_curve.pdf}
% \captionof{figure}{OLMo-1B training loss curve.}
% \label{fig:olmo-1b}
% \end{minipage}
% \end{table}
\begin{table}[t]
\begin{minipage}{0.70\textwidth}
\centering
\captionof{table}{OLMo-1B pretraining performance on downstream tasks. We report the performance after training for 300B tokens.}
\label{table:olmo-1b}
\resizebox{\textwidth}{!}{
\begin{tabular}{lccccc}
\toprule
& \textbf{Train Loss} & \textbf{WikiText} & \textbf{C4} & \textbf{Pile} & \textbf{Avg ppl} \\
\midrule
\textbf{BF16} & 2.551 & 15.234 & 15.538 & 10.563 & 10.083 \\
\textbf{COAT} & 2.568 & 15.384 & 15.695 & 10.672 & 10.176 \\
\midrule
& \textbf{COPA} & \textbf{ARC(Easy)} & \textbf{SciQ} & \textbf{HellaSwag} & \textbf{Avg Acc} \\
\midrule
\textbf{BF16} & 77.0 \% & 57.3 \% & 84.0 \% & 54.5 \% & 68.2 \% \\
\textbf{COAT} & 75.0 \% & 58.1 \% & 83.9 \% & 54.3 \% & 67.8 \% \\
\bottomrule
\end{tabular}
}
\end{minipage}%
\hspace{0.02\textwidth}
\begin{minipage}{0.27\textwidth}
\centering
\includegraphics[width=\textwidth]{figure/OLMo_1B_curve.pdf}
\captionof{figure}{OLMo-1B training loss curve.}
\label{fig:olmo-1b}
\end{minipage}
\end{table}
% \begin{table}[t]
% \begin{minipage}{0.70\textwidth}
% \centering
% \captionof{table}{OLMo-1B pretraining performance on downstream tasks. We report the performance after training for 200B tokens.}
% \label{table:olmo-1b}
% \resizebox{\textwidth}{!}{
% \begin{tabular}{lccccc}
% \toprule
% & \textbf{Train Loss} & \textbf{WikiText} & \textbf{C4} & \textbf{Pile} & \textbf{Avg ppl} \\
% \midrule
% \textbf{BF16} & 2.781 & 16.037 & 16.340 & 10.563 & 14.313 \\
% \textbf{COAT} & 2.794 & 16.265 & 16.351 & 10.672 & 14.429 \\
% \midrule
% & \textbf{COPA} & \textbf{ARC(Easy)} & \textbf{SciQ} & \textbf{HellaSwag} & \textbf{Avg Acc} \\
% \midrule
% \textbf{BF16} & 71.0 \% & 57.0 \% & 82.4 \% & 52.4 \% & 65.7 \% \\
% \textbf{COAT} & 69.0 \% & 58.1 \% & 81.7 \% & 51.5 \% & 65.1 \% \\
% \bottomrule
% \end{tabular}
% }
% \end{minipage}%
% \hspace{0.02\textwidth}
% \begin{minipage}{0.27\textwidth}
% \centering
% \includegraphics[width=\textwidth]{figure/OLMo_1B_curve.pdf}
% \captionof{figure}{OLMo-1B training loss curve.}
% \label{fig:olmo-1b}
% \end{minipage}
% \end{table}
\begin{table}[t]
\begin{minipage}{0.70\textwidth}
\centering
\captionof{table}{OLMo-7B pretraining performance on downstream tasks. We report the performance after training for 250B tokens.}
\label{table:olmo-7b}
\resizebox{\textwidth}{!}{
\begin{tabular}{lccccc}
\toprule
& \textbf{Train Loss} & \textbf{WikiText} & \textbf{C4} & \textbf{Pile} & \textbf{Avg ppl} \\
\midrule
\textbf{BF16} & 2.366 & 12.053 & 12.874 & 8.596 & 11.174 \\
\textbf{COAT} & 2.379 & 12.166 & 12.988 & 8.684 & 11.279 \\
\midrule
& \textbf{COPA} & \textbf{ARC(Easy)} & \textbf{SciQ} & \textbf{HellaSwag} & \textbf{Avg Acc} \\
\midrule
\textbf{BF16} & 83.0\% & 65.7\% & 87.5\% & 56.9\% & 73.2 \% \\
\textbf{COAT} & 81.0\% & 61.9\% & 87.2\% & 60.6\% & 72.7 \% \\
\bottomrule
\end{tabular}
}
\end{minipage}%
\hspace{0.02\textwidth}
\begin{minipage}{0.27\textwidth}
\centering
\includegraphics[width=\textwidth]{figure/OLMo_7B_curve.pdf}
\captionof{figure}{OLMo-7B training loss curve.}
\label{fig:olmo-7b}
\end{minipage}
\end{table}
\begin{table}[h]
\centering
\caption{Evaluation result of fine-tuning Llama-2-7B on math corpus. \texttt{Llama-2-7B} refers to the evaluation metric before fine-tuning. TE refers to TransformerEngine.}
\begin{minipage}{0.70\textwidth}
\label{table:math_corpus_finetune}
\resizebox{\textwidth}{!}{
\begin{tabular}{cccccc}
\toprule
& \textbf{Mathmeticas} & \textbf{SVAMP} & \textbf{NumGLUE} & \textbf{GSM8k} & \textbf{Avg} \\ \midrule
\texttt{Llama-2-7B} & 6.0 & 14.6 & 34.5 & 29.9 & 21.3 \\
\midrule
BF16 & 46.3 & 64.2 & 54.8 & 57.7 & 55.7 \\
TE & 45.3 & 66.1 & 53.5 & 57.7 & 55.6 \\
COAT & 47.8 & 64.4 & 53.3 & 56.6 & 55.5 \\
\bottomrule
\end{tabular}
}
\end{minipage}
\end{table}
\subsubsection{LLM fine-tuning}
We further evaluate our method on LLM fine-tuning task. We focus on math corpus and fine-tune a Llama-2-7B model on the MAmmoTH~\citep{yue2023mammoth} dataset. We train for 3 epochs, and report the downstream tasks performance in ~\ref{table:math_corpus_finetune}. After fine-tuning, COAT still performs consistently with the baselines on the downstream task performance, proving the accurateness of our method.
\subsubsection{VLM Training}
We also evaluate COAT on vision language models. We conduct experiments on VILA ~\citep{lin2024vila} and perform stage-3 SFT of VILA1.5-7B using the same SFT data mixture employed by VILA's original paper~\citep{chen2023sharegpt4v,xu2024visionflan} and set the global batch size to 1024. We pad the sequence length to a multiple of 4 for efficiency consideration. The training loss curve is visualized in Figure~\ref{fig:vila-stage3}. We report the downstream task performance and their average in Table~\ref{table:vila-stage3}. We find that COAT performs on par with BF16 training and is better than the TransformerEngine baseline, which demonstrates the accurateness of our method. We further visualize the VLM captioning experiment in Appendix~\ref{appendix:vision_caption} to prove the effectiveness of COAT on generation tasks.
\begin{figure}[t]
\begin{minipage}{0.61\textwidth}
\centering
\captionof{table}{VILA1.5-7B Stage-3 SFT performance on downstream tasks. * means it has seen the training data.}
\label{table:vila-stage3}
\resizebox{\textwidth}{!}{
\begin{tabular}{lcccccc}
\toprule
Stage 3 & VideoMME & POPE & VizWiz & GQA* & VQAv2* \\
\midrule
BF16 & 42.96 & 86.90 & 61.42 & 64.55 & 81.47 \\
TE & 43.19 & 87.64 & 57.61 & 64.53 & 81.34 \\
COAT & 44.56 & 87.43 & 61.36 & 64.44 & 81.20 \\
\midrule
& & \multicolumn{2}{c}{SEED} & \\
\cmidrule(lr){3-4}
Stage 3 & TextVQA & Image & Video & MMMU Val & Average \\
\midrule
BF16 & 65.60 & 73.40 & 45.65 & 38.56 & 62.80 \\
TE & 64.70 & 73.51 & 43.12 & 35.89 & 61.88 \\
COAT & 64.65 & 73.36 & 43.76 & 37.22 & 62.51 \\
\bottomrule
\end{tabular}
}
\end{minipage}%
\hspace{0.02\textwidth}
\begin{minipage}{0.36\textwidth}
\centering
\includegraphics[width=\textwidth]{figure/VILA_7B_v2_curve.pdf}
\caption{VILA1.5-7B Stage-3 SFT loss curve.}
\label{fig:vila-stage3}
\end{minipage}
\end{figure}
% \subsubsection{ViT pretraining}
\subsection{Memory Saving and Speedup}
We test the memory saving and speedup result of COAT in two settings: The results on a single transformer layer help to accurately analyze the capabilities for memory saving and speedup, while the end-to-end results reflect the practical benefits of our method in real-world applications.
\subsubsection{Memory Saving and Speedup for a Single Transformer Layer}
Table~\ref{table:latency_single_layer} highlights the speedup and memory reduction achieved for a single transformer layer. We conducted experiments with a batch size of 4, varying the hidden sizes between 2048 and 4096, and sequence lengths of 2048 and 4096.
COAT demonstrates better speedup compared with TE, and significantly better memory reduction ability compared with BF16 and TE. Our approach achieves up to $1.57\times$ speedup over BF16 and achieves a consistent $1.65\times$ memory reduction compared to BF16, which is very close to the theoretically $1.69\times$ reported in Table~\ref{table:activation_memory_decompose}.
The speedup ratio becomes larger with larger hidden sizes and longer sequence lengths.
% \begin{figure}[h]
% \centering
% \includegraphics[width=\textwidth]{figure/PerLayerSpeedUp.pdf}
% \caption{Visualization of our precise fp8 precision flow.}
% \label{fig:fp8_precision_flow}
% \end{figure}
\begin{table}[t!]
\centering
\caption{Memory Saving and Speedup for a single Transformer Layer. Memory refers to Activation Memory. Our method achieves better speedup than TransformerEngine and significantly reduces the activation memory footprint by $1.65\times$.}
\label{table:latency_single_layer}
\resizebox{0.95\textwidth}{!}{%
\begin{tabular}{ccccccc|ccccccc}
\toprule
\multicolumn{4}{l}{\textit{Hidden Size = 2048, Batch Size = 4}} \\
\midrule
& \multicolumn{6}{c}{Sequence Length = 2048} & \multicolumn{6}{c}{Sequence Length = 4096} \\
\cmidrule(lr){2-7} \cmidrule(lr){8-13}
& Forward & Backward & Total & Ratio & Memory & Ratio & Forward & Backward & Total & Ratio & Memory & Ratio \\
\midrule
BF16 & 3.36 & 8.47 & 11.83 & $1.00\times$ & 1457 MB & $1.00\times$ & 6.88 & 17.24 & 24.12 & $1.00\times$ & 2914 MB & $1.00\times$ \\
TE & 2.96 & 5.32 & 8.28 & $1.42\times$ & 1210 MB & $1.20\times$ & 5.94 & 11.29 & 17.23 & $1.39\times$ & 2420 MB & $1.20\times$ \\
\rowcolor{lightblue}
COAT & 2.88 & 5.16 & 8.04 & \boldsymbol{$1.47\times$} & 883 MB & \boldsymbol{$1.65\times$} & 5.89 & 10.82 & 16.71 & \boldsymbol{$1.44\times$} & 1766 MB & \boldsymbol{$1.65\times$} \\
\midrule
\multicolumn{4}{l}{\textit{Hidden Size = 4096, Batch Size = 4}} \\
\midrule
& \multicolumn{6}{c}{Sequence Length = 2048} & \multicolumn{6}{c}{Sequence Length = 4096} \\
\cmidrule(lr){2-7} \cmidrule(lr){8-13}
& Forward & Backward & Total & Ratio & Memory & Ratio & Forward & Backward & Total & Ratio & Memory & Ratio \\ \midrule
BF16 & 7.77 & 18.78 & 26.55 & $1.00\times$ & 2914 MB & $1.00\times$ & 16.37 & 38.43 & 54.80 & $1.00\times$ & 5828 MB & $1.00\times$ \\
TE & 6.19 & 11.79 & 17.98 & $1.47\times$ & 2420 MB & $1.20\times$ & 12.66 & 24.58 & 37.24 & $1.47\times$ & 4840 MB & $1.20\times$ \\
\rowcolor{lightblue}
COAT & 5.89 & 10.96 & 16.85 & \boldsymbol{$1.57\times$} & 1766 MB & \boldsymbol{$1.65\times$} & 12.16 & 23.44 & 35.6 & \boldsymbol{$1.53\times$} & 3533 MB & \boldsymbol{$1.65\times$} \\
\bottomrule
\end{tabular}
}
\end{table}
\subsubsection{Speedup and Memory Saving for end-to-end training}
Table~\ref{table:end_to_end_speed_and_memory} presents a detailed comparison of end-to-end memory reduction and speedup results across different configurations for transformer models, specifically Llama-2-7B, Llama-2-13B, and Llama-30B, with variations in the number of GPUs used (1, 2, 4, and 8). It highlights COAT's effectiveness in reducing end-to-end memory footprint and the speedup compared to standard BF16 and TransformerEngine (TE) setups under varying conditions of batch size and context length.
COAT allows full-parameter training of Llama-2-7B on \emph{a single GPU}, where BF16 and TE both out of memory (OOM). Similarly, for the Llama-2-13B and Llama-30B models, our method enables 2-GPU training for Llama-2-13B and 8-GPU training for Llama-30B when Batch Size = 1.
In all multi-GPU training setting, COAT can double the micro-batch size and therefore lead to even higher speedup. For example, our method can achieve $2.25\times$ speedup when training Llama-2-13B on 4-GPUs since we can effectively increase the batch size to 2.
Overall, COAT significantly reduces end-to-end memory usage by up to $1.55\times$ and speeds up the end-to-end training by nearly $1.44\times$. This facilitates full-parameter training on fewer GPUs, which is particularly beneficial for larger language models.
\begin{table}[t]
\centering
\caption{End-to-end memory reduction and speedup results. BS refers to batch size. CL refers to context length. We report token/s per GPU for speed results. $\ddagger$ means CL=1024.}\label{table:end_to_end_speed_and_memory}
\resizebox{0.9\textwidth}{!}{%
\begin{tabular}{cccccc|ccc}
\toprule
\multicolumn{2}{l}{\textbf{Llama-2-7B}} &
\multicolumn{4}{l|}{\emph{Context Length = 2048}} & \multicolumn{3}{l}{\emph{Maximum Batch Size, Context Length = 2048}} \\
\midrule
& & Optimizer & Activations & Peak & Ratio & Max BS & Speed & Ratio \\
\cmidrule(lr){1-9}
\multirow{3}{*}{1 $\text{GPU}_{\text{BS=1}}$} & BF16 & - & - & \textcolor{red}{OOM} & - & - & \textcolor{red}{OOM} & - \\
& TE & - & - & \textcolor{red}{OOM} & - & - & \textcolor{red}{OOM} & - \\
\rowcolor{lightblue}
& COAT & 13.1 GB & 8.1 GB & 79.3 GB & \textcolor{darkgreen}{\ding{51}} & \textcolor{darkgreen}{\textbf{1}} & 5906 token/s & \textcolor{darkgreen}{\ding{51}} \\
\cmidrule(lr){1-9}
\multirow{3}{*}{2 $\text{GPU}_{\text{BS=2}}$} & BF16 & - & - & \textcolor{red}{OOM} & - & 1 & 6130 token/s & 1.00$\times$ \\
& TE & - & - & \textcolor{red}{OOM} & - & 1 & 6842 token/s & 1.11$\times$ \\
\rowcolor{lightblue}
& COAT & 6.5GB & 16.9 GB & 52.8 GB & \textcolor{darkgreen}{\ding{51}} & \textcolor{darkgreen}{\textbf{4}} & 11351 token/s & \textcolor{darkgreen}{\boldsymbol{$1.85\times$}} \\
\cmidrule(lr){1-9}
\multirow{3}{*}{4 $\text{GPU}_{\text{BS=2}}$} & BF16 & 13.1 GB & 25.8 GB & 55.1 GB & $1.00\times$ & 2 & 7730 token/s & $1.00\times$ \\
& TE & 13.1 GB & 21.9 GB & 51.1 GB & 1.08$\times$ & 2 & 9577 token/s & 1.24$\times$ \\
\rowcolor{lightblue}
& COAT & 3.2 GB & 16.9 GB & 35.6 GB & \textcolor{darkgreen}{\boldsymbol{$1.54\times$}} & \textcolor{darkgreen}{\textbf{4}} & 11257 token/s & \textcolor{darkgreen}{\boldsymbol{$1.45\times$}} \\
\cmidrule(lr){1-9}
\multirow{3}{*}{8 $\text{GPU}_{\text{BS=2}}$} & BF16 & 6.5 GB & 25.8 GB & 41.2 GB & $1.00\times$ & 4 & 8238 token/s & $1.00\times$ \\
& TE & 6.5 GB & 21.9 GB & 37.2 GB & 1.11$\times$ & 4 & 11704 token/s & \textcolor{darkgreen}{1.42$\times$} \\
\rowcolor{lightblue}
& COAT & 1.6 GB & 16.9 GB & 27.0 GB & \textcolor{darkgreen}{\boldsymbol{$1.52\times$}} & \textcolor{darkgreen}{\textbf{8}} & 11241 token/s & $1.36\times$ \\
\midrule
\multicolumn{2}{l}{\textbf{Llama-2-13B}} &
\multicolumn{4}{l|}{\emph{Context Length = 2048}} & \multicolumn{3}{l}{\emph{Maximum Batch Size, Context Length = 2048}} \\
\midrule
& & Optimizer & Activations & Peak & Ratio & Max BS & Speed & Ratio \\
\cmidrule(lr){1-9}
\multirow{3}{*}{\vspace{4pt}2 $\text{GPU}_{\text{BS=1}}^{\ddagger}$} & BF16 & - & - & \textcolor{red}{OOM} & - & - & \textcolor{red}{OOM} & - \\
& TE & - & - & \textcolor{red}{OOM} & - & - & \textcolor{red}{OOM} & - \\
\rowcolor{lightblue}
& COAT & 12.6 GB & 10.1 GB & 73.2 GB & \textcolor{darkgreen}{\ding{51}} & \textcolor{darkgreen}{\textbf{1}} & 2137 token/s & \textcolor{darkgreen}{\ding{51}} \\
\cmidrule(lr){1-9}
\multirow{3}{*}{4 $\text{GPU}_{\text{BS=1}}$} & BF16 & 25.1 GB & 20.1 GB & 76.1 GB & $1.00\times$ & 1 & 2345 token/s & $1.00\times$ \\
& TE & 25.1 GB & 17.2 GB & 73.0 GB & 1.04$\times$ & 1 & 2851 token/s & 1.21$\times$ \\
\rowcolor{lightblue}
& COAT & 6.3 GB & 13.2 GB & 49.1 GB & \textcolor{darkgreen}{\boldsymbol{$1.55\times$}} & \textcolor{darkgreen}{\textbf{2}} & 5295 token/s & \textcolor{darkgreen}{\boldsymbol{$2.25\times$}} \\
\cmidrule(lr){1-9}
\multirow{3}{*}{8 $\text{GPU}_{\text{BS=1}}$} & BF16 & 12.6 GB & 20.1 GB & 49.4 GB & $1.00\times$ & 2 & 3907 token/s & $1.00\times$ \\
& TE & 12.6 GB & 17.2 GB & 46.5 GB & 1.06$\times$ & 2 & 5604 token/s & 1.43$\times$ \\
\rowcolor{lightblue}
& COAT & 3.1 GB & 13.2 GB & 32.5 GB & \textcolor{darkgreen}{\boldsymbol{$1.52\times$}} & \textcolor{darkgreen}{\textbf{4}} & 5650 token/s & \textcolor{darkgreen}{\boldsymbol{$1.44\times$}} \\
\midrule
\multicolumn{2}{l}{\textbf{Llama-30B}} &
\multicolumn{4}{l|}{\emph{Context Length = 2048}} & \multicolumn{3}{l}{\emph{Maximum Batch Size, Context Length = 2048}}\\
\midrule
& & Optimizer & Activations & Peak & Ratio & Max BS & Speed & Ratio \\
\cmidrule(lr){1-9}
\multirow{3}{*}{8 $\text{GPU}_{\text{BS=1}}$} & BF16 & - & - & \textcolor{red}{OOM} & - & - & \textcolor{red}{OOM} & - \\
& TE & - & - & \textcolor{red}{OOM} & - & - & \textcolor{red}{OOM} & - \\
\rowcolor{lightblue}
& COAT & 7.8 GB & 24.2 GB & 70.5 GB & \textcolor{darkgreen}{\ding{51}} & \textcolor{darkgreen}{\textbf{1}} & 1363 token/s & \textcolor{darkgreen}{\ding{51}} \\
\bottomrule
\end{tabular}
}
\end{table}
\subsection{Ablation Studies}
\subsubsection{Dynamic Range Expansion's compatibility with other data formats}
\begin{table}[ht]
\centering
\caption{Dynamic Range Expansion is compatible with DE8 (8-bit dynamic quantization).}\label{table:expand_quant_and_DE8}
\small{
\resizebox{0.85\textwidth}{!}{%
\begin{tabular}{l|ccccccc}
\toprule
& \multicolumn{6}{c}{Second Order} \\
\cmidrule{2-7}
First Order & E4M3 & E4M3 + Expand & E5M2 & E5M2 + Expand & DE8 & DE8 + Expand \\
\midrule
E4M3 + Expand & 15.13 & \textbf{\textcolor{red}{12.31}} & 21.96 & \textbf{\textcolor{red}{12.43}} & 14.01 & 18.84 \\
DE8 & 12.11 & 8.27 & 20.02 & 8.43 & \textbf{\textcolor{red}{10.54}} & 16.25 \\
DE8 + Expand & 11.57 & \textbf{\textcolor{darkgreen}{7.47}} & 19.69 & \textbf{\textcolor{darkgreen}{7.65}} & 9.91 & 15.81 \\
\bottomrule
\end{tabular}
}
}
\end{table}
Dynamic Exponent Quantization (DE) was proposed in ~\citep{dettmers2021adam8bit} to quantize the optimizer states since its representation range has a range of 7 orders of magnitude and is very suitable for optimizer states quantization. Therefore, they proposed to quantize both first-order and second-order momentum with 8-bit DE.
We report the quantization error of $\frac{m}{\sqrt{v}}$ in Table~\ref{table:expand_quant_and_DE8}, and find that applying our dynamic range expansion method to 8-bit DE can further reduce the quantization error by $1.41\times$. Specifically, the lowest quantization error is achieved when first-order momentum $m$ is quantized with 8-bit DE + Dynamic Range Expansion, and second-order momentum $v$ is quantized with E4M3/E5M2 + Dynamic Range Expansion. This proves the effectiveness of our method.
% \subsubsection{Different LLM architectures}
\section{Conclusion}
In this work, we present COAT, a memory-efficient FP8 training framework for foundation models by quantizing both optimizer states and activations into FP8 format.
We observe that the FP8 format's representation range is not fully utilized when quantizing optimizer states, prompting us to propose Dynamic Range Expansion to align their dynamic range. We then identify the importance of quantizing non-linear layers, and propose mixed-granularity FP8 precision flow to quantize the activations accurately without introducing too much overhead.
Extensive experiments on LLM and VLM training and fine-tuning demonstrate that COAT can achieve nearly lossless performance. In end-to-end training, COAT achieves a $1.54×$ memory reduction and $1.43×$ speedup compared to BF16, and is comparable or even faster to TransformerEngine's training speed. Our method also enables full-parameter training of billion-scale models on fewer GPUs and is able to double the batch size in realistic settings. These results highlight COAT's capability to enable memory-efficient, large-scale model training without sacrificing accuracy, providing a highly effective solution for memory-constrained environments. Future work could further explore combining our proposed approach with other low-precision gradient compression methods to reduce communication overhead.
\section{Details about optimizer states quantization}\label{appendix:details_optimizer_states}
When performing \texttt{optimizer.step()}, We first dequantize the optimizer states into FP32, then update the optimizer states and weights in FP32 precision. In the end, we quantize the updated optimizer states back to FP8 and store it in GPU memory. This can be formulated as:
\begin{align*}
m_{t-1} &= DQ(m_{t-1}^q, S_{m_{t-1}}) &(\text{Dequantize to FP32}) \\
v_{t-1} &= DQ(v_{t-1}^q, S_{v_{t-1}}) &(\text{Dequantize to FP32}) \\
m_t &= \beta_1 \ast m_{t-1} + (1 - \beta_1) \ast g_t \\
v_t &= \beta_2 \ast v_{t-1} + (1 - \beta_2) \ast g_t^2 \\
\hat{m}_t & = \frac{m_t}{1 - \beta_1^t} \\
\hat{v}_t & = \frac{v_t}{1 - \beta_2^t} \\
w_{t+1} & = w_t - \eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda w_t \right) \\
m_t^q, S_{m_t} &= Q(m_t) &(\text{Quantize to FP8}) \\
v_t^q, S_{v_t} &= Q(v_t) &(\text{Quantize to FP8})
\end{align*}
In the FSDP setting, dynamic range synchronization across GPUs is unnecessary due to the unique distributed training architecture. The approach leverages sharding strategies that eliminate the need for cross-GPU synchronization. Each GPU independently handles its own optimizer state shards. Since dynamic range is calculated based on local device statistics and optimizer states remain locally computed and updated, we do not need to synchronize the dynamic range across GPUs.
\section{Qualitative Example - Vision Language Model captioning}\label{appendix:vision_caption}
\begin{figure}[h]
\begin{minipage}{\textwidth}
\centering
\includegraphics[width=\textwidth]{figure/ImageExample2.pdf}
\captionof{figure}{Comparison of BF16 and COAT on VLM captioning. COAT can accurately summarize the figure and identify the key points in the figure.}
\label{fig:expand_quant_base2}
\end{minipage}
\end{figure}
\section{Visualization of Expand Fuction}
We further visualize Figure~\ref{fig:Optimizer_compander_distribution} by flattening it. This helps to understand the effectiveness of our method since floating point numbers can be more easily understood in a binary manner.
\begin{figure}[h]
\begin{minipage}{\textwidth}
\centering
\includegraphics[width=\textwidth]{figure/ExpandQuantizeComparison_base2.pdf}
\captionof{figure}{Axis is of base 2.}
\label{fig:expand_quant_base2}
\end{minipage}
\end{figure}
\section{Detailed Explanation for Table~\ref{table:activation_memory_decompose}}\label{appendix:details_activation_linear}
We mainly explain why the reduction ratio of linear layers is not exactly 50\% here. 5.66U comes from: QKV Projection - 1U, Attention Projection - 1U, Up\&Gate Projection - 1U, Down Projection - 2.66U. Among them, the attention projection's input is exactly the output of FlashAttention. However, FlashAttention will save its input (QKV) and its output itself, so FlashAttention and Attention Projection actually share the same tensor when saving it. In Python, this tensor will only be saved for 1 time. So we should not further store an FP8 version of the Attention Projection's input~\citep{zhang2025sageattention,zhang2024sageattention2}, since this will even increase the memory usage by 0.5U if we do not change the code of FlashAttention. Although we need to quantize the input to FP8 again in backward pass, we can store the scaling factor to reduce this additional overhead.
Therefore, after FP8 quantization, the memory usage of linear layers comes from: QKV Projection - 0.5U, Attention Projection - 1U, Up\&Gate Projection - 0.5U, Down Projection - 1.33U. They sum up to 3.33U.
\section{Detailed Explanation for Figure~\ref{fig:overview_of_flow_te_and_bar}(c)}
We mainly discuss FP8-LM in this section. FP8-LM reduces the gradient communication precision from FP32 to FP8 and therefore greatly reduces the communication overhead. It also reduces the master weight's precision from FP32 to FP16/BF16, and reduces the optimizer states precision from FP32 to BF16/FP8. For optimizer states and master weights, the memory footprint reduction result is consistent with our bar. However, for gradient, although FP8-LM can quantize the gradient into FP8, it still needs to preserve the main gradient in FP32 for gradient accumulation. Therefore it can not reduce the memory footprint of the weight gradient.
\section{Proof the expand function is optimal}\label{appendix:proof_expand_function_optimal}
We prove that the choice of $f(x) = x^k$ is optimal given certain basic assumptions for the expansion function $f(x)$.
\begin{enumerate}
\item $f(x)$ is an odd function, where $f(x) + f(-x) = 0$, since the representation range of FP8 is symmetric. Therefore, in the following discussion, we only consider the case when $x \geq 0$.
\item $f(x)$ is continuously differentiable.
\item $f(x)$ is monotonically increasing. This is because we want to keep the relative order before and after quantization or expansion.
\item $f(0) = 0$ and $f(1) = 1$. $f(0) = 0$ since $f(x)$ is odd. $f(1) = 1$ since we can always consider $g(x) = \frac{f(x)}{f(1)}$ without changing its expansion ability.
\item For any $x > y > 0, r > 0$, $\frac{f(x)}{f(y)} = \frac{f(rx)}{f(ry)}$. This is used to make sure the function is scale-invariant. For example, if you multiply the unquantized value $X$ by $10$, the value after expansion and quantization is the same as if you do not multiply the value by $10$.
\end{enumerate}
We prove that based on these assumptions, the only function that meets these assumptions is $f(x) = \text{sign}(x) |x|^k$.
\textbf{Proof:}
We only consider the case when $x > 0$. We denote $f(2) = m > 1$. From assumption 5, we have for any $p > r > s > q$, if $pq = rs$, then $f(p)f(q) = f(r)f(s)$, since
\[
\frac{f(p)}{f(r)} = \frac{f\big(p \cdot (s / p)\big)}{f\big(r \cdot (s / p)\big)} = \frac{f(s)}{f(q)}.
\]
And for any $p = rs$,
\[
f(p) = f(p)f(1) = f(r)f(s).
\]
Therefore, for any $x = 2^t, t \in \mathbb{N}$,
\[
f(x) = f(2^t) = f(2)^t = m^t.
\]
It is also true that for any $x = 2^{-t}, t \in \mathbb{N}$,
\[f(x) = \frac{f(1)}{f(\frac{1}{x})} = \frac{f(1)}{f(2^t)} = \frac{1}{m^t} = m^{-t},\text{so that } f(x) = m^{-t}.\]
Therefore, for any $N \in \mathbb{N}$, $x = 2^t$, $t = \sum\limits_{i=-N}^{N} c_i \cdot 2^i$, $c_i \in \{0, 1\}$,
\[
f(x) = f(\prod_{i=-N}^{N} 2^{c_i \cdot 2^i}) = \prod_{i=-N}^{N} f(c_i \cdot 2^i) = \prod_{i=-N}^{N} m^{c_i \cdot 2^i} = m^{\sum\limits_{i=-N}^{N} c_i \cdot 2^i} = m^t.
\]
For $\forall x > 0$, it can be written in the form of $x = 2^t, t\in \mathbb{R}$. Consider the binary representation of \(t\), denoted by \( t = \sum\limits_{i=-\infty}^{\infty} c_i \cdot 2^i \), where \( c_i \in \{0, 1\} \). For any integer \( N > 0 \), there exists two series \( \{a_i\}_{-N}^N \) and \( \{b_i\}_{-N}^N \) that satisfies
$a_{-N} = 0, b_{-N} = 1, \text{and } a_i = b_i \text{ for } -N+1 \leq i \leq N,$
\[
t_a \leq t < t_b,
\]
where $t_a = \sum\limits_{i=-N}^{N} a_i \cdot 2^i$ and $t_b = \sum\limits_{i=-N}^{N} b_i\cdot 2^i$.
Since \( f(x) \) is continuous, the values of \( f(x) \) at \( x_a = 2^{t_a} \) and \( x_b = 2^{t_b} \) can be used to approximate \( f(x) = f(2^t) \), and by taking the limit as \( N \to \infty \), the equality holds:
\[
f(x) = \lim_{N \to \infty} f(x_a) = \lim_{N \to \infty} f(x_b).
\]
Therefore for $\forall~ t \in \mathbb{R}$, \( f(2^t) = m^t \) for some \( m > 0 \). Therefore $\forall~ x > 0$, $f(x) = m^{\log_2(x)}$.
Using the exponential and logarithmic identity \( m^{\log_2(x)} = x^{\log_2(m)} \), this simplifies to:
\[
f(x) = x^{\log_2(m)},
\]
where \( \log_2(m) \) is a constant determined by $f(2) = m$. If we denote $\log_2(m) = k$, then $f(x) = x^k$.
\section{Implementation Details of COAT}
This section highlights the key differences and details of the implementation of the COAT LLaMA model. The primary focus is on the techniques used to handle quantization and optimization efficiently.
\subsection{Model Architecture}
The COAT LLaMA implementation introduces three new modules, which are not present in the Hugging Face implementation, specifically designed to optimize for FP8 precision flow:
\begin{enumerate}
\item \textbf{BeforeAttention} module (inherit from torch.autograd.Function) calculates the LayerNorm and QKV Projection before RoPE and correctly handles how quantized inputs interact with the residual connection. If Attention is quantized~\citep{zhang2024sageattention,zhang2024sageattention2,shah2024flashattention3}, then the output of this module should also be quantized to accommodate the input of Attention.
\item \textbf{AfterAttention} (inherit from torch.autograd.Function) calculate the attention projection layer and ensures quantized outputs are correctly scaled and summed up with residual connection.
\item \textbf{MLPResidual} (inherit from torch.autograd.Function) handles quantization within the feed-forward network (FFN) efficiently. It also incorporates residual connections in this module to handle it efficiently.
\end{enumerate}
These modules ensure that the computation flow after quantization is correct. We fuse the group scaling method into the \textit{add} operator in residual connection, therefore we need to manually control how the gradient flows in the network.
\subsection{Triton-Based FP8 Kernels for Linear Layers}
The linear layers utilize Triton to implement custom FP8 matrix multiplication kernels. This implementation introduces the following optimizations:
\begin{enumerate}
\item
Block Size Tuning: The block sizes for dimensions \( M \), \( N \), and \( K \) are autotuned, with options chosen from \{128, 256\}. This dynamic tuning optimizes performance for the specific hardware environment.
\item
Group Scaling Technique: The output of linear layer should be BF16 or FP8 per-group quantized tensor. If it should be per-group quantized, to implement it efficiently in Triton, we reshape the output (BLOCK\_M, BLOCK\_N) to (BLOCK\_M, BLOCK\_N // G, G), where G is the group size. We then calculate the maximum value along the last dimension to get the scaling factor.
\end{enumerate}
\subsection{Triton-Based FP8 Kernels for Non-Linear Layers}
\begin{enumerate}
\item
Transposed FP8 Output: Besides generating the quantized FP8 output, these layers also produce a transposed FP8 output. This is critical for backward computations, where column-major inputs are required to calculate weight gradients.
\item
Block Size Tuning: Block sizes for dimensions \( M \) and \( N \) are autotuned, with options chosen from \{32, 64\}.
\end{enumerate}
\subsection{Optimizer States}
We implement a fused kernel for optimizer states update in CUDA. Since optimizer states use $1\times 128$ quantization group size, every CUDA block calculate the result of a quantization group and contains 128 threads, each thread corresponds to one element. We uses warp shuffle functions for reduction to get the maximum absolute value and minimum absolute value within a quantization group to perform dynamic range expansion. We
Direct computation of $X_{\text{min}}^k$ can result in underflow and numerical instability since its magnitude is usually about $10^{-12}$ and $k > 3$. To mitigate this issue, we introduce an additional FP32 value $C_X$, defined as $C_X = \sqrt{X_{\text{min}} \cdot X_{\text{max}}}$, to ensure numerical stability. By normalizing $X$ with $C_X$, the values $\frac{X_{\text{min}}}{C_X}$ and $\frac{X_{\text{max}}}{C_X}$ are computed, which are reciprocals of each other and closer to 1 compared to the original $X_{\text{min}}^k$. This normalization enhances numerical robustness and reduces the risk of instability.
\subsection{Others}
\begin{enumerate}
\item
Backward Gradient Handling: Backward gradients are computed in BF16 precision and are not quantized. This ensures that the gradient computations remain stable and accurate, avoiding numerical issues that can arise with quantization.
Therefore for the backward kernels, the activation stored from forward is in FP8, while the input gradient and the output gradient are both in BF16 precision.
\item
Gradient Accumulation Optimization: During the first gradient accumulation step, the scaling factors of weight matrices are stored. This avoids the need to recompute scaling factors in subsequent steps, reducing computational overhead. The FP8 weights are not stored to prevent additional memory usage, which is critical for large-scale models.
\end{enumerate}
\section{Group Size Analysis}\label{appendix: group size analysis}
For activation, we calculate the quantization error of the output of last layer in forward pass, and the gradient before the first layer in backward pass, comparing with BF16. Based on Table below, we find that group size = 16 achieves a good balance for memory efficiency and quantization error.
\begin{table}[h!]
\centering
\begin{tabular}{ccccc>{\columncolor{lightblue}}ccc}
\toprule
& \multicolumn{2}{c}{\textbf{Only Quantize Linear}} & \multicolumn{5}{c}{\textbf{Also Quantize Activation}}\\
\cmidrule{2-3} \cmidrule{4-8}
\textbf{MSE} & \textbf{1 $\times$ 128 Group} & \textbf{Per Tensor} & \textbf{4} & \textbf{8} & \textbf{16} & \textbf{32} & \textbf{64} \\
\midrule
\textbf{Memory Overhead} & - & - & 50\% & 25\% & 12.5\% & 6.5\% & 3.2\% \\
\textbf{Forward} & 372 & 438 & 464 & 480 & 492 & 512 & 520 \\
\textbf{Backward} & 285 & 286 & 319 & 334 & 363 & 361 & 371 \\
\bottomrule
\end{tabular}
\caption{Quantization overhead, forward, and backward MSE for different configurations and group sizes.}
\end{table}
We also find that applying per-group quantization to linear layers is harmful for speed, while not improving the accuracy by too much. Even with system optimizations, this finer-grained method reduces the speed of linear layers by about 40\% (when group size = 128, speed is reduced from 1300 TFlops to about 900 TFlops), greatly reducing the benefit of FP8. On the other hand, the accuracy gain is only marginal. Therefore we do apply per-tensor quantization for linear layers.
\section{Efficiency Breakdown for optimizer and activation}
We break down the improvement of each component and present the result in this table. Quantizing the optimizer states to FP8 is helpful in reducing the memory footprint, and can potentially increase the batch size to fully utilize the GPU. Quantizing the activation to FP8 can accelerate the training, and this number can be further improved by increasing the maximum batch size.
\begin{table}[H]
\centering
\caption{Llama-2-7B Training Memory and Throughput Analysis}
\begin{tabular}{lccccccc}
\toprule
\textbf{Configuration} & \textbf{Optimizer} & \textbf{Activation} & \textbf{Peak} & \textbf{Max BS} & \textbf{Throughput} & \textbf{Speedup} \\
\midrule
\multicolumn{6}{l}{\textbf{Llama-2-7B on 2 GPU}} \\
BF16 & 26.2 GB & 25.8 GB & OOM & 1 & 6130 tokens/s & 1.00x \\
FP8 Optimizer & 6.4 GB & 25.8 GB & 61.9 GB & 2 & 7258 tokens/s & 1.18x \\
FP8 Optimizer + FP8 Activation & 6.4 GB & 16.9 GB & 35.6 GB & 4 & 11351 tokens/s & 1.85x \\
\midrule
\multicolumn{6}{l}{\textbf{Llama-2-7B on 4 GPU}} \\
BF16 & 13.1 GB & 25.8 GB & 55.1 GB & 2 & 7730 tokens/s & 1.00x \\
FP8 Optimizer & 3.2 GB & 16.9 GB & 44.7 GB & 4 & 8384 tokens/s & 1.08x \\
FP8 Optimizer + FP8 & 3.2 GB & 16.9 GB & 35.6 GB & 4 & 11257 tokens/s & 1.45x \\
\midrule
\multicolumn{6}{l}{\textbf{Llama-2-7B on 8 GPU}} \\
BF16 & 6.5 GB & 25.8 GB & 41.2 GB & 4 & 8238 tokens/s & 1.00x \\
FP8 Optimizer & 1.6 GB & 16.9 GB & 36.1 GB & 4 & 8444 tokens/s & 1.02x \\
FP8 Optimizer + FP8 Activation & 1.6 GB & 16.9 GB & 27.0 GB & 8 & 11241 tokens/s & 1.36x \\
\bottomrule
\end{tabular}
\end{table}
\section{Overhead introduced by Group Scaling}
We show that the max reduction on each 1×G “can be seamlessly fused with the previous operation, adding minimal overhead”. Take RMSNorm as an example, the speed of the RMSNorm kernel only increases by less than 5\% when quantization is fused with it. Therefore, this adds only minimal overhead to the kernel.
\begin{table}[!h]
\centering
\caption{Performance Comparison: With and Without Quantization}
\begin{tabular}{lcc}
\toprule
\textbf{Matrix Size} & \textbf{With Quantize} & \textbf{Without Quantize} \\
\midrule
(4096, 1024) & 10.1 ms & 8.7 ms \\
(4096, 2048) & 29.6 ms & 28.7 ms \\
(4096, 4096) & 55.6 ms & 53.2 ms \\
(4096, 8192) & 122.6 ms & 117.5 ms \\
\bottomrule
\end{tabular}
\end{table}
\section{Understand how each component influences the model performance}
We separately quantize the optimizer and activations to better understand how each component influences the model performance. We find that quantizing activation incurs almost no degradation due to our well-designed quantization flow, and the main degradation comes from optimizer quantization. We choose to quantize both to maximize memory saving.
\begin{table}[!h]
\centering
\begin{tabular}{lcccccc}
\toprule
& \textbf{Train PPL} & \textbf{COPA} & \textbf{ARC (Easy)} & \textbf{SciQ} & \textbf{HellaSwag} & \textbf{Average} \\
\midrule
\textbf{BF16} & 2.995 & 60.0 & 45.6 & 67.3 & 33.7 & 51.6 \\
\textbf{FP8 O + A} & 3.008 & 61.0 & 44.2 & 67.6 & 33.7 & 51.5 \\
\textbf{FP8 O} & 2.998 (+0.003) & 60.0 (+0.0) & 44.5 & 66.0 & 34.1 & 51.1 \\
\textbf{FP8 A} & 2.999 (+0.004) & 63.0 (+3.0) & 44.0 & 68.4 & 34.1 & 52.3 \\
\bottomrule
\end{tabular}
\caption{Performance comparison across different metrics and configurations.}
\label{tab:performance}
\end{table}
\begin{thebibliography}{73}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Adler et~al.(2024)Adler, Agarwal, Aithal, Anh, Bhattacharya, Brundyn, Casper, Catanzaro, Clay, Cohen, et~al.]{adler2024nemotron}
Bo~Adler, Niket Agarwal, Ashwath Aithal, Dong~H Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, et~al.
\newblock Nemotron-4 340b technical report.
\newblock \emph{arXiv preprint arXiv:2406.11704}, 2024.
\bibitem[Anil et~al.(2019)Anil, Gupta, Koren, and Singer]{anil2019sm3}
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer.
\newblock Memory efficient adaptive optimization.
\newblock \emph{Advances in Neural Information Processing Systems}, 32, 2019.
\bibitem[Ba(2016)]{ba2016layer}
Jimmy~Lei Ba.
\newblock Layer normalization.
\newblock \emph{arXiv preprint arXiv:1607.06450}, 2016.
\bibitem[Cai et~al.(2020)Cai, Gan, Zhu, and Han]{cai2020tinytl}
Han Cai, Chuang Gan, Ligeng Zhu, and Song Han.
\newblock Tinytl: Reduce activations, not trainable parameters for efficient on-device learning.
\newblock \emph{arXiv preprint arXiv:2007.11622}, 2020.
\bibitem[Chen et~al.(2020)Chen, Gai, Yao, Mahoney, and Gonzalez]{chen2020statistical}
Jianfei Chen, Yu~Gai, Zhewei Yao, Michael~W Mahoney, and Joseph~E Gonzalez.
\newblock A statistical framework for low-bitwidth training of deep neural networks.
\newblock \emph{Advances in neural information processing systems}, 33:\penalty0 883--894, 2020.
\bibitem[Chen et~al.(2021)Chen, Zheng, Yao, Wang, Stoica, Mahoney, and Gonzalez]{chen2021actnn}
Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez.
\newblock Actnn: Reducing training memory footprint via 2-bit activation compressed training.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 1803--1813. PMLR, 2021.
\bibitem[Chen et~al.(2023)Chen, Li, Dong, Zhang, He, Wang, Zhao, and Lin]{chen2023sharegpt4v}
Lin Chen, Jisong Li, Xiaoyi Dong, Pan Zhang, Conghui He, Jiaqi Wang, Feng Zhao, and Dahua Lin.
\newblock Sharegpt4v: Improving large multi-modal models with better captions.
\newblock \emph{arXiv preprint arXiv:2311.12793}, 2023.
\bibitem[Chen et~al.(2024)Chen, Liang, Huang, Real, Wang, Pham, Dong, Luong, Hsieh, Lu, et~al.]{chen2024lion}
Xiangning Chen, Chen Liang, Da~Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et~al.
\newblock Symbolic discovery of optimization algorithms.
\newblock \emph{Advances in neural information processing systems}, 36, 2024.
\bibitem[Clark et~al.(2018)Clark, Cowhey, Etzioni, Khot, Sabharwal, Schoenick, and Tafjord]{clark2018arc}
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord.
\newblock Think you have solved question answering? try arc, the ai2 reasoning challenge.
\newblock \emph{arXiv preprint arXiv:1803.05457}, 2018.
\bibitem[Cobbe et~al.(2021)Cobbe, Kosaraju, Bavarian, Chen, Jun, Kaiser, Plappert, Tworek, Hilton, Nakano, et~al.]{cobbe2021gsm8k}
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et~al.
\newblock Training verifiers to solve math word problems.
\newblock \emph{arXiv preprint arXiv:2110.14168}, 2021.
\bibitem[Dai et~al.(2021)Dai, Venkatesan, Ren, Zimmer, Dally, and Khailany]{dai2021vsquant}
Steve Dai, Rangha Venkatesan, Mark Ren, Brian Zimmer, William Dally, and Brucek Khailany.
\newblock Vs-quant: Per-vector scaled quantization for accurate low-precision neural network inference.
\newblock \emph{Proceedings of Machine Learning and Systems}, 3:\penalty0 873--884, 2021.
\bibitem[Davies et~al.(2021)Davies, Veli{\v{c}}kovi{\'c}, Buesing, Blackwell, Zheng, Toma{\v{s}}ev, Tanburn, Battaglia, Blundell, Juh{\'a}sz, et~al.]{davies2021deepmind}
Alex Davies, Petar Veli{\v{c}}kovi{\'c}, Lars Buesing, Sam Blackwell, Daniel Zheng, Nenad Toma{\v{s}}ev, Richard Tanburn, Peter Battaglia, Charles Blundell, Andr{\'a}s Juh{\'a}sz, et~al.
\newblock Advancing mathematics by guiding human intuition with ai.
\newblock \emph{Nature}, 600\penalty0 (7887):\penalty0 70--74, 2021.
\bibitem[Dettmers et~al.(2021)Dettmers, Lewis, Shleifer, and Zettlemoyer]{dettmers2021adam8bit}
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer.
\newblock 8-bit optimizers via block-wise quantization.
\newblock \emph{arXiv preprint arXiv:2110.02861}, 2021.
\bibitem[Dubey et~al.(2024)Dubey, Jauhri, Pandey, Kadian, Al-Dahle, Letman, Mathur, Schelten, Yang, Fan, et~al.]{dubey2024llama3}
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et~al.
\newblock The llama 3 herd of models.
\newblock \emph{arXiv preprint arXiv:2407.21783}, 2024.
\bibitem[Fishman et~al.(2024)Fishman, Chmiel, Banner, and Soudry]{fishman2024fp8trillion}
Maxim Fishman, Brian Chmiel, Ron Banner, and Daniel Soudry.
\newblock Scaling fp8 training to trillion-token llms.
\newblock \emph{arXiv preprint arXiv:2409.12517}, 2024.
\bibitem[Fu et~al.(2024)Fu, Dai, Luo, Li, Ren, Zhang, Wang, Zhou, Shen, Zhang, et~al.]{fu2024videomme}
Chaoyou Fu, Yuhan Dai, Yondong Luo, Lei Li, Shuhuai Ren, Renrui Zhang, Zihan Wang, Chenyu Zhou, Yunhang Shen, Mengdan Zhang, et~al.
\newblock Video-mme: The first-ever comprehensive evaluation benchmark of multi-modal llms in video analysis.
\newblock \emph{arXiv preprint arXiv:2405.21075}, 2024.
\bibitem[Gao et~al.(2020)Gao, Biderman, Black, Golding, Hoppe, Foster, Phang, He, Thite, Nabeshima, et~al.]{gao2020pile}
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et~al.
\newblock The pile: An 800gb dataset of diverse text for language modeling.
\newblock \emph{arXiv preprint arXiv:2101.00027}, 2020.
\bibitem[Gordon et~al.(2012)Gordon, Kozareva, and Roemmele]{gordon-etal-2012-copa}
Andrew Gordon, Zornitsa Kozareva, and Melissa Roemmele.
\newblock {S}em{E}val-2012 task 7: Choice of plausible alternatives: An evaluation of commonsense causal reasoning.
\newblock In Eneko Agirre, Johan Bos, Mona Diab, Suresh Manandhar, Yuval Marton, and Deniz Yuret (eds.), \emph{*{SEM} 2012: The First Joint Conference on Lexical and Computational Semantics {--} Volume 1: Proceedings of the main conference and the shared task, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation ({S}em{E}val 2012)}, pp.\ 394--398, Montr{\'e}al, Canada, 7-8 June 2012. Association for Computational Linguistics.
\newblock URL \url{https://aclanthology.org/S12-1052}.
\bibitem[Goyal et~al.(2017)Goyal, Khot, Summers-Stay, Batra, and Parikh]{goyal2017vqav2}
Yash Goyal, Tejas Khot, Douglas Summers-Stay, Dhruv Batra, and Devi Parikh.
\newblock Making the v in vqa matter: Elevating the role of image understanding in visual question answering.
\newblock In \emph{Proceedings of the IEEE conference on computer vision and pattern recognition}, pp.\ 6904--6913, 2017.
\bibitem[Groeneveld et~al.(2024)Groeneveld, Beltagy, Walsh, Bhagia, Kinney, Tafjord, Jha, Ivison, Magnusson, Wang, et~al.]{groeneveld2024olmo}
Dirk Groeneveld, Iz~Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya~Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et~al.
\newblock Olmo: Accelerating the science of language models.
\newblock \emph{arXiv preprint arXiv:2402.00838}, 2024.
\bibitem[Gurari et~al.(2018)Gurari, Li, Stangl, Guo, Lin, Grauman, Luo, and Bigham]{gurari2018vizwiz}
Danna Gurari, Qing Li, Abigale~J Stangl, Anhong Guo, Chi Lin, Kristen Grauman, Jiebo Luo, and Jeffrey~P Bigham.
\newblock Vizwiz grand challenge: Answering visual questions from blind people.
\newblock In \emph{Proceedings of the IEEE conference on computer vision and pattern recognition}, pp.\ 3608--3617, 2018.
\bibitem[Hendrycks \& Gimpel(2016)Hendrycks and Gimpel]{hendrycks2016gelu}
Dan Hendrycks and Kevin Gimpel.
\newblock Gaussian error linear units (gelus).
\newblock \emph{arXiv preprint arXiv:1606.08415}, 2016.
\bibitem[Hoffmann et~al.(2022)Hoffmann, Borgeaud, Mensch, Buchatskaya, Cai, Rutherford, Casas, Hendricks, Welbl, Clark, et~al.]{hoffmann2022chinchillia}
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de~Las Casas, Lisa~Anne Hendricks, Johannes Welbl, Aidan Clark, et~al.
\newblock Training compute-optimal large language models.
\newblock \emph{arXiv preprint arXiv:2203.15556}, 2022.
\bibitem[Hudson \& Manning(2019)Hudson and Manning]{hudson2019gqa}
Drew~A Hudson and Christopher~D Manning.
\newblock Gqa: A new dataset for real-world visual reasoning and compositional question answering.
\newblock In \emph{Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, pp.\ 6700--6709, 2019.
\bibitem[Kalamkar et~al.(2019)Kalamkar, Mudigere, Mellempudi, Das, Banerjee, Avancha, Vooturi, Jammalamadaka, Huang, Yuen, et~al.]{kalamkar2019bf16study}
Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma~Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, et~al.
\newblock A study of bfloat16 for deep learning training.
\newblock \emph{arXiv preprint arXiv:1905.12322}, 2019.
\bibitem[Kingma(2014)]{kingma2014adam}
Diederik~P Kingma.
\newblock Adam: A method for stochastic optimization.
\newblock \emph{arXiv preprint arXiv:1412.6980}, 2014.
\bibitem[Li et~al.(2024)Li, Chen, and Zhu]{li2024optimizer4bit}
Bingrui Li, Jianfei Chen, and Jun Zhu.
\newblock Memory efficient optimizers with 4-bit states.
\newblock \emph{Advances in Neural Information Processing Systems}, 36, 2024.
\bibitem[Li et~al.(2023{\natexlab{a}})Li, Wang, Wang, Ge, Ge, and Shan]{li2023seed}
Bohao Li, Rui Wang, Guangzhi Wang, Yuying Ge, Yixiao Ge, and Ying Shan.
\newblock Seed-bench: Benchmarking multimodal llms with generative comprehension.
\newblock \emph{arXiv preprint arXiv:2307.16125}, 2023{\natexlab{a}}.
\bibitem[Li et~al.(2023{\natexlab{b}})Li, Du, Zhou, Wang, Zhao, and Wen]{li2023pope}
Yifan Li, Yifan Du, Kun Zhou, Jinpeng Wang, Wayne~Xin Zhao, and Ji-Rong Wen.
\newblock Evaluating object hallucination in large vision-language models.
\newblock \emph{arXiv preprint arXiv:2305.10355}, 2023{\natexlab{b}}.
\bibitem[Lin et~al.(2022)Lin, Zhu, Chen, Wang, Gan, and Han]{lin2022device256kb}
Ji~Lin, Ligeng Zhu, Wei-Ming Chen, Wei-Chen Wang, Chuang Gan, and Song Han.
\newblock On-device training under 256kb memory.
\newblock \emph{Advances in Neural Information Processing Systems}, 35:\penalty0 22941--22954, 2022.
\bibitem[Lin et~al.(2024)Lin, Yin, Ping, Molchanov, Shoeybi, and Han]{lin2024vila}
Ji~Lin, Hongxu Yin, Wei Ping, Pavlo Molchanov, Mohammad Shoeybi, and Song Han.
\newblock Vila: On pre-training for visual language models.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, pp.\ 26689--26699, 2024.
\bibitem[Liu et~al.(2022{\natexlab{a}})Liu, Zheng, Wang, Cen, Chen, Han, Chen, Liu, Tang, Gonzalez, et~al.]{liu2022gact}
Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu~Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, et~al.
\newblock Gact: Activation compressed training for generic network architectures.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 14139--14152. PMLR, 2022{\natexlab{a}}.
\bibitem[Liu et~al.(2022{\natexlab{b}})Liu, Cheng, Huang, Xing, and Shen]{liu2022NU2Qnonuniform}
Zechun Liu, Kwang-Ting Cheng, Dong Huang, Eric~P Xing, and Zhiqiang Shen.
\newblock Nonuniform-to-uniform quantization: Towards accurate quantization via generalized straight-through estimation.
\newblock In \emph{Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, pp.\ 4942--4952, 2022{\natexlab{b}}.
\bibitem[Loshchilov(2017)]{loshchilov2017adamw}
I~Loshchilov.
\newblock Decoupled weight decay regularization.
\newblock \emph{arXiv preprint arXiv:1711.05101}, 2017.
\bibitem[Merity et~al.(2016)Merity, Xiong, Bradbury, and Socher]{merity2016wikitext}
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher.
\newblock Pointer sentinel mixture models.
\newblock \emph{arXiv preprint arXiv:1609.07843}, 2016.
\bibitem[Micikevicius et~al.(2017)Micikevicius, Narang, Alben, Diamos, Elsen, Garcia, Ginsburg, Houston, Kuchaiev, Venkatesh, et~al.]{micikevicius2017mixedprecision}
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et~al.
\newblock Mixed precision training.
\newblock \emph{arXiv preprint arXiv:1710.03740}, 2017.
\bibitem[Micikevicius et~al.(2022)Micikevicius, Stosic, Burgess, Cornea, Dubey, Grisenthwaite, Ha, Heinecke, Judd, Kamalu, et~al.]{micikevicius2022fp8format}
Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenthwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, et~al.
\newblock Fp8 formats for deep learning.
\newblock \emph{arXiv preprint arXiv:2209.05433}, 2022.
\bibitem[Mishra et~al.(2022)Mishra, Mitra, Varshney, Sachdeva, Clark, Baral, and Kalyan]{mishra2022numglue}
Swaroop Mishra, Arindam Mitra, Neeraj Varshney, Bhavdeep Sachdeva, Peter Clark, Chitta Baral, and Ashwin Kalyan.
\newblock Numglue: A suite of fundamental yet challenging mathematical reasoning tasks.
\newblock \emph{arXiv preprint arXiv:2204.05660}, 2022.
\bibitem[Novikov et~al.(2023)Novikov, Bershatsky, Gusak, Shonenkov, Dimitrov, and Oseledets]{novikov2023fewbitbackward}
Georgii~Sergeevich Novikov, Daniel Bershatsky, Julia Gusak, Alex Shonenkov, Denis~Valerievich Dimitrov, and Ivan Oseledets.
\newblock Few-bit backward: Quantized gradients of activation functions for memory footprint reduction.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 26363--26381. PMLR, 2023.
\bibitem[NVIDIA(2024{\natexlab{a}})]{nvidia_h100}
NVIDIA.
\newblock Nvidia h100 tensor core gpu, 2024{\natexlab{a}}.
\newblock URL \url{https://www.nvidia.com/en-us/data-center/h100/}.
\newblock Accessed: 2024-09-19.
\bibitem[NVIDIA(2024{\natexlab{b}})]{nvidia_transformerengine}
NVIDIA.
\newblock Transformerengine: An efficient library for training transformer models, 2024{\natexlab{b}}.
\newblock URL \url{https://github.com/NVIDIA/TransformerEngine}.
\newblock Accessed: 2024-09-19.
\bibitem[{Open Compute Project}(2023)]{OCP8bit2023}
{Open Compute Project}.
\newblock Ocp 8-bit floating point specification (ofp8), revision 1.0, December 2023.
\newblock URL \url{https://www.opencompute.org/documents/ocp-8-bit-floating-point-specification-ofp8-revision-1-0-2023-12-01-pdf-1}.
\newblock Accessed: 2024-10-09.
\bibitem[Patel et~al.(2021)Patel, Bhattamishra, and Goyal]{patel2021svamp}
Arkil Patel, Satwik Bhattamishra, and Navin Goyal.
\newblock Are nlp models really able to solve simple math word problems?
\newblock \emph{arXiv preprint arXiv:2103.07191}, 2021.
\bibitem[Peng et~al.(2023)Peng, Wu, Wei, Zhao, Yang, Liu, Xiong, Yang, Ni, Hu, et~al.]{peng2023fp8lm}
Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze~Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, et~al.
\newblock Fp8-lm: Training fp8 large language models.
\newblock \emph{arXiv preprint arXiv:2310.18313}, 2023.
\bibitem[Raffel et~al.(2020)Raffel, Shazeer, Roberts, Lee, Narang, Matena, Zhou, Li, and Liu]{raffel2020c4}
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter~J Liu.
\newblock Exploring the limits of transfer learning with a unified text-to-text transformer.
\newblock \emph{Journal of machine learning research}, 21\penalty0 (140):\penalty0 1--67, 2020.
\bibitem[Rasley et~al.(2020)Rasley, Rajbhandari, Ruwase, and He]{rasley2020deepspeed}
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He.
\newblock Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters.
\newblock In \emph{Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery \& Data Mining}, pp.\ 3505--3506, 2020.
\bibitem[Shah et~al.(2024)Shah, Bikshandi, Zhang, Thakkar, Ramani, and Dao]{shah2024flashattention3}
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao.
\newblock Flashattention-3: Fast and accurate attention with asynchrony and low-precision.
\newblock \emph{arXiv preprint arXiv:2407.08608}, 2024.
\bibitem[Shazeer(2020)]{shazeer2020glu}
Noam Shazeer.
\newblock Glu variants improve transformer.
\newblock \emph{arXiv preprint arXiv:2002.05202}, 2020.
\bibitem[Shazeer \& Stern(2018)Shazeer and Stern]{shazeer2018adafactor}
Noam Shazeer and Mitchell Stern.
\newblock Adafactor: Adaptive learning rates with sublinear memory cost.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 4596--4604. PMLR, 2018.
\bibitem[Shoeybi et~al.(2019)Shoeybi, Patwary, Puri, LeGresley, Casper, and Catanzaro]{shoeybi2019megatron}
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.
\newblock Megatron-lm: Training multi-billion parameter language models using model parallelism.
\newblock \emph{arXiv preprint arXiv:1909.08053}, 2019.
\bibitem[Singh et~al.(2019)Singh, Natarajan, Shah, Jiang, Chen, Batra, Parikh, and Rohrbach]{singh2019textvqa}
Amanpreet Singh, Vivek Natarajan, Meet Shah, Yu~Jiang, Xinlei Chen, Dhruv Batra, Devi Parikh, and Marcus Rohrbach.
\newblock Towards vqa models that can read.
\newblock In \emph{Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, pp.\ 8317--8326, 2019.
\bibitem[Smith et~al.(2022)Smith, Patwary, Norick, LeGresley, Rajbhandari, Casper, Liu, Prabhumoye, Zerveas, Korthikanti, et~al.]{smith2022nlg530b}
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et~al.
\newblock Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model.
\newblock \emph{arXiv preprint arXiv:2201.11990}, 2022.
\bibitem[Soldaini et~al.(2024)Soldaini, Kinney, Bhagia, Schwenk, Atkinson, Authur, Bogin, Chandu, Dumas, Elazar, et~al.]{soldaini2024dolma}
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, et~al.
\newblock Dolma: An open corpus of three trillion tokens for language model pretraining research.
\newblock \emph{arXiv preprint arXiv:2402.00159}, 2024.
\bibitem[Team et~al.(2024)Team, Mesnard, Hardin, Dadashi, Bhupatiraju, Pathak, Sifre, Rivi{\`e}re, Kale, Love, et~al.]{team2024gemma}
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivi{\`e}re, Mihir~Sanjay Kale, Juliette Love, et~al.
\newblock Gemma: Open models based on gemini research and technology.
\newblock \emph{arXiv preprint arXiv:2403.08295}, 2024.
\bibitem[Touvron et~al.(2023)Touvron, Lavril, Izacard, Martinet, Lachaux, Lacroix, Rozi{\`e}re, Goyal, Hambro, Azhar, et~al.]{touvron2023llama}
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timoth{\'e}e Lacroix, Baptiste Rozi{\`e}re, Naman Goyal, Eric Hambro, Faisal Azhar, et~al.
\newblock Llama: Open and efficient foundation language models.
\newblock \emph{arXiv preprint arXiv:2302.13971}, 2023.
\bibitem[Wang et~al.(2022)Wang, Yuan, Rimanic, He, Dao, Chen, Re, and Zhang]{wang2022aqsgd}
Jue Wang, Binhang Yuan, Luka Rimanic, Yongjun He, Tri Dao, Beidi Chen, Christopher Re, and Ce~Zhang.
\newblock Fine-tuning language models over slow networks using activation compression with guarantees.
\newblock \emph{arXiv preprint arXiv:2206.01299}, 2022.
\bibitem[Wang et~al.(2018)Wang, Choi, Brand, Chen, and Gopalakrishnan]{wang2018trainingfp8}
Naigang Wang, Jungwook Choi, Daniel Brand, Chia-Yu Chen, and Kailash Gopalakrishnan.
\newblock Training deep neural networks with 8-bit floating point numbers.
\newblock \emph{Advances in neural information processing systems}, 31, 2018.
\bibitem[Wang \& Kang(2023)Wang and Kang]{wang2023DDQgradient}
Shuai Wang and Yi~Kang.
\newblock Gradient distribution-aware int8 training for neural networks.
\newblock \emph{Neurocomputing}, 541:\penalty0 126269, 2023.
\bibitem[Welbl et~al.(2017)Welbl, Liu, and Gardner]{welbl2017sciq}
Johannes Welbl, Nelson~F Liu, and Matt Gardner.
\newblock Crowdsourcing multiple choice science questions.
\newblock \emph{arXiv preprint arXiv:1707.06209}, 2017.
\bibitem[Wortsman et~al.(2023)Wortsman, Dettmers, Zettlemoyer, Morcos, Farhadi, and Schmidt]{wortsman2023int8clip}
Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, and Ludwig Schmidt.
\newblock Stable and low-precision training for large-scale vision-language models.
\newblock \emph{Advances in Neural Information Processing Systems}, 36:\penalty0 10271--10298, 2023.
\bibitem[Xi et~al.(2023)Xi, Li, Chen, and Zhu]{xi2023int4training}
Haocheng Xi, Changhao Li, Jianfei Chen, and Jun Zhu.
\newblock Training transformers with 4-bit integers.
\newblock \emph{Advances in Neural Information Processing Systems}, 36:\penalty0 49146--49168, 2023.
\bibitem[Xi et~al.(2024)Xi, Chen, Zhao, Zheng, Chen, and Zhu]{xi2024jetfire}
Haocheng Xi, Yuxiang Chen, Kang Zhao, Kaijun Zheng, Jianfei Chen, and Jun Zhu.
\newblock Jetfire: Efficient and accurate transformer pretraining with int8 data flow and per-block quantization.
\newblock \emph{arXiv preprint arXiv:2403.12422}, 2024.
\bibitem[Xu et~al.(2024)Xu, Feng, Shao, Ashby, Shen, Jin, Cheng, Wang, and Huang]{xu2024visionflan}
Zhiyang Xu, Chao Feng, Rulin Shao, Trevor Ashby, Ying Shen, Di~Jin, Yu~Cheng, Qifan Wang, and Lifu Huang.
\newblock Vision-flan: Scaling human-labeled tasks in visual instruction tuning.
\newblock \emph{arXiv preprint arXiv:2402.11690}, 2024.
\bibitem[Yang et~al.(2024)Yang, Kim, and Kim]{yang2024gluspike}
Jaewoo Yang, Hayun Kim, and Younghoon Kim.
\newblock Mitigating quantization errors due to activation spikes in glu-based llms.
\newblock \emph{arXiv preprint arXiv:2405.14428}, 2024.
\bibitem[Yue et~al.(2023)Yue, Qu, Zhang, Fu, Huang, Sun, Su, and Chen]{yue2023mammoth}
Xiang Yue, Xingwei Qu, Ge~Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu~Su, and Wenhu Chen.
\newblock Mammoth: Building math generalist models through hybrid instruction tuning.
\newblock \emph{arXiv preprint arXiv:2309.05653}, 2023.
\bibitem[Yue et~al.(2024)Yue, Ni, Zhang, Zheng, Liu, Zhang, Stevens, Jiang, Ren, Sun, et~al.]{yue2024mmmu}
Xiang Yue, Yuansheng Ni, Kai Zhang, Tianyu Zheng, Ruoqi Liu, Ge~Zhang, Samuel Stevens, Dongfu Jiang, Weiming Ren, Yuxuan Sun, et~al.
\newblock Mmmu: A massive multi-discipline multimodal understanding and reasoning benchmark for expert agi.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, pp.\ 9556--9567, 2024.
\bibitem[Zellers et~al.(2019)Zellers, Holtzman, Bisk, Farhadi, and Choi]{zellers2019hellaswag}
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.
\newblock Hellaswag: Can a machine really finish your sentence?
\newblock \emph{arXiv preprint arXiv:1905.07830}, 2019.
\bibitem[Zhang \& Sennrich(2019)Zhang and Sennrich]{zhang2019root}
Biao Zhang and Rico Sennrich.
\newblock Root mean square layer normalization.
\newblock \emph{Advances in Neural Information Processing Systems}, 32, 2019.
\bibitem[Zhang et~al.(2024)Zhang, Huang, Zhang, Wei, Zhu, and Chen]{zhang2024sageattention2}
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen.
\newblock Sageattention2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization, 2024.
\newblock URL \url{https://arxiv.org/abs/2411.10958}.
\bibitem[Zhang et~al.(2025)Zhang, Wei, Zhang, Zhu, and Chen]{zhang2025sageattention}
Jintao Zhang, Jia Wei, Pengle Zhang, Jun Zhu, and Jianfei Chen.
\newblock Sageattention: Accurate 8-bit attention for plug-and-play inference acceleration.
\newblock In \emph{International Conference on Learning Representations (ICLR)}, 2025.
\bibitem[Zhao et~al.(2024)Zhao, Zhang, Chen, Wang, Anandkumar, and Tian]{zhao2024galore}
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian.
\newblock Galore: Memory-efficient llm training by gradient low-rank projection.
\newblock \emph{arXiv preprint arXiv:2403.03507}, 2024.
\bibitem[Zhou et~al.(2017)Zhou, Wang, Wen, He, and Zou]{zhou2017balancedquant}
Shu-Chang Zhou, Yu-Zhi Wang, He~Wen, Qin-Yao He, and Yu-Heng Zou.
\newblock Balanced quantization: An effective and efficient approach to quantized neural networks.
\newblock \emph{Journal of Computer Science and Technology}, 32:\penalty0 667--682, 2017.
\bibitem[Zhu et~al.(2020)Zhu, Gong, Yu, Liu, Wang, Li, Yang, and Yan]{zhu2020towardsint8}
Feng Zhu, Ruihao Gong, Fengwei Yu, Xianglong Liu, Yanfei Wang, Zhelong Li, Xiuqi Yang, and Junjie Yan.
\newblock Towards unified int8 training for convolutional neural network.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, pp.\ 1969--1979, 2020.
\end{thebibliography}