COAT: Compressing Optimizer states and Activations for Memory-Efficient FP8 Training

We thank the reviewer for acknowledging our novel contributions and insightful questions. Below we respond to the questions. We will revise the paper's final versions based on the reviewers' comments.

Q1: Justification of the novelty of this work

A1: COAT's key contributions include:

Optimizer Quantization:

• COAT first identifies the dynamic range under-utilization problem for optimizer state quantization, which is crucial for optimizer states quantization. We proposes the dynamic range expansion method to address the problem, and prove that our polynomial expand function is the theoretically optimal expand function. We also empirically verify that our choice can significantly reduce the MSE and do not lead to performance degradation.

Activation Quantization:

• Propose to use distinct quantization policies for linear and non-linear layers to best preserve accuracy and optimize memory efficiency, addressing gaps left by previous work that primarily focused on compute efficiency. Previous FP8 training works only focus on compute efficiency and overlook the importance of memory efficiency. While such insights have been extensively studied in inference, they are rarely explored in training due to the complexity of the training workflow and error accumulation.

COAT’s innovations address critical challenges in FP8 quantization and system optimization. COAT democratizes LLM training, reducing GPU requirements while maintaining high throughput, making it highly practical for real-world scenarios.

Q2: Break down the improvement of each component in ablation study.

A2: Thanks for the constructive advice! We present a more detailed analysis for Table 7 to do an ablation study on the improvement of each component. We conduct experiments on Llama-2-7B and vary the number of devices from 2, 4, and 8.

We find that quantizing the optimizer states to FP8 is more helpful when the number of GPU <= 4, and quantizing the activation to FP8 is more beneficial when the number of GPU >= 8.

Q3: Empirical comparisons about FP8-LM in Section 6 and Table 2.

A3: We would like to clarify that COAT is orthogonal to FP8-LM. FP8-LM mainly quantizes the (1) gradients to FP8 to reduce communication overhead, (2) parameters to FP16 to save memory, and (3) first-order momentum to FP8 to save memory. In contrast, we mainly focus on quantizing the first-order and second-order optimizer states and activations to FP8.

We updated Table 2 and observed that FP8-LM has an equivalent effect on activation memory as TE. As shown in Table 2, COAT delivers 1.51x higher activation memory saving over FP8-LM. It is possible to combine COAT with FP8-LM to boost training efficiency further. However, it is out of the scope of this work and we leave it to future research.

Q7: Details about the max reduction on each 1×G group in our Group Scaling design “can be seamlessly fused with the previous operation, adding minimal overhead”

A7: We take the RMSNorm + Linear as an example. Input size is (M, N), group size is G. In our method, the RMSNorm kernel has 6 parts:

