COAT: Compressing Optimizer states and Activations for Memory-Efficient FP8 Training
Authors: Haocheng Xi
Han Cai
Ligeng Zhu
Yao Lu
Kurt Keutzer
Jianfei Chen
Song Han
Venue: ICLR 2025 Poster (2025)
Keywords: FP8 training
quantization
low-precision training
memory efficient training
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a memory efficient fp8 training method for large language models.

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 (**C**ompressing **O**ptimizer States and **A**ctivations for FP8 **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) **Dynamic Range Expansion**, which aligns optimizer state distributions more closely with the FP8 representation range, thereby reducing quantization error, and (2) **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. Code will be released upon publication.

Published as a conference paper at ICLR 2025
COAT: COMPRESSING OPTIMIZER STATES AND A CTIVATION FOR MEMORY-EFFICIENT FP8 TRAINING
Haocheng Han Cai², Han
Ligeng Zhu?, Yao Lu?, Kurt Keutzer', Jianfei Chen | Song Han
University of California, Berkeley NVIDIA 3 MIT
Tsinghua University
& https://nvlaba: gi 1o/CCAT/
ABSTRACT
FP8 training has emerged 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 (Compressing Optimizer States and Activations for FP8 Training). novel FP8 training frame- work designed to significantly reduce memory footprint when training large mod- els. COAT addresses current limitations through two key innovations: (1) Dy- namic Range Expansion, which aligns optimizer state distributions more closely with the FP8 representation range, thereby reducing quantization error, and (2) HixeedGraaulae
A Activation Quantization, which optimizes activation mem- ory using a combination of per-tensor and per-group quantization strategies. Ex- periments demonstrate that COAT effectively reduces end-to-end training memory footprint by 1.54x 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.43x end-to-end training speedup compared to BF16, performing on par with or surpassing Trans- formerEngine's i speedup. COAT enables efficient full -ull-parrmeter training of large models on fewer GPUs, and facilitates doubling the batch size in distributed train- ing settings, providing a practical solution for scaling large-scale model training. The code is available at https://github..om// Vlabs/COAT
FP8
FPS Per-Group Quant
Optimizer
Weight Gradient
Weight Weight Gradient Activation Optimizer
BF16
Per-Tensor Quant
Activation
Weight
FP8
LGB 70L0
BMSNemm
RMSNom
RMSNurn
RMSSNom
56.0
512
HuiiAmuriie
Anicetion
Arieasiay
14.0
0.0
19716
PP81 Muurs
(a) TramffomneEnggin
(b) COAT
(6) Memory Decompositon
Figure 1: (a.b) Comparing the quantization flow of Transformer Engine and COAT. Both the opti- mizer states and activations are quantized to FPS in COAT. (c) End-to-end Ger memory com- parison when training Llama-2-133 23BB on 8X0 H100 using FSDP.
Parr of the work done during an internship at NVIDIA.
1 "Corrsponndig author

---

Published as a conference paper at ICLR 2025
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 (Dubey et al. 2024] Adler et al. 2024 Team et al. 2024 Lin et al., 2024). However, the training of such models, which often comprise billions of parameters, demands substantial com- putational resources and memory. This presents substantial challenges, making the training of these foundation models very challenging (Smith et al.
2022 Hoffmann et al. (022).
Low-precision training has emerged as promising approach to make FMs training more effi cient MMciievvicius et al. 2017] Wang et al. 2018 Zhu et al., 2020, Xi et al. [2023] Wortsman et al.) 2023 Xi ct al. (024). 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 (Kalamkar et al. 2019; Micikevicius al. (017) is the most prevalent
low-precision method, and is widely adopted in large-scale training frameworks like DeepSpeed (Rasley et al. 2020) and Megatron-LMM (Shoeybi et al. 2019).
With the advent of Nvidia's H100 GPU (NVIDIA 2024a), FP8 training Micikevicius et al. (2022) is emerging as the next-generation low-precision technique. Compared BF16, FP8 training has the potential to (I) double the speed and (2) halve the memory footprint. To achieve practical speedup, Transformer Engine (NVIDIA] 20246) 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 I FP8-LM (Peng et al.) (023) 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 activations and still leave the optimizer's '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-oree -oddrr momentum is more sensitive to quan- tization than first-order momentum (Fishman al.] (004),, and activations' large spikes also make them hard to quantize to FP8 (Yang et al.] 2024). This potential accuracy degradation makes them missing a crucial opportunity optimize memory further.
In this work, we propose COAT: Compressing Optimizer states and Activations for memory- efficient FP8 Training to address the aforementioned issue. COAT significantly reduces the over- all memory footprint by quantizing optimizer states and activations into FP8. For optimizer states, we observe that FP8 format's representation range under-utilized when quantizing them, as il- lustrated in Figure 2 a).. To address this, we introduce a novel Dynamic Range Expansion method which adjusts the distribution of optimizer states better fit within the FP8 range, thereby minimiz- ing quantization error. For activations, we propose Mixed fixee-Granularity Activation Quantization to achieve efficient and accurate quantization. We apply ine-grained quantization to non-linear layers and apply per-tensor quantization to linear layers. Per-tensor guantization for matrix multiplications is more efficient and better suited for TensorCores, while fine- muue-graindd quantization helps maintain accuracy. These two approaches tackle high memory consumption while ensuring minimal perfor- mance degradation. We provide an overview of COAT in Figure ( for demonstration.
We demonstrate the accurate performance of COAT training. LLM fine-tuning, and VLM training. COAT achieves nearly lossless performance on all of these tasks. For efficiency results, COAT achieves 1.54 X end-to-end memory reduction compared with BF16, and 1.43x end-to-end training speed up on Llama 7B, 13B, and 30B models compared to BF16. COAT also doubles the batch size in all realistic distributed training settings, which is crucial for higher speedup and support for longer context length, leading to more efficient training process for large-scale models.
wide range of tasks, including LLM pre-
2 RELATED WORK
Low-precision Training Low precision training (Wang et al. 2018: Chen et al. 2020: Lin et al.
2022 Wortsman et al. 2023 Xi et al. 2024) has become prominent technique in modern deep leannigg offering reductions in both computational costs and memory requirements. FP16 (half- precision) training (Micikevicius et al. 2017) is the most prevalent low-precision method nowadays.
2

---

