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
Primary area: foundation or frontier models, including LLMs
Keywords: FP8 training, quantization, low-precision training, memory efficient training

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730017026954

thread_type:
official_review

<thread_object>
creator: reviewer_j45s

<content>
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.

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. 
- The ablation studies do not break down the improvement of each component.
- All discussions about FP8-LM are textual, lacking empirical comparisons in Section 6 and formal analysis in Table 2.

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.

(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.

Besides, 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. 

(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. 

(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. 

(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?

(Q6) Minor comments:
- Line 192: “can not” should be “cannot”.
- Footnote 2: “RoPE and FlashAttention is …” should be plural.
- In Table 2 and Appendix, it might be more appropriate to set 1U to be FP8, rather than BF16.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731087416345

thread_type:
official_review

<thread_object>
creator: reviewer_gvqo

<content>
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.

Extensive 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.

strengths:
In my opinion, this paper has the following main strenghts:
- **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.
- **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:
- **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.
- **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.

In 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.
I provisionally give a "Borderline accept" and wait for the next phase to discuss with other reviewers.

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.
- 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.

Small issues:
- There is a missing ref. in line 372: "performance in [] 4".
- 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731095229989

thread_type:
official_review

<thread_object>
creator: reviewer_p9h9

<content>
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.

strengths:
- The experiments demonstrate significant space improvements in practice and slight speed up.
- Figures illustrates methodology and motivation clearly.
- 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.
- Table 2 shows the actual improvements in memory clearly.
- 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.

The 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.

Another 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.

Finally, 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.

[1] Balanced Quantization: An Effective and Efficient Approach to Quantized Neural Networks. Zhou et al. 2017.

questions:
None.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731284381624

thread_type:
official_review

<thread_object>
creator: reviewer_4rzz

<content>
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.

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.
2. Dynamic range expansion is well-motivated by the analysis from Figure 2.
3. 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.
2. It is unclear to me how/whether the dynamic range is synchronized across GPUs in the distributed training setting.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732326872734

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers

comment:
We thank all reviewers for the constructive advice and valuable feedback! We are happy that all reviewers appreciate the practical value of our work to the LLM community. Here, we would like to highlight the algorithm-wise contributions of COAT:

- We **identify the range under-utilization problem** for optimizers, which prevents the optimizer states from being quantized to FP8. While previous work [1] shows that naively quantizing optimizers leads to performance degradation, COAT's **dynamic range expansion** addresses this problem. It significantly reduces training memory by 4x while maintaining on-par performance, as shown in Table 3 / Figure 5.

- In addition to empirical experiments, we also **theoretically prove** that our polynomial expansion function is optimal for quantization (in Appendix E). This neat theory results translate into excellent training results—COAT's performance closely matches the BF16 baselines across various LLM and VLM tasks.

- COAT provides the first adaptive quantization flow for LLM training. It effectively reduces the quantization error by employing different quantization strategies for different layers, making it an important step to transform most training dataflows into FP8 without losing accuracy.

In summary,  COAT delivers significant value to the LLM community by enabling FP8 quantization to be effectively used in real-world training scenarios. Its contributions directly address the scalability and cost-efficiency challenges researchers and practitioners face by allowing large language models (LLMs) to train effectively on setups with fewer GPUs while also having a high training throughput. We have also released our codebase to anonymous GitHub to benefit the community (https://anonymous.4open.science/r/COAT-Anonymous-E13B/).  

[1] Peng, Houwen, et al. "Fp8-lm: Training fp8 large language models." arXiv preprint arXiv:2310.18313 (2023).

## Proof the expand function is optimal

We prove that the choice of $f(x) = x^k$ is optimal given certain assumptions for the expansion function $f(x)$.

1. **$f(x)$ is an odd function**, where $f(x) + f(-x) = 0$, since the FP8's range is symmetric. Therefore, in the following discussion, we only consider $x \geq 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 or expansion.
4. **$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.
5. **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.$

We prove that based on these assumptions, the only function that meets these assumptions is $f(x) = \text{sign}(x) |x|^k$.

### 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{\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.  (1)
$$
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(2^{-t}) = m^{-t}.  (2)$$  
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\left(\prod_{i=-N}^{N} 2^{c_i \cdot 2^i}\right) = \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.  (3)
$$
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 exist two series $a_i, b_i, -N \leq i \leq N$ that satisfy 
$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 $2^{t_a}$ and $2^{t_b}$ can be used to approximate $f(x) = f(2^t)$:
$$
f(x) = \lim_{N \to \infty} f(2^{t_a}) = \lim_{N \to \infty} f(2^{t_b}).  (4)
$$
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)},  (5)
$$
where $f(2) = m$ is a constant. If we denote $\log_2(m) = k$, then $f(x) = x^k$.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces COAT, an FP8 training framework that improves memory efficiency and training speed for large-scale models while maintaining BF16-level performance. Key contributions include dynamic range expansion, aligning optimizer states to FP8 value ranges, and mixed-granularity activation quantization, combining per-tensor and per-group strategies. While the techniques are incremental and lack substantial novelty, their practical significance is clear, with strong experimental results and potential for adoption by practitioners. The paper could benefit from more detailed ablations and comparisons with additional relevant related works as requested by reviewers. Overall, I recommend borderline acceptance for its practical contributions to efficient training.