1. Load the FP8 input of size (1, N) and BF16 scaling factor of size (1, N // G) to chip
2. Dequantize the input to float32
3. Calculate the RMSNorm result
4. Calculate the maximum value for every G value, to get a tensor of shape (1, N // G).
5. Store the output of step 3 to HBM.
6. Store the maximum value in step 4 to HBM

The overhead we introduce is step 4 and step 6. If we fuse these 2 steps, the latency of the RMSNorm kernel only increases by less than 5%.

Q8: Minor comments:

A7: Thanks for the suggestion and we have revised accordingly.

Q4: Dynamically scaling the value has been proposed for a 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 the exponential expansion function is chosen.

A4: To our best knowledge, the range underutilization problem of optimizer state quantization is first identified in COAT [1] [2]. This under-utilization problem cannot be resolved by traditional quantization methods, which typically scale tensors by dividing them with a scale factor based on their maximum value. While this normalizes the tensor, it does not alter the ratio between the maximum and minimum values (i.e., the dynamic range), leading to significant quantization errors. We also empirically verify that this method decreases the MSE of m / \sqrt{v} by 1.63x, as shown in Table 1 in our paper.

For the expansion function chosen by COAT, We mathematically prove that our polynomial expand function $f(x) = x^k$ is optimal and attach the proof in Appendix G. Importantly, if the expansion function is not chosen carefully, the MSE error will become suboptimal.

[1] Dettmers, Tim, et al. "8-bit optimizers via block-wise quantization." arXiv preprint arXiv:2110.02861 (2021).

[2] Li, Bingrui, Jianfei Chen, and Jun Zhu. "Memory efficient optimizers with 4-bit states." Advances in Neural Information Processing Systems 36 (2024).

Q5: How group sizes are determined in the experiments

A5: We choose to set the group size = 16 for activation quantization. This choice is based on balancing the memory efficiency overhead and the quantization error. With a smaller group size, the quantization error is reduced but will introduce high memory overhead on the scaling factors (bfloat16).

We calculate the quantization error of the output of the last transformer layer in the forward pass, and the gradient before the first transformer layer in the backward pass, compared with BF16. Based on the Table below, we find that group size = 16 achieves a good balance for memory efficiency and quantization error. We include these results in Appendix H to improve our paper.

Q6: Compare per-tensor quantization and per-group quantization for matrix multiplications

We find that per-group quantization is not suitable for matrix multiplication since it makes linear layers much slower because of frequent quantized and dequantize operations, and the accuracy gain is not significant. This phenomenon is also observed in Qserve[1].

Meanwhile, using very fine-grained quantization for linear layers does not bring much benefit compared with per-tensor quantization in terms of MSE, as reported in the table above in Q5. Therefore we do not apply per-group quantization to linear layers.

[1] Lin, Yujun, et al. "Qserve: W4a8kv4 quantization and system co-design for efficient llm serving." arXiv preprint arXiv:2405.04532 (2024).

Dear Reviewer J45S,

We are looking forward to further communication and please let us know if our response have fully addressed your concerns.

Sincerely,

Authors

Q5: Whether polynomial activations can have the same benefits, for example.

A5: Applying dynamic range expansion to activation will significantly slow the training, but the accuracy gain is relatively marginal. We report the latency of the quantization operator, which demonstrates a nearly threefold increase in quantization overhead. And even if we apply dynamic range expansion to activation, the quantization error is not reduced significantly. Consequently, we opted not to adopt the dynamic range expansion technique to activations in our method.

Q6: Justifying the engineering choices, such as fine-grained ablation study or theoretical analysis of the mixed-grained quantization

A6: To balance overhead and accuracy, we set the group size to 16. We explored group sizes ranging from 4 to 64 to compare their respective overheads and quantization errors, as shown in the table below.

When the group size of FP8 quantization is 16, each quantization group requires 8×16+16=144 bits in total, resulting in 9 bits per element. This introduces a 12.5% memory overhead compared to the ideal 8-bit-per-element scenario. The quantization error is measured using the Mean Squared Error (MSE) of the output values in the last layer between BF16 and FP8 training.

Smaller group sizes (e.g., 4) yield lower quantization errors but incur higher overhead. Larger group sizes (e.g., 64) reduce overhead but increase forward or backward quantization errors. Based on these statistics, we chose to set the group size to 16.

We add the related discussions to Appendix H to improve our paper.

Q7: More careful literature review for dynamic range under-utilization problem

A7: Thanks for the suggestions. We have added more literature [1-3] in lines 141-148, Section 2 in our updated paper and will further revise it in our final version.

[1] Zhou, Shu-Chang, et al. "Balanced quantization: An effective and efficient approach to quantized neural networks." Journal of Computer Science and Technology 32 (2017): 667-682.

[2] Wang, Shuai, and Yi Kang. "Gradient distribution-aware INT8 training for neural networks." Neurocomputing 541 (2023): 126269.

[3] Liu, Zechun, et al. "Nonuniform-to-uniform quantization: Towards accurate quantization via generalized straight-through estimation." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

We thank the reviewer for acknowledging our contributions and the insightful comments. We respond to the questions below. We will revise the paper's final versions based on the reviewers' comments.

Q1: Clarify the scientific contribution of this work:

A1: While we appreciate the reviewers recognizing the “scientific metric/pragmatic approach” of COAT, we would like to politely yet firmly emphasize the algorithm-wise contribution of COAT:

• Dynamic range expansion that firstly identifies the range under-utilization issue and designs theoretical optimal expansion function (General response for details).

• Different quantization policies for linear/non-linear layers best mitigates the accuracy gap and maximize the memory efficiency while previous work only focuses on compute efficiency.

COAT has not only theoretically proven the optimal expansion function, but also extensively verified the performance on a wide range of tasks, as well as the choice of quantization strategy for different layers. We hope that reviewers can provide a holistic view of COAT’s scientific contribution.

Q2: Little interest in the community of “learning representations”,

A2: COAT as an efficient training technique is a good fit for the ICLR community. While the ICLR community initiates from the interests of “learning representations”, From ICLR 2025 Call for papers [1], one important relevant topic of ICLR is

• “infrastructure, software libraries, hardware, etc” and “considers a broad range of subject areas including ……. applications in vision, audio, speech, language …… in ML.”

In the past acceptance of ICLR, awesome projects such as LBA Inference [2], ZeRO++ [3], and LUQ [4] have demonstrated the importance of extensive system optimizations in identifying the best practices for real-world scenarios. COAT proposes a novel FP8 training framework for large-scale model training and democratizes the language models, thus surely fitting the scope of the ICLR conference.

[1] https://iclr.cc/Conferences/2025/CallForPapers

[2] Towards Cheaper Inference in Deep Networks with Lower Bit-Width Accumulators. https://openreview.net/forum?id=oOwDQl8haC

[3] ZeRO++: Extremely Efficient Collective Communication for Large Model Training. https://openreview.net/forum?id=gx2BT0a9MQ

[4] Accurate Neural Training with 4-bit Matrix Multiplications at Standard Formats. https://openreview.net/forum?id=yTbNYYcopd

Q3: studied or tried other invertible transformations that spread out the values nicely across the representable values.

A3: Thanks for the insightful question. In fact, the transformation (expansion function) has to be chosen carefully. In COAT, We theoretically prove that the choice of $f(x) = x^k$ is optimal for the expansion function (details in newly added Appendix F). Other invertible transformations are suboptimal in theory and lead to larger MSE errors as we will show below.

Following the suggestion, we also test other invertible transformations. We choose the $\frac{e^{kx} - 1}{e^k - 1}$ function and use $k$ to control the strength. To avoid numerical problems, we divide $x$ by its maximum value before expansion, therefore the input is restricted to (0, 1]. We find that the quantization error of the exponential function underperforms our polynomial expansion function.

Q4: Would direct training in this transformed low-bit regime work?

A4: We find that direct training when representing optimizer states with FP8 does not work for training, and the loss curve will quickly diverge after even 100 steps. The scaling factor and the expand function are crucial to stabilize the training process.

Q5: Main differences between the proposed Group Scaling and the TransformerEngine's Delayed Scaling

A5: Delay scaling estimates the scaling factor of recent activations history, while our group scaling uses the current step’s activation to calculate the scaling factor. Our group scaling and delay scaling both reduce the latency when calculating the scaling factor.

In a linear layer, suppose we train for $N$ steps, with activations $X_1, \dots, X_N$. When the activation $X_{N+1}$ arrives, delay scaling calculates the scaling factor for $X_{N+1}$ by recording the maximum value from a sliding window of recent activations, $X_{N-w}, \dots, X_N$, where $w$ is the window size. Our group scaling will calculate the scaling factor based on the maximum value of $X_{N+1}$.

We propose group scaling instead of delayed scaling because of stability issues, which will expand below.

Q6: Delayed Scaling causes instabilities, but it is not clear what they are and why they arise.

A6: In pretraining scenarios, large activation spikes can produce unusually high activation values. If scaling factors are determined solely based on historical statistics, they may be insufficient to accommodate the larger values of XN+1X_{N+1}, which demand a significantly higher scaling factor. This mismatch results in excessive quantization during the current step, causing substantial quantization errors in the training process.

In contrast, our method uses the statistics of $X_{N+1}$ to calculate an appropriate scaling factor, ensuring accurate quantization and mitigating the impact of these activation spikes.

Q3: More discussion about implementation details

A3: We have uploaded our codebase to anonymous GitHub for preview: https://anonymous.4open.science/r/COAT-Anonymous-E13B/ and will open source when less anonymous. The COAT LLaMA implementation introduces three FP8-optimized modules: BeforeAttention, AfterAttention, and MLPResidual. These modules, built using torch.autograd.Function, ensuring accurate gradient flow computation. LayerNorm input quantization is fused with the residual connection to boost performance.

Custom FP8 matrix multiplication kernels for linear layers are implemented with Triton and optimized using triton.autotune, with M, N, and K dimensions tuned from {128, 256} based on hardware. Group scaling reshapes outputs to (BLOCK_M, BLOCK_N//G, G) before linear layers, computes scaling factors for each group, and stores them efficiently.

Non-linear layers also leverage Triton-based FP8 kernels. The forward pass generates quantized FP8 and transposed FP8 outputs, critical for efficient backward computations requiring column-major inputs. This avoids costly reordering during backpropagation, enhancing performance. Block sizes {32, 64} are autotuned to ensure optimal performance across all FP8 computations.

To replicate our result using our anonymous code, you can first install coat and the fp8 optimizer kernel. Then download the dataset of OLMo.

To reproduce the memory reduction result, run:

BF16 baseline:

MEMORY_BENCH=1 torchrun --nproc_per_node=4 scripts/train.py configs/reproduce/OLMo-7B-reproduce-MemBench.yaml

COAT:

MEMORY_BENCH=1 torchrun --nproc_per_node=4 scripts/train.py configs/coat/OLMo-7B-COAT-Both-MemBench.yaml

To reproduce the memory reduction result, run the following:

BF16:

SPEED_BENCH=1 torchrun --nproc_per_node=4 scripts/train.py configs/reproduce/OLMo-7B-reproduce-SpeedBench.yaml

COAT:

SPEED_BENCH=1 torchrun --nproc_per_node=4 scripts/train.py configs/coat/OLMo-7B-COAT-Both-SpeedBench.yaml

To reproduce the OLMo pretraining experiment, run: torchrun --nproc_per_node=8 scripts/train.py configs/coat/OLMo-1B-COAT-BOTH.yaml

Q4: Difference with other related works about activation quantization?

A4: Recent activation quantization approaches collectively focus on the memory saving but ignore the compute efficiency, and most of them are designed for neural networks and not optimized for recent large language models. COAT’s design the activation quantization with FP8 precision flow, thus enjoying both memory- and compute-efficiency. COAT also extensively verifies its accuracy on LLM related tasks.

For example, ActNN[1] introduced a 2-bit activation compressed training framework using random quantization, achieving a 12x reduction in activation memory footprint1. GACT[2] extended this concept to support various machine learning tasks and architectures, providing up to 8.1x reduction in activation memory. However, these methods mostly focus on convolutional networks and are not applicable to LLMs. AQ-SGD[3] proposed compressing activation changes rather than direct values for pipeline-parallel training to reduce the communication overhead. Few-Bit Backward[4] quantizes the non-linear activation functions using optimal piecewise-constant approximations to reduce the memory footprint of activation functions.. Jetfire[5] 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. We modify lines 127-138 to discuss these related works in detail.

[1] Chen, Jianfei, et al. "Actnn: Reducing training memory footprint via 2-bit activation compressed training." International Conference on Machine Learning. PMLR, 2021. [2] Liu, Xiaoxuan, et al. "Gact: Activation compressed training for generic network architectures." International Conference on Machine Learning. PMLR, 2022. [3] Wang, Jue, et al. "Fine-tuning language models over slow networks using activation quantization with guarantees." Advances in Neural Information Processing Systems 35 (2022): 19215-19230. [4] Novikov, Georgii Sergeevich, et al. "Few-bit backward: Quantized gradients of activation functions for memory footprint reduction." International Conference on Machine Learning. PMLR, 2023. [5] Xi, Haocheng, et al. "Jetfire: Efficient and accurate transformer pretraining with int8 data flow and per-block quantization." arXiv preprint arXiv:2403.12422 (2024).

We thank the reviewer for acknowledging our novel contributions and insightful questions. Below we respond to the questions. We will revise the paper's final versions based on the reviewers' comments.

Q1: Novelty/Fit for ICLR

A1: From ICLR 2025 Call for papers [1], one important relevant topic of ICLR is

• “infrastructure, software libraries, hardware, etc” and “considers a broad range of subject areas including ……. applications in vision, audio, speech, language …… in ML.”

COAT proposes a novel FP8 training framework for large-scale model training and democratizes the language models, thus undoubtedly fit the scope of ICLR conference. In past acceptance of ICLR, awesome projects such as LBA Inference [2], ZeRO++ [3], and LUQ [4] have demonstrated the importance of extensive system optimizations in identifying the best practices for real-world scenarios. The insights brought to the community are equally important as the techniques used.

Further, 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 faced by researchers and practitioners by allowing large language models (LLMs) to train effectively on setups with fewer GPUs while also having a high training throughput.

Q2: combination of state-of-art techniques, with a strong engineering focus, xxx. which are not themselves new techniques

A2: We arguably disagree with the comments and clarify as follows:

The dynamic range expansion is motivated by a clear rationale, addressing the under-utilization of the dynamic range in optimizer states. Without this expansion, significant quantization errors can occur. If the expand function is not chosen carefully, it may also result in sub-optimal results.

In COAT, the polynomial expansion comes from a theoretical formula and we have proved its optimal for quantization optimizer states (details in Appendix F). To our best knowledge, the formula, as well as the range under-utilization is first discussed for optimizer states.

Applying different quantization policies for different layers is rarely discussed and most previous studies have focused on the inference. Existing solutions such as TransformerEngine[5], FP8-LM[6], and FP8 to Trillion Token[7] only use per-tensor quantization for linear layers and do not quantize non-linear layers to best preserve accuracy. COAT makes an important step to transform most dataflows into FP8 for better efficiency.

Therefore, we firmly believe that COAT will also make a valuable contribution to the conference and inspire future low-bit training researchers.

[1] https://iclr.cc/Conferences/2025/CallForPapers

[2] Towards Cheaper Inference in Deep Networks with Lower Bit-Width Accumulators https://openreview.net/forum?id=oOwDQl8haC

[3] ZeRO++: Extremely Efficient Collective Communication for Large Model Training https://openreview.net/forum?id=gx2BT0a9MQ

[4] Accurate Neural Training with 4-bit Matrix Multiplications at Standard Formats: https://openreview.net/forum?id=yTbNYYcopd

[5] https://github.com/NVIDIA/TransformerEngine

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

[7] Fishman, Maxim, et al. "Scaling FP8 training to trillion-token LLMs." arXiv preprint arXiv:2409.12517 (2024).

Dear Reviewer gVQo,

We are looking forward to further communication and please let us know if our response have fully addressed your concerns.

Sincerely,

Authors

Dear Reviewer gVQo,

We truly appreciate your time and effort in reviewing our submission and offering valuable suggestions. The author-reviewer discussion period is closing soon, so we would like to confirm whether our responses have effectively addressed your questions and concerns.

Thank you for your time and insights.

Best regards,

Authors

Dear Reviewer gVQo,

Thank you again for your detailed suggestions and comments. As the deadline of discussion period is approaching, could you kindly review our responses and inform us whether you have further questions?

Thank you!

Sincerely,

The Authors

Thanks for your valuable comments and feedback. We respond to the questions below.

Q1: COAT still underperforms bf16 training on most tasks

A1: COAT achieves on-par performance with BF16, and we would like to clarify the misunderstanding regarding COAT’s performance. For LLM results, 3 out of 8 cases outperform BF16 by an average of 1.0%, 3 out of 8 underperform by an average of 1.3%, and 2 are comparable to BF16. Similarly, for VLM, 2 out of 9 cases outperform BF16, 3 out of 9 underperform, and 4 are comparable. Overall, COAT achieves accuracy on par with BF16.

Q2: How each component influences the model performance

A2: This is a constructive suggestion! We separately quantize the optimizer and activations to better understand how each component influences the model performance. We find that quantizing activation incurs almost no degradation due to our well-designed quantization flow, and the main degradation comes from optimizer quantization. We choose to quantize both to maximize memory saving.

Q3: How/whether the dynamic range is synchronized across GPUs in the distributed training setting.

A3: To clarify, the dynamic range is NOT synchronized across GPUs. COAT incurs no communication overhead. During distributed training, only gradients are communicated between GPUs while activations and optimizers stay on each GPU and are never sent across. Therefore, we calculate the dynamic range based on the statistics of the optimizer states on each device individually, eliminating the need for cross-GPU synchronization. For FSDP related implementation, we add more details in Section 3, lines 198-200 and Appendix A.

Dear reviewer 4rZz,

We are looking forward to further communication and please let us know if our response have fully addressed your concerns.

Sincerely,

Authors

Dear Reviewer 4rZz,

We truly appreciate your time and effort in reviewing our submission and offering valuable suggestions. The author-reviewer discussion period is closing soon, so we would like to confirm whether our responses have effectively addressed your questions and concerns.

Thank you for your time and insights.

Best regards,

Authors

Dear Reviewer 4rZz,

Thank you again for your detailed suggestions and comments. As the deadline of discussion period is approaching, could you kindly review our responses and inform us whether you have further questions?

Thank you!

Sincerely,

The Authors

Thank you for providing this proof. I have a few comments:

• The argument for $x=2^{-t}$ mistakenly uses $x^t=1$, but $x^t=2^{-t^2}$. Though this can be easily fixed as follows: $f(x) = \frac{f(x)}{f(1)} = \frac{f(x x^{-1})}{f(x^{-1})} = f(x^{-1})^{-1}$, and since $f(2^t)=m^t$, we have $f(2^{-t})=f(2^{t})^{-1} = m^{-t}$.
• The variables $x_a$ and $x_b$ were introduced without explanation. I believe the definitions for $t_a$ and $t_b$ were intended for them instead.
• The claim that $f(x)=\text{sign}(x)|x|^k$ is the only function (or optimal function) that meets the assumptions should be made more precise (at least in the revision). What does optimality mean here? The claim may hold on $x \in \mathbb{R}$, but does this imply that it also holds on the constrained domain $x = \sum_{i=-N}^N a_i 2^i$ for all $a_i \in \{0, 1\}$ and finite $N$? This is just a technicality and I believe that the claim for the unconstrained case suffices for the purposes of this paper, but it could be the case that there exists a better (or another) $f$ for the constrained case.

Also, thank you for providing the anonymous and professionally-written code.

Thank you for the comment. We noticed some minor mistakes in our original proof. We have fixed the issues and updated our manuscript. The conclusion remains the same.

Q1 & Q2: Thank you for pointing this out. The correct proof is updated in Equation (2):

$f(x) = \frac{f(1)}{f(\frac{1}{x})} = \frac{f(1)}{f(2^t)} = \frac{1}{m^t} = m^{-t},\text{so that } f(x) = m^{-t}.$

We intend to use $t_a$ and $t_b$ for this part. We also correct these terms.

Q3: We also updated the proof after Equation (3) consequently. The binary representation should be expressed as $x = 2^t, t = \sum_{i=-\infty}^\infty c_i \cdot 2^i, c_i \in \{0, 1\}$. We replace $a_i$ with $c_i$ to avoid the repeated use of letters, and the binary representation should be discussed in terms of t.

The intention of “optimality” was to argue that $f(x) = \operatorname{sign} (x) |x|^k$ is the only function that satisfies all the 5 assumptions on the unconstrained domain $x \in \mathbb{R}$.

On the constrained domain where $x = 2^t, t = \sum_{i=-N}^N c_i \cdot 2^i, c_i \in \{0, 1\}$, our claim also holds. Equation 3 has already demonstrated that in this constrained domain, we have $f(2^t) = m^t$, where $m = f(2)$. This means that the only function on this constrained domain that satisfies all the 5 assumptions should have $f(x) = x^k, where k = \log_{2}{m}$. This proves the optimality of this function.