Published as a conference paper at ICLR 2025
It introduces loss scaling to address FP16's narrower representation range problem. BF16 training (Kalamkar et al. 2019) 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 computed in FP16 or BF16 precision, while master weights, gradients, and optimizers are stored in FP32.
FP8 training (Fishman et al. 2024 Micikevicius al. 2022) 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) (NVIDIA) 2024b) is the first framework designed for FP8 mixed-precision training that employs FP8 Tensor- core for linear layer calculation. FP8-LM (Peng et al., 2023) extends FP8 quantization 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 mo- mentum in FP16, limiting the full potential of FP8's memory advantages.
Memory Efficient Optimizers While 32-bit optimizer Kingma 2014 Loshchilov 2017) states widely adopted, several research efforts have been made reduce the memory footprint of op- timizer states through quantization. The 8-bit Adam (Dettmers et al. 2021) introduces novel data format called dynamic exponent (DE) for guantization, but the adoption of this new data format lim- its its flexibility. 4-bit Optimizer (Li et al. 2024) further pushes the limit of optimizer quantization 4-bit by addressing the zero-point problem, but restricted to fine-tuning tasks. FP8-LM (Peng al. 2023) quantizes the first-order momentum to FP8 while leaving second-order momentum in FPI6, which limits the overall memory savings, (Fishman et al. 2024) finds that second-order momentum is more sensitive quantization, and proposes to quantize it using E5M2 format.
In addition to quantization, there are other approaches that aim to reduce the memory footprint the optimizer states (Shazeer & Stern 20188] Anil et al. [2019; Chen et al. [2024] Zhao et al. (0024), such as low-rank decomposition and optimizer simplification that only store the first-order momentum.
These methods orthogonal approach.
PRELIMINARIES
Quantization compresses high- gigh-precision tensor to low-precision achieve speedup and save memory footprint, at the cost of lower precision and lower representation range.
FP8 Quantization
FP8 formut consists of two encodings E4M3 and E5M2 (Open Compute Project 2023). E4M3
has higher precision, while F5M2 has a larger representation range. We define H4M3's min and max values as A E4M3
= 2 and EAM: = 448, while E5M2's min and max values are A min = 2 16
ESM2
min
and A E5M2 57344. To quantize an FP32 tensor X into E4M3 precision, we use a quantizer Q(-)
min
to map the tensor into FP8's representation range. This process can be formulated as
X Sx
max (IXrpszl) Sx
X Sx Q(XFF32), where XFPS
A4M3
max
where Sx is the scaling factor, [-] refers cound-to-nearess Quantize into E5M2 follows a similar procedure, where we only replace A max
B4M3 with A E5M2 To map the quantized tensor back to FP32 pre- max
cision, the dequantize operation DQ(+) can be formulated as Xp = (Q((xrs Sx) Sx X
Optimizer Update Rule Optimizers are widely used in deep learning update parameters. The most common gradient-basse optimizer is Adam/Adam (Kingmal 2014 Loshchilov (017), which uses first-order m and second-order momentum achieve better convergence. The update rule of Adam W at time step can be formulated as:
m, = B m + (1 -1)9t-1
B201-1 + -32)9
m, m4
Vt
02
1
3
m
W++1 wt 17
Xwa
+
where m, is the first-order momentum, v is the second-order momentum, g. is the gradient, 8 and 32 are betas of AdamW, 7) is learning rate, is weight decay, and is used to prevent NaN or Inf.
3

---

Published as a conference paper at ICLR 2025
Refore Expansion
Order
Scood Order
Expansion
onder 5k is larger ths Eimt-eolle
First
10a FF
10a
7P8 EAMS Reprosmtaiion Rangs, heavily under unitized
Expecially sccund-onler
100
10a
TP8 Range is well utilioed
Sccond order
10a
5200
10a
10
60
80
100
20
40
5
15
20
Ccoup todra
Grrup Index (b)
Figure 2: (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 well utilized.
(c) Distribution of k for optimizer states. The second order's is larger than the first order's 's k, since the recond-orde momentum's 's dynamic range is smaller.
To perform optimizer state quantization, we adopt per-group quantization for both first-order and second-order momentum, following previous works (Dettmers et al. 2021 Li et al.) 2024). 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 A
DYNAMIC RANGE EXPANSION FOR ACCURATE OPTIMIZER QUANTIZATION
In subsequent sections, explain how COAT utilizes FP8 quantization achieve memory-efficient FP8 training without compromising accuracy. Section focuses on optimizer states quantization, while Section 3 discusses activation quantization.
4.1 UNDERSTANDING THE ISSUE OF CURRENT OPTIMIZER STATES QUANTIZATION METHOD
Under per-group quantization, we find that one significant drawback of current quantization methods that, they can not fully utilize the representation range of FP8 and therefore lead to a large quanti- zation error. Take the E4M3 data format sentable value and minimum representable value is A E4M3 / /ABM3
example, the ratio between E4M3's maximum repre-
B4M3 448 512 = 229376 2 10
Therefore, for quantization group X, if we want to fully utilize the 256 representable value of FP8, we hope the dynamic range of the quantization group X should cover the entire span between A E4M3 and A F4M3 max
To make it more formally, we define dynamic range as the ratio between the maximum absolute value and the minimum absolute value within quantization group X:
Definition (dynamic range) Given a set that consists of G real numbers X = {x1, T2, rch. the dynamic range R is defined as
max([z1),
|2
x
Rx
min [r1l. [2
rG]
where |-| denotes the absolute value.
That to say, E4M3's dynamic range is REAMS = 448 x 512 229376 2 x 105, 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 Figure[jaa. the dynamic range for first-order momentum is typically less than 1e4, and for second-order momentum, it is
AAttalll among them small amount of values represents NaN and Inf.
4

---

Published as a conference paper at ICLR 2025
Quantization error of quantization settings. +Expand means applying our Dynamic Range Expansion method.
Table
ve under different
Bolee Expead Alar Expund
Beterr Expeol
Alle Expeed
FPM Reprovatiaio Raage
F1S Reproseitice Reng
MSE m
Second Order
First Order
E4M3--Epann
E5M2
ESM2+Expand
E4M3
20.10
18.08
25.65
18.16
ExM3-Expand
13.13
12.31
12.43
E5M2
37.02
35.96
40.30
36.00
Figure 3:
Dynamic Range Expansion can
Expund
17.79
15.48
23.84
15.57
better utilize E4M3 representation range.
usually less than lel-both far below the available range of FP8. As a result, a substantial portion of FP8's representational capacity is wasted, leading to large quantization error.
4.2 DYNAMIC RANGE EXPANSION
To address this problem, we introduce a expand function f(-) before quantization to expand the dynamic range of the quantization group and align with E4M3, which be formalized as XFP8, Sx = (XFPs))) The expand function we use is defined as
f(x) sign(>))=)**,
where l is used to control the strength of the expansion. In Appendix, F we prove that it is optimal. For quantization group X, after applying the expand function f to X, the dynamic range becomes
R((X- max(]f(X)) min(If(X)] minn[signi/ xI) min((xX)
(xxx(sign(X)X**)
max ( (Rx)*.
Therefore, when k: > 1, Rxx will enlarged and become closer to the ideal RE4M3 The optimal k satisfy that (Rx)* RE4M3, which means that k: = logR (Rrams). With this optimal k. f(X) can fully utilize the representation range of F4M3, while the original X can only utilize small portion of it. As shown in Figure 2 c), the second-order momentum typically has larger & value (5 ~ 15) compared to the first-order momentum's 's (1 ~ 3). This corresponds to our previous observation, that second-order momentum usually has smaller dynamic range compared with first- order momentum, so it requires larger i to align well with E4M3's s dynamic range.
We calculate i on-the-fly for every optimizer step and for every quantization group for accuracy consideration. When dequantize, apply the inverse of expand function fx) I after de- quantization to recover its original value, which can be expressed as X = = (DQ((88s)
We apply regular quantizer and our dynamic range expansion method to both the first-order momen- tum m and second-order momentum v. As visualized in Figure 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 In Adam optimizer step, as stated Eq. Vo+c the actual effective term for weight update, so we report the MSE of yete to quantity 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 Appendix A more results.
5 MIXED-GRANULARITY ACTIVATION QUANTIZATION
5.1 DECOMPOSE THE ACTIVATION MEMORY FOOTPRINT
In the forward pass of neural networks, activations must be preserved for the backward pass calcu- late gradients. As illustrated in Table non-linear layers such as LayerNorm/RMSN (Ba 2016
Zhang & Sennrich 2019) and Activation Functions (Hendrycks & Gimpel 2016 Shazeer 2020
5

---

Published as a conference paper at ICLR 2025
Table 2: Activation memory footprint of different operators. U is a unit to measure memory usage, where IU = Batch Size x Sequence Length x Hidden Size x 2 bytes (for BF16). For Llama-style model, Act Fune refers to SiLU & Multiply, and Linear refers to the summation of QKV/Atn/Up//aate/Da
jippl-- projection. RMSNorm is upcast to float32 in transformes (imple-- mentation) 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 D
Non-Linear
Attention
Reduction Ratio
RMSNorm
Act Fune
RoPE
FlashAttn
Linear
Total
Ideal
Achieved
BF16
4U
8U
2U
3U
5.66U
22.66U
1.00x
1.00x
Llama-style
TE
8U
2U
3U
3,33U
18.33U
1.23x
1,20x
COAT
10
4U
20
3U
3333
13.330
1.69x
1.65x
typically account for approximately 50% of the memory footprint in the Llama model series Tou- vron et al. (2023). In contrast, linear layers contribute less than 25%. Therefore, it is essential to optimize both linear and non-linear layers reduce activation memory footprint
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 (Yang et al. 2024] Fishman al. 2024). This necessitates us develop efficient and accurate quantization techniques.
5.2
MIXED GRANULARITY FP8 PRECISION FLOW
To address the inefficiency and inaccurate problem, propose to use 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.
We propose vary the quantization granularity across different layers to balance precision and efficiency in a mixed- grannlarrtllrrtty manner. For non-linear layers, VS- VS-Quant (Dai et al. 2021 or Per- Block Quant (Xi et al. 2024) methods are well-suited due their fine ame-grained and precise nature. For linear layers, we apply per-tensor quantization to maximize the performance of Tensor Cores
We observe that quantizing the input of layernorm across multiple token axes is detrimental to accuracy. As illustrated in Figure 4 4a). 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 x B as proposed in (Xi jet al., 2024). we propose to use per-group quantization with group size x G, where 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 where we display the full precision flow for a Llama-style decoder layer, both forward and backward pass.
D²
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 x G element and storing the results as intermediate values; (2) applying max reduction on the intermediate tensor 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 Gx smaller than the original tensor. As illustrated in Figure (b)) Group Scaling successfully reduces the max reduction overhead compared with just-in-ting scaling-
"RoPEs and Flash Attentioss are more as a distinct module and should be handled separately. FFinergrriind methods can also he applied since they are generally compatible with FP8 precision flow.

---

Published as a conference paper at ICLR 2025
4x4 quant for LayerNorm Einggne
1x16 quant for LayerNorm Mixed Granularity (Ours)
Tensor Size
Just-in-Time Time Scaling
Group Scaling (Ours)
600
11008-16384
1.37
0.10
450
11008-8192
0.89 ms
11x
0.08 ms
4096-16384
0.55
0.07
150
4096-8192 92
0.32 ms
0.06 ms
13
Layer
(b) Time to enleulate per-tensor max value, Our Group Sealing greatly reduce reduction overhead, it suffer from Delayed Scaling's instability issuc.
4 quantization of LayerNoorm' 's input lends large quantization orvor. The error becomns larger in deeper layers.
Figure (a) Quantization Error in forward pass. (b) Time comparison of various scaling methods.
In comparison, TransformerEnggnee 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, 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 simpler and more flexible alternative that can potentially offer improved numerical stability and precision while not being significantly slower than Delayed Scaling.
find that using the current tensor's 's statisties to compute
EXPERIMENTS
6.1 ACCURACY EXPERIMENTS
Setups We compare COAT with BF16 training and TransformerEnginne (NVIDIA 2024b) base- lines. For all experiments, 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 tnn-itnnng, and Vision Language model (VLM) training. For LLM pertaining, we report the perplexity Wikitext 103 (Merity et al. 2016). C4 (Raffel et al., 2020). and Pile (Gao al. 2020), and the accuracy on COPA (Gordon al. 2012), ARC (Clark et. 2018), SciQ (Welbl et al. 2017), and HellaSwag (Zellers et al.) 2019). For LLM fine-tuning, we conduct experiments
in math corpus, and evaluate on Mathmeticas (Davies et al.
2021).
SVAMP (Patel et al. 2021)
NumGLUE (Mishra et al.] 2022). and GSM8K (Cobbe et al. the score on VideoMME (Fu et al. 2024), POPE (Li et al. GQA (Hudson & Manning (019). VQAv2 (Goyal al.]
2021
For VLM training, we report
2023b),
VizWiz (Gurari et al. 2018
2017), TextVQA (Singh al. 2019
SEED (Li et al.) 2023a), and MMMU Validation Set (Yue et al. (024).. We use x 128 per-group quantization for optimizer states and 1 16 per-group quantization for non-linear layer activations.
6.1.1 LLM PRETRAINING
To evaluate our method when pretraining LLMs, We train OLMo- (Groeneveld al. 2024) and OLMo-7B on Dolma (Soldaini et al. 2024p. Following the official report, we use global batch size of 4M tokens (2048 macro batch size, with sequence length of 2048 tokens). We use PyTorch FSDP in our experiments.
For OLMo- 1B, conduct pretraining for 300B tokens, which corresponds to 75k training steps. We report the training curve in Figure 5 and report the perplexity and accuracy result in Table 3 The training curve and downstream task performance validating the effectiveness of COAT. For OLMo--B experiment, we pretrain for 250B tokens, which corresponds 40k training steps.we we report the training curve in Figure 6 and downstream task performance in Table 4 where COAT also aligns well with the baseline and is nearly lossless.
consistent with BF16 training baseline,
6.1.2 LLM FINE-TUNING
We further evaluate our method on LLM fine-tuning task. We focus on math corpus and tne-tune a Llama-2-7B model the МAmmoTH (Yue et al. 2023) dataset. We train for 3 epochs, and report the downstream tasks performance in S After fine-tuning, COAT still performs consistently with the baselines the downstream task performance, proving the accurateness of our method.

---

Published as a conference paper at ICLR 2025
Table 3: OLMo- pretraining performance on downstream tasks. We report the performance after training for 300B tokens.
4.5
Training Loss Curve of O1Mo-11
Br36
Train Loss
WikiTeext
C4
Pile
Avg ppl
BF16
2,551
15.234
15,538
10.563 10.672
10.083
COAT
2.568
15.384
15.695
10.176
1.5
CОРА
ARC(Easy)
SciQ
HellaSwag
Avg Acc
120
200
240
320
BF16
77.0 %
57.3 %
84.0 %
54.5 %
68.2 %
Figure 5: OLMo- 1B train- ing loss curve.
COAT
75.0 %
58.1 %
83.9 %
54.3 %
67.8 %
Table 4: OLMo-7B pretraining performance on downstream tasks. We report the performance after training for 250B tokens.
4.00
Trariing Loss Curve of OLNo 78
Train Loss
WikiTeet
C4
Pile
Avg ppl
3.00
BF16
2.366
12.053
12.874
8.596
11.174
COAT
2.379
12.166
12.988
8.684
11.279
СОРА
ARC(Easy)
SciQ
HellaSwag
Avg Acc
150
BF16
83.0%
65.7% 61.9%
87.5%
56.9%
73.2
Figure 6: OLMo-7B train-
COAT
81.0%
87.2%
60.6%
72.7 %
ing loss curve.
6.1.3
VLM TRAINING
We also evaluate COAT vision language models. We conduct experiments VILA (Lin al. 2024) and perform stage-3 SFT of VILA1.5-7B using the same SFT data mixture employed by VILA's original paper (Chen al. 2023 Xu
(2024) and set the global batch size 1024. We pad the sequence length to a multiple of 4 for efficiency consideration. The training loss curve is visualized in Figure We report the downstream task performance and their average in Table We find that COAT performs on par with BF16 training and is better than the TransformerEnginee base- line, which demonstrates the accurateness our method. We further visualize the VLM captioning experiment in Appendix B to prove the effectiveness of COAT on generation tasks.
6.2 MEMORY SAVING AND SPEEDUP
We test the memory saving and speedup result of COAT in two settings: The results 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.
6.2.1 MEMORY SAVING AND SPEEDUP FOR A SINGLE TRANSFORMER LAYER
Table highlights the speedup and memory reduction achieved for 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.57x speedup over BF16 and achieves a consistent 1.65 memory reduction compared to BF16, which is very close to the theo- retically 1.69x reported in Table The speedup ratio becomes larger with larger hidden sizes and longer sequence lengths.
6.2.2 SPEEDUP AND MEMORY SAVING FOR END-TO-END TRAINING
Table 8 presents a detailed comparison of end-to-end memory reduction and speedup results across different configurations for transformer models, specifically Lamm--2-B, Llama-2-133 13B. and Llama- 30B, with variations in the number of GPUs used (1, 2, and 8). It highlights COAT's effectiveness
8

---

Published as a conference paper at ICLR 2025
Table 5: Evaluation result fine-tuning Llama-2-7B on math corpus. Llama -amm--2--B refers the evaluation metric before fine-tuning, TE refers to TransformerEngine.
Mathmeticas
SVAMP
NumGLUE
GSM8K
Avg
1.1 ama 2 7B
6.0
14.6
34.5
29.9
21.3
BFI6
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
Training Loss Curve VILA-7B
Table 6: VILA1.5-7B Stage-3 SFT performance on down-
2.0
stream tasks. # means has seen the training data.
1.6
Stage
VidcoMME
POPE
VizWiz
GQA*
VQAv2*
BF16
42.96
86.90
61.42
64.55
81.47
TE COAT
43.19
87 64
$7.61
64,53
81.34
1.2-
44.56
87.43
61.36
64.44
81.20
TE
SEED
0.8
Stage 3
TextQQA
Image
Video
MMMU Val
Average
2000
BF16
65.60
73.40
45.65
38.56
62.80
TE
64.70
73.51
43.12
35.89
61.88
Figure 7: VILAA15-7B Stage-3 SFT
COAT
64.65
73.36
43.76
37.22
62.51
loss curve.
Table 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.65x.
Hidden 2048, Size
Sequence Length
2048
Sequence Length 4096
Forward
Backward
Total
Memory
Forward
Backward
Total
Ratio
Momory
BF16
3.36
8.47
11.83
1.00x
1457 MB
1.00x
6.88
17.24
24.12
1.00x
2914 MB
1.00x
TB
2.96
5.32
8.28
1.42x
1210 MB
1.20x
5.94
11:29
17.23
1.30x
2420 MB
1.20x
Hidden Sizr
4096, Barch Size
Sequence Length 2048
Sequence Length 4096
Forward
Backward
Total
Ratio
Memory
Ratio
Forward
Backward
Total
Ratio
Momory
Ratio
BF16
7.77
18.78
26.55
1.00x
2914 MB
1.0%
16.37
38.43
54.80
1.00x
5828 MB
1.00x
11.79
2420
12.66
37.24
1,20x
COAT
10.96
16.85
1.57x
1766 MB
1.65X
12:16
23.44
35.6
1.53x
3533 MB
1.65X
in reducing end-to-end memory footprint and the speedup compared to standard BF16 and Trans- formerEngine (TE) setups under varying conditions of batch size and context length.
COAT allows full-parameter training of Llama-2-7B on single GPU, where BF16 and TE both out of memory (OOM). Similarly, for the Llama-2-4 2-13B and Llama-30B models, our method enables 2-GPU training for Llama-2-133 1-1BB-2-11BBBB 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, method can achieve 2.25x speedup when training Llama-2-133 2-1B on 4-GPUs since we can effectively increase the batch size to
Overall, COAT significantly reduces end-to-end memory usage by up 1.55x and speeds up the end-to-end training by nearly 1.44x. This facilitates full full-parameter training on fewer GPUs, which particularly beneficial for larger language models.

---

Published as a conference paper at ICLR 2025
Table 8: 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. $ means CL=1024.
Llama-2-7B
Context Length 2048
Maximan Batch Size, Context Length 2048
Optmizer
Activations
Peak
Ratio
Max BS
Speed
Ratio
BF16
OOM
OOM
I GPU
TE COAT
OOM
OOM
13.1 GB
8.1 GB
79.3 GB
1
5906 token's
BF16
OOM
6130 token/s
1.00x
2 GPU
TE
OOM
1
6842 token's
1.11x
COAT
6.5GB
16.9
52.8 GB
11351 token's
1.85x
BF16
13.1
25.8
55.1 GB
1.00x
2
7730 token's
1.0x
4 GPUss=2
TE
13.1 GB
21.9 GB
51.1 GB
1.08
2
9577 token's
1.24x
COAT
3.2 GB
16.9
35.6 GB
1.54x
11257 token's
1.45x
BF16
6.5 GB
25.8
41.2 GB
1.00x
8238 token's
1.00x
8 GPUpse
TB
6.5 GB
21.9 GB
37.2 GB
1.11x
4
11704 token's
1.42%
COAT
GB
16.9
27.0 GB
1.52x
8
11241 token's
1.36x
Llama-2-13B
Context Length 2048
Maximun Batch Size, Context Length
2048
Optimizer
Activations
Peak
Ratio
Max BS
Speed
Ratio
BF16
OOM
OOM
2 GPU B5ml
OOM
OOM
COAT
12.6 GB
10.1 GB
73.2 GB
2137 token/s
BF16
25.1 GB
20.1 GB
76.1 GB
1.00x
2345 token's
1.00x
GPURsa
25.1
17.2
73.0
1.04x
2851 token's
1.21x
COAT
6.3 GB
13.2 GB
49.1 GB
1.55X
2
$295
token's
2.25x
BF16
12.6
20.1 GB
49.4 GB
1.00x
3907 token's
1.0x
8 GPUss=
TR
12.6 GB
17.2 GB
46.5 GB
1.06x
2
5604 token/s
1.43x
COAT
3.1 GB
13.2 GB
32.5 GB
1.52x
4
5650 token's
1.44x
Llama-30B
Context Lengih = 2048
Maximum Batch Size, Context Length = 2048
Optimizer
Activations
Ratio
Speed
Ratio
BF16 TE COAT
0OM
OOM
8 GPUss=1
OOM
OOM
7.8 GB
24.2
70.5 GB
1363 token's
Table 9: Dynamic Range Expansion is compatible with DE8 (8-bit dynamic quantization).
Sccond Order
First Order
E4M3
E4M3
Expand
E5M2
E5M2
Expand
DE8
DE8 + Expand
E4M3 Expand
15.13
12.31
21.96
12.43
14.01
18.84
DE8
12.11
8.27
20.02
8.43
10.54
16.25
DF8 Expand
11.57
7.47
19.69
7.65
9.91
15.81
6.3
ABLATION STUDIES
6.3.1
DYNAMIC RANGE EXPANSION'S 'S COMPATIBILITY WITH OTHER DATA FORMATS
Dynamic Exponent Quantization (DE) was proposed in DDettners et al. 2021 to quantize the optimizer states since its representation range has range of orders of magnitude and is very suitable for optimizer states quantization. Therefore, they proposed quantize both first-order and second-order momentum with 8-bit DE.
We report the quantization error of / in Table and find that applying our dynamic range expan- sion method to 8-bit DE can further reduce the quantization error by 1.41 x. Specifically, the lowest quantization error is achieved when first-order momentum is quantized with 8-bit DE + Dynamic Range Expansion, and second-order momentum is quantized with E4M3/E5M2 + Dynamic Range Expansion. This proves the effectiveness of our method.
7 CONCLUSION
In this work, we present COAT, a memory FP8 training framework for foundation models by quantizing both optimizer states and activations into FP8 format. We observe that the FP8 for-
10

---

Published as a conference paper at ICLR 2025
mat's representation range is not fully utilized when quantizing optimizer states, prompting to propose Dynamic Range Expansion to align their dynamic range. We then identify the importance of quantizing non-linear layers, and propose activations accurately without introducing too much overhead.
FP8 precision flow quantize the
Extensive experiments LLM and VLM training and fine-tuning demonstrate that COAT achieve nearly lossless performance. In end-to-end training, COAT achieves 1.54 memory reduc- tion and 1.43 speedup compared BF16, and is comparable or even faster to TransformerEnginee training speed. Our method also enables full-parameter training of billion-scale models on fewer GPUs and is able double the batch size realistic settings. These results highlight COAT's capa- bility to enable wemory-efficieee large-scale model training without sacrificing accuracy, providing a highly effective solution for memory-constrane combining our proposed approach with other low-precision gradient compression methods to reduce communication overhead.
environments. Future work could further explore

---

Published as a conference paper at ICLR 2025
REFERENCES
Bo Adler, Niket Agarwal, Ashwath Aithal, Dong H Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, al. Nemotron-4 340b technical report. arXiv preprint arXiv:2406. /704, 2024.
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimiza- tion, Advances in Neural Information Processing Systems, 32, 2019.
Jimmy Lei Ba. Layer normalization. arXiv preprint arXiv:1607.064450 2016.
Han Cai, Chuang Gan, Ligeng Zhu, and Song Han. Tinytl: Reduce activations, not trainable param- eters for efficient on-device learning. arXiv preprint arXiv:2007. /622,, 2020.
Jianfei Chen, Yu Gai, Zhewei Yao, Michael W Mahoney, and Joseph E Gonzalez. A statistical framework for low-bitwidth training of deep neural networks. Advances in neural information processing systems, 33:883-894, 2020.
Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez. Acton: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning, pp. 1803-1813, PMLR, 2021.
Lin Chen, Jisong Li, Xiaoyi Dong, Pan Zhang, Conghui He, Jiaqi Wang, Feng Zhao, and Dahua Lin. Sharegpt4v: Improving large multi-modal models with better captions. arXiv preprint arXiv:2311.12795 2023.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsich, Yifeng Lu, al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXhv:1803.05457, 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 4166, 2021.
Steve Dai, Rangha Venkatesan, Mark Ren, Brian Zimmer, William Dally, and Brucek Khailany. Vs-quant: Per-vector scaled quantization for accurate low-precision neural network inference. Proceedings of Machine Learning and Systems, 3:873-884, 2021.
Alex Davies, Petar Veličković, Lars Buesing, Sam Blackwell, Daniel Zheng, Nenad Tomašev, Richard Tanburn, Peter Battaglia, Charles Blundell, András Juhász, et al. Advancing mathematics by guiding human intuition with ai. Nature, 600(7887):70-74, 2021.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXiv:21 2021.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788. 2024.
Maxim Fishman, Brian Chmiel, Ron Banner, and Daniel Soudry. Scaling fp8 training to trillion- token 1lms. arXiv preprint wxXi:2449.115117 2024.
Chaoyou Fu, Yuhan Dai, Yondong Luo, Lei Li, Shuhuai Ren, Renrui Zhang, Zihan Wang, Chenyu Zhou, Yunhang Shen, Mengdan Zhang, et al. Video-mme: The first-ever comprehensive evalua- tion benchmark of multi-moda 1ims in video analysis. arXiv preprint arXiv:2405.21077 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv: 100.00077, 2020.
12

---

Published as a conference paper at ICLR 2025
Andrew Gordon, Zornitsa Kozareva, and Melissa Roemmele. SemEval-2012 task 7: Choice of plau- sible alternatives: An evaluation of commonsense causal reasoning. In Eneko Agirre, Johan Bos, Mona Diab, Suresh Manandhar, Yuval Marton, and Deniz Yuret (eds.), *SEM 2012: The First Joint Conference on Lexical and Computational Semantics Volume I: Proceedings of the main conference and the shared task, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation (SemEval 2012), pp. 394-398, Montréal, Canada, 7-8 June 2012. Asso- anthology org/ Argg/122- S12 1052
ciation for Computational Linguistics. URL httph: pittts://acce http:s/s///aaaahloooyyyooo
Yash Goyal, Tejas Khot, Douglas Summers- Dhruv Batra, and Devi Parikh. Making the v in vqa
matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 6904-6913, 2017. 2017.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang. al. Olmo: Accelerating the science of language models. arXiv preprint arXiv:2402.00:: 00038, 2024.
Danna Gurari, Qing Li. Abigale Stangl, Anhong Guo, Chi Lin, Kristen Grauman, Jiebo Luo, and Jeffrey P Bigham. Vizwiz grand challenge: Answering visual questions from blind people. In Proceedings the IEEE conference on computer vision and pattern recognition, pp. 3608-3617, 2018.
Dan Hendrycks and Kevin Gimpel. 16Xiv:1606.0e0 2016.
Gaussian
error
linear units (gelus). arXiv preprint
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Drew A Hudson and Christopher Manning. Gqa: A new dataset for real-world visual reasoning and compositional question answering. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 6700-6709, 2019.
Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, et al. A study of 6 for deep learning training. arXiv preprint arXiv: 1e.55.22350 12322, 2019.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint 2014.
Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 36, 2024.
Bohao Li, Rui Wang, Guangzhi Wang, Yuying Ge, Yixiao Ge, and Ying Shan. Seed-bench: Bench- marking multimodal Iims with generative comprehension. arXiv preprint arXiv:2307.16125, 2023a.
Yifan Li, Yifan Du, Kun Zhou, Jinpeng Wang, Wayne Xin Zhao, and Ji-Rong Wen. Evaluating object hallucination in large vision- languaean models. arXiv preprint 2xXi:23005.10355 2023b.
Ji Lin, Ligeng Zhu, Wei-Ming Chen, Wei-Chen Wang, Chuang Gan, and Song Han. On-device training under 256kb memory. Advances in Neural Information Processing Systems, 35:22941- 22954, 2022.
Ji Lin, Hongxu Yin, Wei Ping, Pavlo Molchanov. Mohammad Shoeybi, and Song Han. Vila: On pre- training for visual language models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
2024.
Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, et al. Gact: Activation compressed training for generic network architectures. In International Conference on Machine Learning, pp. 14139-14152. PMLR, 2022a.
13

---

Published as a conference paper at ICLR 2025
Zechun Liu, Kwang-Ting Cheng, Dong Huang, Eric P Xing, and Zhiqiang Shen. Nonuniform-to uniform quantization: Towards accurate quantization via generalized ttraight-thooug estimation. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp- 4942-4952 2022b.
I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv:1711.051041 2017.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprini arXiv: 160.07843, 2016.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, al. Mixed precision training. arXiv preprint 171..377000034400 2017.
Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Comea,, Pradeep Dubey, Richard Grisen- thwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, et al. Fp8 formats for deep learning. arXiv preprint arXiv:2209.05433 2022.
Swaroop Mishra, Arindam Mitra, Neeraj Varshney, Bhavdeep Sachdeva, Peter Clark, Chitta Baral, and Ashwin Kalyan. Numglue: A suite of fundamental yet challenging mathematical reasoning tasks. arXiv preprint arXiv:2204.0566 2022.
Georgii Sergeevich Novikov, Daniel Bershatsky, Julia Gusak, Alex Shonenkov, Denis Valerievich Dimitrov, and Ivan Oseledets. Few-bit backward: Quantized gradients of activation functions for memory footprint reduction. In International Conference on Machine Learning, pp. 26363- 26381. PMLR, 2023.
NVIDIA. Nvidia h100 tensor core gpu, 2024a. URL https: http:/wwwwwwiiia..////
data data-Genter/h100/ Accessed: 2024-09-19.
NVIDIA. Transformerenginn: An efficient library for training transformer models, 2024b. URL
https://github: http:////.huubb..cm/////t IA www.:a//w..aaaaooeeeee.n...
Accessed: 2024-09-19.
Open Compute Project
Ocp 8-bit floating point specification (ofp8), revision 1.0,
December 2023.
URL https::/ http:////ww.eeeeeeooeeeeeeo//ooo
sep-8-bit-floatiin
Accessed: 2024-10-09.
Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are nlp models really able solve simple math word problems? arXiv preprint arXiv:2103.07191 2021.
Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze Liu, Yifan Xiong, Ziyue Yang. Bolin Ni, Jingcheng Hu, et al. Fp8-Im: Training fp8 large language models. arXiv preprint arXiv:2310.188 18313, 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21((40):1-67, 2020.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti- mizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery Data Mining, pp. 3505-3506, 2020.
Jay Shah, Ganesh Bikshandi. Ying Zhang, Vijay Thakkar, Pradcep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608, 2024.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202. 2020.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604 PMLR, 2018.
14

---

Published as a conference paper at ICLR 2025
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv:1909.0800:. 2019.
Amanpreet Singh, Vivek Natarajan, Meet Shah, Yu Jiang, Xinlei Chen, Dhruv Batra, Devi Parikh, and Mareus Rohrbach. Towards vqa models that can read. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8317-8326, 2019.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deep- speed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint ar arXiv:2011.11990 2022.
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, et al. Dolma: An open corpus of three trillion tokens for language model pretraining research. arXiv preprint arXiv:2402.00159, 2024.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.088 08295, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv: 13971,
Jue Wang, Binhang Yuan, Luka Rimanic, Yongjun He, Tri Dao, Beidi Chen, Christopher Re, and Ce Zhang. Fine-tuning language models slow networks using activation compression with guarantees. arXiv preprint arXiv:2206.01299 2022.
Naigang Wang, Jungwook Choi, Daniel Brand, Chia- Yu Chen, and Kailash Gopalakrishnan. Train- ing deep neural networks with 8-bit floating point numbers. Advances in neural information processing systems, 31, 2018.
Shuai Wang and Yi Kang. Gradient distribution-awarr int8 training for neural networks. Neurocom- puting, 541:126266 2023.
Johannes Welbi, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. arXiv preprint 1707.06990, 2017.
Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, and Ludwig Schmidt. Stable and low- low-precision training for large-scale vision-language models. Advances in Neural Information Processing Systems, 36: 16:10271-10299 2023.
Haocheng Xi, Changhao Li, Jianfei Chen, and Jun Zhu. Training transformers with 4-bit integers. Advances in Neural Information Processing Systems, 36:49 19:49146-4918,, 2023.
Haocheng Xi, Yuxiang Chen, Kang Zhao, Kaijun Zheng, Jianfei Chen, and Jun Zhu. Jetfire: Eff- cient and accurate transformer pretraining with int8 data flow and per-block quantization. arXiv preprint arXiv:2403. 200.:22203.12222.... 2024.
Haocheng Xi, Shuo Yang, Yilong Zhao, Chenfeng Xu, Muyang Li, Xiuyu Li, Yujun Lin, Han Cai, Jintao Zhang, Dacheng Li, et al. Sparse videogen: Accelerating video diffusion transformers with spatial-remporal sparsity. arXiv preprint 2rXiv:2502.01772 2025.
Zhiyang Xu, Chao Feng. Rulin Shao, Trevor Ashby, Ying Shen, Di Jin, Yu Cheng, Qifan Wang, and Lifu Huang. Vision-flan: Scaling human-labeled tasks in visual instruction tuning. arXiv preprint 24X2 2024.
Jaewoo Yang, Hayun Kim, and Younghoon Kim. Mitigating quantization errors due to activation spikes in glu-based lms.. arXiv preprint arXiv:2405. 2nfe:2405.44225 /4428, 2024.
15

---

Published as a conference paper at ICLR 2025
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309.05653, 2023.
Xiang Yue, Yuansheng Ni, Kai Zhang, Tianyu Zheng, Ruoqi Liu, Ge Zhang, Samuel Stevens, Dongfu Jiang, Weiming Ren, Yuxuan Sun, et al. Mmmu: A massive multi-discipline multi- modal understanding and reasoming benchmark for expert agi. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9556-9567, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? arXiv preprint
2019.
Biao Zhang and Rico Sennrich, Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen. Sageattention2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization, 2024. URL
https://arxit org/ abs/ 2411 1..11.41 10958
Jintao Zhang, Jia Wei, Pengle Zhang, Jun Zhu, and Jianfei Chen. Sageattention: Accurate 8-bit attention for plug-and and-play inference acceleration. In International Conference Learning Rep- resentations (ICLR), 2025a.
Jintao Zhang, Chendong Xiang, Haofeng Huang, Haocheng Xi, Jia Wei, Jun Zhu, and Jianfei Chen. Spargeattn: Accurate sparse attention accelerating any model inference, 2025b.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficien Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507 2024.
Shu-Chang Zhou, Yu-Zhi Wang, He Wen, Qin- Yao He, and Yu-Heng Zou. Balanced quantization: An effective and efficient approach quantized neural networks. Journal of Computer Science and Technology. 32:667-682, 2017.
Feng Zhu, Ruihao Gong, Fengwei Yu, Xianglong Liu, Yanfei Wang, Zhelong Li, Xiuqi Yang, and Junjie Yan. Towards unified int8 training for convolutional neural network. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1969-1979, 2020.
16

---

Published as a conference paper at ICLR 2025
A DETAILS ABOUT OPTIMIZER STATES QUANTIZATION
When performing opt imi zer 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:
M D
(Dequantize FP32) (Dequantize to FP32)
4- DQ(() Sen-1)
m4 81 m4-1 + (1 (1) 9t
$2 1t1 (1 -2) gi
m
m, 1- B
0. Ut
32
m, + xaws
W111 wi
Vortc
m4, Sm; ((m)
(Quantize to FP8)
ui. Sy, Q(v))
(Quantize to FP8)
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. Fach 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.
B QUALITATIVE EXAMPLE VISION LANGUAGE MODEL CAPTIONING
C VISUALIZATION OF EXPAND FUCTION
We further visualize Figure g by flattening it. This helps to understand the effectiveness of our method since floating point numbers can be more easily understood binary manner.
D DETAILED EXPLANATION FOR TABLE 2
We mainly explain why the reduction ratio of lincar layers is not exactly 50% here. 5.66U comes from: QKV Projection IU, Attention Projection 1U, Up&Gate Projection 1U, Down Projection 2.66U. Among them, the attention projection's ' input is exactly the output of FlashAttention. How- ever, FlashAttention will save its input (QKV) and its output itself, FlashAttention and Attention Projection actually share the same tensor when saving it. In Python, this tensor will only be saved for time. So we should not further store FPS version of the Attention Projection's 's input (Zhang) et al. 2025a 2024), since this will even increase the memory usage by 0.5U if we do not change the code of FlashAttention Although we need quantize the input to FP8 again in backward pass, 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& Projection 0.5U, Down Projection 1.33U. They sum up 3.33U.
E DETAILED EXPLANATION FOR FIGURE T
We mainly discuss FP8-LM this section. FP8-LM reduces the gradient communication precision from FP32 to FP8 and therefore greatly reduces the communication overhead. It also reduces the
17

---

Published as a conference paper at ICLR 2025
Instruction: Briefly describe this figure.
The image depicts a fantastical scene set in a mountainous landscape. Dominsting the center of image large tree, twisted and grarled form standing on rocky outcnopping. The tree is surrounded by a sereme lake, its calm reflecting the vibrant colors of sky above. The sky spectacle in itself, painted with hues of orange and red as the sun sets behind mountains, casting warm glow over the back Buckground add sense of depth and grandeur to the image, their peaks reaching up
BE16
scene.
sky. The overall composition tranquil yet powerful natural setting, possibly inspired by mythology or fantasy.
image suggests
The image presents captivating seene of solitary tree standing rocky outeropping, The trec, with its twisted and gnarled form, vibrant green. contrasting with surrounding landscape. The rocky outeropping on which it stands is rugged and uneven. udding the overall wildness of the scene. The buckgroumd breathtaking view of mountainous landscape under hues of orange and red, possibly indicating sunrise sunset. The sun is partially obscured by the clouds, casting a warm glow that illuminates the scene. The image rendered digital style, which adds skillfilll used color and lighting to create sense of tranquility and solitude, with the tree standing as solitary figure against the vastness the lundscape. overall composition the image bulanced hurmenious, each element complementing the cohesive whole.
COAL:
abluze with
realism and depth
The artist has
create
Question: Why People love this?
People love this because visually appealing and appetizing image of delicious burger. The burger presented paper bug, which adds rustic touch presentation. The combination of juicy burger with a red pepper and fiery back crentes an enticing The ceptures the
BF16:
and mouth-watcring burger, and the excitement of freshly cooked meal. The vibrant colors and the dramatic image even more cuptivating, appreciated by many people.
of classic American favorite,
backdrop
why likely
COAT: People love this image of burger bun with various toppings, such onions, tomatoes, and lettuce, because represents delicious and visually appealing meal. The burger is placed on a piece of paper, which adds an interesting and casual touch to the presentation. The background. with anticipation for the meal. The combination of the mouth-watering burger and the striking background can make the image appealing and appetizing to many people, especially those enjoy burgers similar dishes.
flame,
dramatic and exciting atmosphere, which evoke
and
Figure 8: Comparison of BF16 and COAT VLM captioning. COAT can accurately summarize the figure and identify the key points in the figure.
Fist vode:
Sovesluskr
Figure 9: Axis is of base 2.
master weight's precision from FP32 FP16/BF16, and reduces the optimizer states precision from FP32 to BF16/FPS. For optimizer states and master weights. the memory footprint reduction result 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 reduce the memory footprint the weight gradient.
F PROOF THE EXPAND FUNCTION IS OPTIMAL
We prove that the choice of f(x) = * is optimal given certain basic assumptions for the expansion function f(x).
1. f(x) is an odd function, where f(x)+((-a)) = 0, since the representation range of FP8 is symmetric. Therefore, in the following discussion, we only consider the case when 2 0.
2. f(x) is continuously differentiable.
3. f(x) is monotonically increasing. This is because we want to keep the relative order before and after quantization expansion.
18

---

Published as a conference paper at ICLR 2025
4. f(0) 0 and f(1) 1. f(0) 0 since f(x) is odd. f(1) = since we can always consider g(x) f(1)
f(x) without changing its expansion ability.
[ firz) This is used to make sure the function is 5. For any > y 0, r 0, Tly
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.
We prove that based on these assumptions, the only function that meets these assumptions is f(z) (ggn(>)aa)))
Proof: We only consider the case when O. We denote f(2) = > 1. From assumption 5, we have for any p > r s q, if pq = rs, then f()) f(r)fss, since
f(p) f(p (s/p)) f(s)
[(r) ((r-(s/p) f(q)
And for any p rs, Therefore, for any = 2", E.
f(p) = f(D)f(1) = f(r)f(s).
f(x) = =
= f(2)' = m
is also that for any == 2- E,
f(1) 1(1)
f(x)
so that f(x) m
[(2]
[(2()
m
N
Therefore, for any N e, 2
Σ c 22-2'. 10.1).
i=-N
N
N
N
2z m
f(x) = f( II
20.2
II
f(c: - 2¹) = II
m 2
m---N
N
N
N
For Væ > 0, can be written in the form of x = 2', E R. Consider the binary representation of
t, denoted by t Σ C: 2', where c; E {0,, 1}. For any integer N > 0, there exists two series
1=-00
( and (b,) NN that satisfies a-N = 0, N 1, and a; b, for N N+IsisN,
La < lb:
N
where to Σ a 2 and to
Σ b, 2 Since f(x) continuous, the values of f(z at Fa = 2 and In = 2 can be used to approximate f(x) = f(2"), and by taking the limit as
i=-N
N- + co, the equality holds:
f(z) lim f(xa) lim f(xs).
No0
Therefore for V / E RR, ((2") m² for some m > 0. Therefore V x > 0, f(x) (og2()) Using the exponential and logarithmic identity (x) = (m)) this simplifies to:
f(x) x (m))
where (m))(m) is constant determined by f(2) m. If we denote logy(m) = k, then f(x) = t
G 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 effi- ciently.
19

---

Published as a conference paper at ICLR 2025
G.1 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:
1. BeforeAttention module (inherit from torch.autogra.. aorchauttorad.Fuunciio Function) calculates the LayerNorm and QKV Projection before RoPE and correctly handles how quantized inputs interact with the residual connection. If Attention is quantized (Zhang et al. [2025a) 2024; Shah et al.) 2024 Zhang al. 2025b; Xi et al. 2025), then the output of this module should also be quantized to accommodate the input of Attention.
2. AfterAttention (inherit from torch. Function)) calculate the attention projection layer and ensures quantized outputs are correctly scaled and summed up with residual con- nection.
3. MLPResidual (inherit from lorch.autograd.o Function) handles quantization within the feed- forward network (FFN) efficiently. also incorporates residual connections in this module to handle it efficiently.
These modules ensure that the computation flow after quantization correct. We fuse the group scaling method into the add operator in residual connection, therefore need to manually control how the gradient flows in the network.
G.2 TRITON- BASED FP8 KERNELS FOR LINEAR LAYERS
The linear layers utilize Triton to implement custom FP8 matrix multiplication kernels. This imple- mentation introduces the following optimizations:
Block Size Tuning: The block sizes for dimensions M, NN, and K are autotuned, with op- tions chosen from {128, 256). This dynamic tuning optimizes performance for the specific hardware environment.
2. Group Scaling Technique: The output of linear layer should be BF16 or FP8 per 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 get the scaling factor.
G.3 TRITON-BASED FP8 KERNELS FOR NON-LINEAR LAYERS
1. Transposed FP8 Output: Besides generating the quantized FP8 output, these layers also produce transposed FP8 output. This is critical for backward computations, where column-major inputs are required to calculate weight gradients.
2. Block Size Tuning: Block sizes for dimensions M and N are autotuned, with options chosen from (32, 64).
G.4 OPTIMIZER STATES
We implement fused kernel for optimizer states update CUDA. Since optimizer states use 1x 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 get the maximum absolute value and minimum absolute value within quantization group to perform dynamic range expansion. We
Direct computation of X is usually about 10-12 and k > 3. To mitigate this issue, we introduce an additional FP32 value Cx, defined Cx VXmin X Xmaxx ensure numerical stability. By normalizing X with Cx, the values X and Xms Cx Cx computed, which are reciprocals of each other and closer to 1 compared to the original X - This normalization enhances numerical robustness and reduces the risk of instability.
result in underflow and numerical instability since its magnitude
20

---

Published as a conference paper at ICLR 2025
G.5 OTHERS
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.
2. Gradient Accumulation Optimization: During the first gradient accumulation step, the scal- ing factors of weight matrices are stored. This avoids the need 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.
H 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 backward pass, comparing with BF16. Based on Table below, we find that group size 16 achieves good balance for memory efficiency and quantization error.
Only Quantize Linear
Also Quantize Activation
MSE
1 x 128 Group
Per Tensor
8
16
64
Memory Overhead Forward
50%
25%
12.5%
6.5%
3.2%
372
438
464
480
492
512
520
Backward
285
286
319
334
363
361
371
Table 10: Quantization overhead, forward, and backward MSE for different configurations and group sizes.
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 FPS. On the other hand, the accuracy gain is only marginal. Therefore we do apply per-tensor quantization for linear layers.
EFFICIENCY BREAKDOWN FOR OPTIMIZER AND ACTIVATION
We break down the improvement of each component and present the result this table. Quantizing the optimizer states to FP8 is helpful in reducing the memory footprint, and potentially increase the batch size to fully utilize the GPU. Quantizing the activation to FPS can accelerate the training, and this number can be further improved by increasing the maximum batch size.
21

---

Published as a conference paper at ICLR 2025
Table 11: Llama-2-7B Training Memory and Throughput Analysis
Configuration
Optinizer
Activation
Peak
Max BS
Throughput
Speedup
Llama-2-B: on GPU BF16
26.2 GB
25.8 GB
OOM
6130 tokens/s
1.00x
FP8 Optimizer
6.4 GB
25.8 GB
61.9 GB
7258 tokens/s
1.18x
FP8 Optimizer + FP8 Activation
6,4 GB
16.9 GB
35.6 GB
4
11351 tokens/s
1.85x
Llama-2-7B GPU
BF16
13.1 GB
25.8 GB
55.1 GB
7730 tokens/s
1.00x
FP8 Optimizer
3.2 GB
16.9 GB
44.7 GB
8384 tokens/s
1 1.08x
FP8 Optimizer + FP8
3,2 GB
16.9 GB
35.6 GB
4
11257 tokens/s
1.45x
LJama-2-7B on 8 GPU
BF16
6.5 GB
25.8 GB
41.2 GB
8238 tokens/s
1.00x
FP8 Optimizer
1.6 GB
16.9 GB
36.1 GB
8444 tokens/s
1.02x
FP8 Optimizer + FP8 Activation
1.6 GB
16.9 GB
27.0 GB
11241 tokens/s
1.36x
J OVERHEAD INTRODUCED BY GROUP SCALING
We show that the max reduction on each IxG "can be seamlessly fused with the previous operation, adding minimal overhead". Take RMSNorm as an example, the speed of the RMSNorm kemne only increases by less than 5% when quantization is fused with it. Therefore, this adds only minimal overhead to the kernel.
Table 12: Performance Comparison: With and Without Quantization
Matrix Size
With Quantize
Without Quantize
(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
K UNDERSTAND HOW EACH COMPONENT INFLUENCES THE MODEL PERFORMANCE
We separately quantize the optimizer and activations to better understand how each component influ- ences 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.
Train PPL
CОРА
ARC (Easy)
SciQ
HellaSwag
Average
BF16
2.995
60.0
45.6
67.3
33.7
51.6
FP8 O A
3.008
61.0
44.2
67.6
33.7
51.5
FPS O
2.998 (+0.003)
60.0 (+0.0)
44.5
66.0
34.1
51.1
FPS A
2.999 (+0.004)
63.0 (+3.0)
44.0
68.4
34.1
52.3
Table 13: Performance comparison across different metrics and configurations.
L MORE RELATED WORKS
Due page limit, some related works presented in Appendix.
22

---

Published as a conference paper at ICLR 2025
Table 14: Compared with DeepSeck-Ve
Non-Linear
Atlention
Reduction Ratio
RMSNorm
Act Fune
RoPE
FlashAttn
Linear
Total
Ideal
Achieved
BF16
4U
SU
2U
3U
5.66U
22.66U
1.0%x
1.00x
Llama-ttyl
TE
20
8U
2U
3.33U
18.33U
1.23x
1.20x
DeepSeek-V3 COAT
20
4U
2U
3U
3.08U
14.08U
161x
10
4U
2U
3U
3.330
13.33U
1.69x
1.65x
Activation Quantization Recent research has focused on reducing memory footprint Cai et al. (2020) during neural network training through activation quantization. ActNN (Chen et al. 2021) introduced 2-bit activation compressed training framework using random quantization, achieving 12x reduction in activation memory footprintl. GACT (Liu et al. 2022a) extended this concept to support various machine learning tasks and architectures, providing up to 8.1x reduction in ac- tivation memory for CNNs, transformers, and GNNs. However, these methods mostly focus on convolutional networks and not applicable LLMs. AQ-SGD (Wang al., 2022 proposed compressing activation changes rather than direct values for pipeline-paralle training to reduce the communication overhead in the pipeline-paralle parallel setting. Few-Bit Backward(Noviko) et al. 2023)
quantizes the non-linear activation functions using optimal connsatt approximations to reduce the memory footprint while maintaining convergence. Jetfire (Xi al. (024) 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.
Bit-
Balance Quantization Zhou al. (2077) recursively partitions the parameters by percentiles to reduce quantization error. DDSQ Wang & Kang (2023) dynamically changes the quantization parameters according different gradient distributions. N2UQ Liu et al. (2022)) improves the quantization precision via learning input thresholds and nonuniform map- ping. These approaches are aware of the dynamic range problem in quantization and propose more adaptive, distribution-awant techniques that can preserve the nuanced information contained in full precision. They share fundamental goal of expanding the effective dynamic range of low-precision neural network representations. These methods adopt learnable non uniform quantization lookup table, while our work relies on the natural FP8 data format.
23


Reviews:
[{"review_id": "XbIcovSVLc", "reviewer": "Reviewer_4rZz", "summary": "This paper introduces a novel fp8 training framework named COAT that improves throughput and memory footprint of (end-to-end) large-scale training respectively by 1.54x and 1.43x. In detail, COAT is built upon two key techniques of (1) dynamic range expansion that allows for abetter utilization of fp8 format’s representation range, and (2) mixed-granularity activation quantization that reduces the memory footprint for activations. The authors demonstrate that COAT achieves superior memory/compute efficiency while preserving the bf16 performance on three tasks of LLM pretraining, fine-tuning, and VLM training.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Given that H100 is a fairly new hardware, fp8 training has not been explored extensively in previous literature. At the same time, the potential memory/compute efficiency gain from fp8 training can be significant, so I believe this paper addresses an important problem and makes a progress in this topic.\n2. Dynamic range expansion is well-motivated by the analysis from Figure 2.\n3. I believe ~1.5x improvements in memory/compute efficiency is significant.", "weaknesses": "1. Even if the performance gap looks not too significant, COAT still underperforms bf16 training on most tasks. At the same time, the ablation study only focuses on memory/compute efficiency, but not the final model performance. Providing insights on how each component influences the model performance and insights on the source of the performance degradation can be useful.\n2. It is unclear to me how/whether the dynamic range is synchronized across GPUs in the distributed training setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5CB2V5Acpq", "reviewer": "Reviewer_P9H9", "summary": "This work proposes a quantization improvement for transformers that mainly targets the optimizer’s states and activations. The authors propose a FP8 per-group quantization of the optimizer’s states along with a “dynamic range expansion” that utilizes the quantized space more effectively. The FP8 per-group quantization is also applied on the activations with \"mixed-granularity\", i.e., different layers quantize their inputs differently (e.g., layernorm vs. linear). These improvements can save up to 1.5x memory over TransformerEngine and is slightly faster as well without sacrificing performance vs. FP16 quantization.", "questions": "None.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The experiments demonstrate significant space improvements in practice and slight speed up.\n- Figures illustrates methodology and motivation clearly.\n- The authors motivate the quantization method by looking at the low-bit values as a range of representable values, and show that optimizer states are sparse in this representation and do not span the range nicely. A transformation $f(x)=x^k$ can be applied before quantization and inverted after dequantization easily, allowing for the quanitized values to span the range of representable values more uniformly.\n- Table 2 shows the actual improvements in memory clearly.\n- The mixed-granularity approach is simple yet effective in practice.", "weaknesses": "The scientific contribution of this work is limited. Despite the practical improvements, it might be of little interest to the community of “learning representations”, but perhaps of great interest to LLM practitioners, which is certainly a significant part of the AI community. Thus, I will not wholly judge the paper by its scientific merit and shall take a pragmatic approach in my rating here, but the point regarding the contribution still stands.\n\nThe authors could have, for example, expanded on Sec 4.1. They could have studied or tried other inveritble transformations that spread out the values nicely across the representable values. Would direct training in this transformed low-bit regime work? Perhaps the author can study whether polynomial activations can have the same benefits, for example. This would make the paper more interesting.\n\nAnother thing is that the idea behind the dynamic range is quite straightfoward to conceive, so I believe a more careful literature review can reveal very similar techniques. For example, a quick search reveals [1], which studies a similar problem.\n\nFinally, the engineering improvements to transformer quantization are indeed non-trivial but can be found by extensive experimentations and ablation studies, which might not be possible for those who cannot afford to run such extensive experiments. Thus, I believe more effort should go into justifying the engineering choices found by the authors, which would be greatly appreciated and would motivate similar practices in the future. For example, a more fine-grained ablation studiy or theoretical analysis of the mixed-grained quantization can be great. This should also inform future readers on whether this particular quantization method is the best one for some particular setting.\n\n[1] Balanced Quantization: An Effective and Efficient Approach to Quantized Neural Networks. Zhou et al. 2017.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bSGmIfkEdo", "reviewer": "Reviewer_gVQo", "summary": "This paper introduces COAT, a scheme for efficient low-precision training of large models. It is based on compressing optimizer states and model activation through two mechanisms: _(i) Dynamical Range Expansion_ for fully exploiting the representation range of the FP8 format for the optimizer states) and _(ii) Mixed-Granularity Activation Quantization_, which uses a combination of per-tensor and per-group quantization strategies to compress model's activations.\n\nExtensive experiments are conducted on LLMs (OLMo-1/7B, LLaMa 2/7B) and VLMs (VILA1.5-7B), showing solid and consistent improvements in (i) memory efficiency and (ii) training speed, at nearly lossless performance.", "questions": "- Can you discuss better the approaches related to Activation Quantization? From lines 128-132 is not clear what are the improvements and the novelty of the proposed approach with respect to previous works.\n- Can you provide a deeper explanation of what are the main differences between the proposed Group Scaling and the TransformerEngine's Delayed Scaling? In the text it is mentioned that it causes instabilities, but it is not clear what they are and why they arise.\n\nSmall issues:\n- There is a missing ref. in line 372: \"performance in [] 4\".\n- Table 6 is a little bit confusing in reporting the memory saving and the speedup, e.g. it is not immediately clear that the first \"Ratio\" column refers to runtime and the second to memory usage. I suggest to make it clearer, for example by adding a vertical line after the first \"Ratio\" column.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "In my opinion, this paper has the following main strenghts:\n- **Clear presentation:** the paper is very clearly written, both the introduction and the related work sections introduce the necessary notions the proposed approach is built on and improves upon. All the figures/tables are of good quality and are well displayed in a printed version of the manuscript.\n- **Practical significance:** as highlighted by the authors, efficient training of large models is a timely problem, the ideas presented in this work are relatively simple but proven effective in realistic scenarios. Overall, I believe that the proposed approach has the potential to be integrated in classical schemes for large-scale training.", "weaknesses": "The main points of discussion of this submission are in my opinion related to:\n- **Novelty/Fit for ICLR:** the proposed techniques are not strictly novel in terms of research. Indeed, while the approach itself is technically new (e.g. at the best of my knowledge has not been proposed in previous works) it uses a combination of state-of-art techniques, with a strong engineering focus on how to integrate them in the most effective way. For example, the proposed _Dynamic Range Expansion_ is a simple expansion of the domain values of the optimizer states to the range of the used representation, and the _Mixed-Granularity Activation Quantization_ uses VS-Quant and Per-Block Quant for non-linear layers and per-tensor quantization for linear layers, which are not themselves new techniques.\n- **Missing implementation details:** given that the work is purely experimental and that reproducing the results strongly relies on efficient implementation, one would expect more implementation details. The provided details are sufficient for an high-level understanding of the proposed approach but still require significant effort to reproduce the results. Ideally, one would like to have a working code attached to the submission, such that reviewers can try it. Indeed, at the current state, it is not possible for a reviewer to reproduce the results.\n\nIn essence, I believe that this paper has mainly (good) incremental engineering contributions, while I have some concerns about its contributions in term of research ideas. Overall, I do question if it is a good fit for ICLR, rather than the validity of the approach or its practical usefulness.\nI provisionally give a \"Borderline accept\" and wait for the next phase to discuss with other reviewers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zzX9OcPL7D", "reviewer": "Reviewer_J45S", "summary": "This work presents COAT, a new framework for FP8 training. There are two key designs. The first is a dynamic range expansion that tries to align the value range of optimizer states to that of FP8 representation, which improves the utilization of the FP8 value range and reduces quantization errors. The second is a mixed-granularity activation quantization, which combines (less nuanced) per-tensor quantization to linear layers and (more nuanced) per-group quantization to non-linear layers. Experimental results show that COAT achieves nearly the same convergence as BF16 training, yet is more time and memory efficient.", "questions": "(Q1) Dynamically scaling the value been has been proposed for long. For instance, almost all integer quantization approaches will scale the values with the max value or the norm of tensors. I suggest the authors discuss the novelty of their work, such as why exponential expand function is chosen.\n\n(Q2) Per-group quantization can be done by first slicing the original tensor and then performing a per-slice quantization, which is implementation-specific. For me, a more important problem is how to determine the group size automatically. Currently there are no related discussions. And it is unclear how group sizes are determined in the experiments (e.g., based on hyper-parameter tuning or adaptive selection). If it is manually tuned, a sensitivity experiment would be nice to have.\n\nBesides, in lines 83-85, the authors said that “per-tensor quantization for matrix multiplications is more efficient and better suited for TensorCores, while fine-grained quantization helps maintain accuracy.” I was expecting an empirical assessment on this, but failed to see any. Intuitively, fine-grained per-group quantization should always achieve lower quantization errors, then we should always prefer it to the per-tensor one if the efficiency gap is not significant. \n\n(Q3) The convergence experiments are incredible. However, it would be better to break down the improvement of each design in your work. In addition, breaking down the speedup and memory reduction for different modules in the model can help readers better understand the effectiveness. \n\n(Q4) FP8-LM, one of the most relevant works, is only briefly discussed, while not included in the analysis and experiments in the main body. \n\n(Q5) In lines 333-334, the authors said that the max reduction on each $1 \\times G$ “can be seamlessly fused with the previous operation, adding minimal overhead”. Can you provide more details about how this is achieved?\n\n(Q6) Minor comments:\n- Line 192: “can not” should be “cannot”.\n- Footnote 2: “RoPE and FlashAttention is …” should be plural.\n- In Table 2 and Appendix, it might be more appropriate to set 1U to be FP8, rather than BF16.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written, well elaborating the problems of FP8 training as well as the proposed techniques to address the problems. Experiments on LLM pre-training, LLM fine-tuning, and VLM SFT demonstrate the sound convergence of COAT. And positive results present the improvement in terms of memory reduction and training speedup.", "weaknesses": "- The novelty of this work needs a justification. \n- The ablation studies do not break down the improvement of each component.\n- All discussions about FP8-LM are textual, lacking empirical comparisons in Section 6 and formal analysis in Table 2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)