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

## Abstract

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. The code is available at https://github.com/NVlabs/COAT .

## 1 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., ; Adler et al., ; Team et al., ; Lin et al., ). However, the training of such models, which often comprise billions of parameters, demands substantial computational resources and memory. This presents substantial challenges, making the training of these foundation models very challenging (Smith et al., ; Hoffmann et al., ).

Low-precision training has emerged as a promising approach to make FMs training more efficient (Micikevicius et al., ; Wang et al., ; Zhu et al., ; Xi et al., ; Wortsman et al., ; Xi et al., ). 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., ; Micikevicius et al., ) is the most prevalent low-precision method, and is widely adopted in large-scale training frameworks like DeepSpeed  (Rasley et al., ) and Megatron-LM (Shoeybi et al., ).

With the advent of Nvidia’s H100 GPU (NVIDIA, ), FP8 training Micikevicius et al. () is emerging as the next-generation low-precision technique. Compared to BF16, FP8 training has the potential to (1) double the speed and (2) halve the memory footprint. To achieve practical speedup,
Transformer Engine (NVIDIA, ) 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 1, FP8-LM (Peng et al., ) 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 second-order momentum in higher precision. The memory problem of activations becomes even more critical when optimizer, gradient, and weights are sharded across multiple GPUs using ZeRO or FSDP. Besides, second-order momentum is more sensitive to quantization than first-order momentum (Fishman et al., ), and activations’ large spikes also make them hard to quantize to FP8 (Yang et al., ). This potential accuracy degradation makes them missing a crucial opportunity to 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 overall memory footprint by quantizing optimizer states and activations into FP8.* For optimizer states, we observe that FP8 format’s *representation range is under-utilized* when quantizing them, as illustrated in Figure 2(a). To address this, we introduce a novel *Dynamic Range Expansion* method which adjusts the distribution of optimizer states to better fit within the FP8 range, thereby minimizing quantization error. For activations, we propose *Mixed-Granularity Activation Quantization* to achieve efficient and accurate quantization. We apply fine-grained quantization to non-linear layers and apply per-tensor quantization to linear layers. Per-tensor quantization for matrix multiplications is more efficient and better suited for TensorCores, while fine-grained quantization helps maintain accuracy. These two approaches tackle high memory consumption while ensuring minimal performance degradation. We provide an overview of COAT in Figure 1(b) for demonstration.

We demonstrate the accurate performance of COAT on a wide range of tasks, including LLM pretraining, LLM fine-tuning, and VLM training. COAT achieves nearly lossless performance on all of these tasks. For efficiency results, COAT achieves $1.54\times$ end-to-end memory reduction compared with BF16, and $1.43\times$ 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 a more efficient training process for large-scale models.

## 2 Related Work

##### Low-precision Training

Low precision training  (Wang et al., ; Chen et al., ; Lin et al., ; Wortsman et al., ; Xi et al., ) has become a prominent technique in modern deep learning, offering reductions in both computational costs and memory requirements.
FP16 (half-precision) training  (Micikevicius et al., ) is the most prevalent low-precision method nowadays. It introduces loss scaling to address FP16’s narrower representation range problem. BF16 training  (Kalamkar et al., ) refines this approach, as BF16 has a larger representation range and is more stable for large-scale training. In these approaches, forward and backward passes are computed in FP16 or BF16 precision, while master weights, gradients, and optimizers are stored in FP32.

FP8 training  (Fishman et al., ; Micikevicius et al., ) 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, ) is the first framework designed for FP8 mixed-precision training that employs FP8 Tensorcore for linear layer calculation. FP8-LM  (Peng et al., ) extends FP8 quantization to gradients and optimizer states, further improving the training throughput. However, they fail to reduce the memory usage of activations stored for the backward pass using FP8, and leave second-order momentum in FP16, limiting the full potential of FP8’s memory advantages.

##### Memory Efficient Optimizers

While 32-bit optimizer  (Kingma, ; Loshchilov, ) states are widely adopted, several research efforts have been made to reduce the memory footprint of optimizer states through quantization. The 8-bit Adam  (Dettmers et al., ) introduces a novel data format called dynamic exponent (DE) for quantization, but the adoption of this new data format limits its flexibility. 4-bit Optimizer  (Li et al., ) further pushes the limit of optimizer quantization to 4-bit by addressing the zero-point problem, but is restricted to fine-tuning tasks. FP8-LM  (Peng et al., ) quantizes the first-order momentum to FP8 while leaving second-order momentum in FP16, which limits the overall memory savings.  (Fishman et al., ) finds that second-order momentum is more sensitive to quantization, and proposes to quantize it using E5M2 format.

In addition to quantization, there are other approaches that aim to reduce the memory footprint of the optimizer states  (Shazeer & Stern, ; Anil et al., ; Chen et al., ; Zhao et al., ), such as low-rank decomposition and optimizer simplification that only store the first-order momentum. These methods are orthogonal to our approach.

##### Activation Quantization

Activations consume a significant amount of memory during training Cai et al. (), with memory usage scaling in proportion to sequence length and batch size. ActNN (Chen et al., ) is a memory-efficient training framework designed to reduce the memory footprint by quantizing activations to 2 bits. Few-bit Backward (Novikov et al., ) focus on quantizing activation functions while achieving higher performance. Jetfire proposes INT8 data flow to quantize the activation of both linear and non-linear layers to reduce memory footprint.

## 3 Preliminaries

##### FP8 Quantization

Quantization compresses a high-precision tensor to low-precision to achieve speedup and save memory footprint, at the cost of lower precision and lower representation range.
*FP8 format* consists of two encodings - E4M3 and E5M2 (Open Compute Project, ). E4M3 has higher precision, while E5M2 has a larger representation range. We define E4M3’s min and max values as $\Delta_{\min}^{\text{E4M3}}=2^{-9}$ and $\Delta_{\max}^{\text{E4M3}}=448$, while E5M2’s min and max values are $\Delta_{\min}^{\text{E5M2}}=2^{-16}$ and $\Delta_{\min}^{\text{E5M2}}=57344$. To *quantize* an FP32 tensor $X$ into E4M3 precision, we use a quantizer $Q(\cdot)$ to map the tensor into FP8’s representation range. This process can be formulated as

$$ $X_{\text{FP8}},S_{X}=Q(X_{\text{FP32}}),\>\text{where}\>X_{\text{FP8}}=\left \lceil\frac{X_{\text{FP32}}}{S_{X}}\right\rfloor,\>S_{X}=\frac{\max\left( \lvert X_{\text{FP32}}\rvert\right)}{\Delta_{\max}^{\text{E4M3}}},$ $$

where $S_{X}$ is the scaling factor, $\lceil\cdot\rfloor$ refers to round-to-nearest. Quantize into E5M2 follows a similar procedure, where we only replace $\Delta_{\max}^{\text{E4M3}}$ with $\Delta_{\max}^{\text{E5M2}}$. To map the quantized tensor back to FP32 precision, the *dequantize* operation $DQ(\cdot)$ can be formulated as
$X_{\text{FP32}}=DQ(X_{\text{FP8}},S_{X})=S_{X}X_{\text{FP8}}.$

##### Optimizer Update Rule

Optimizers are widely used in deep learning to update parameters. The most common gradient-based optimizer is Adam/AdamW (Kingma, ; Loshchilov, ), which uses first-order $m$ and second-order momentum $v$ to achieve better convergence. The update rule of AdamW at time step $t$ can be formulated as:

$$ $\begin{aligned} m_{t}&=\beta_{1}m_{t-1}+(1-\beta_{1})g_{t-1}\\ \hat{m}_{t}&=\frac{m_{t}}{1-\beta_{1}^{t}}\\ \end{aligned}\quad\quad\quad\begin{aligned} v_{t}&=\beta_{2}v_{t-1}+(1-\beta_{ 2})g_{t-1}^{2}\\ \hat{v}_{t}&=\frac{v_{t}}{1-\beta_{2}^{t}}\end{aligned}$ $$

$$ $\displaystyle w_{t+1}=w_{t}-\eta\left(\frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}}+ \epsilon}+\lambda w_{t}\right)$ (1) $$

where $m_{t}$ is the first-order momentum, $v_{t}$ is the second-order momentum, $g_{t}$ is the gradient, $\beta_{1}$ and $\beta_{2}$ are betas of AdamW, $\eta$ is learning rate, $\lambda$ is weight decay, and $\epsilon$ is used to prevent NaN or Inf.

To perform optimizer state quantization, we adopt per-group quantization for both first-order and second-order momentum, following previous works (Dettmers et al., ; Li et al., ).
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. More details can be found in Appendix A.

Figure: 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 is well utilized. (c) Distribution of $k$ for optimizer states. The second order’s $k$ is larger than the first order’s $k$, since the second-order momentum’s dynamic range is smaller.
Refer to caption: x2.png

## 4 Dynamic range expansion for accurate optimizer quantization

In subsequent sections, we explain how COAT utilizes FP8 quantization to achieve memory-efficient FP8 training without compromising accuracy.
Section 4 focuses on optimizer states quantization,
while Section 5 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 is that, they can not *fully utilize the representation range* of FP8 and therefore lead to a large quantization error. Take the E4M3 data format as an example, the ratio between E4M3’s maximum representable value and minimum representable value is $\Delta_{\max}^{\text{E4M3}}/\Delta_{\min}^{\text{E4M3}}=448\div\frac{1}{512}=2
29376\approx 2\times 10^{5}.$ Therefore, for a quantization group $X$, if we want to fully utilize the 256 representable value(^1^11Actually among them a small amount of values represents NaN and Inf.) of FP8, we hope the dynamic range of the quantization group $X$ should cover the entire span between $\Delta_{\min}^{\text{E4M3}}$ and $\Delta_{\max}^{\text{E4M3}}$.

To make it more formally, we define *dynamic range* as the ratio between the maximum absolute value and the minimum absolute value within a quantization group $X$:

###### Definition 1 (dynamic range)

Given a set that consists of $G$ real numbers $X=\{x_{1},x_{2},\dots,x_{G}\}$, the dynamic range ${\mathcal{R}}$ is defined as

$$ ${\mathcal{R}}_{X}=\frac{\max(\lvert x_{1}\rvert,\lvert x_{2}\rvert,\dots, \lvert x_{G}\rvert)}{\min(\lvert x_{1}\rvert,\lvert x_{2}\rvert,\dots,\lvert x _{G}\rvert)},$ $$

where $\lvert\cdot\rvert$ denotes the absolute value.

That is to say, E4M3’s dynamic range is ${\mathcal{R}}_{\text{E4M3}}=448\times 512=229376\approx 2\times 10^{5}$. However, in practice, many quantization groups within the optimizer states fail to effectively map values across this wide range. We observe that optimizer states are highly sparse, with fewer than 1% of values having large magnitudes, while the majority are relatively small and closely clustered. Most groups exhibit low dynamic ranges since large values are so few. As visualized in Figure 2(a), the dynamic range for first-order momentum is typically less than 1e4, and for second-order momentum, it is usually less than 1e1—both far below the available range of FP8. As a result, a substantial portion of FP8’s representational capacity is wasted, leading to a large quantization error.

###### Definition 1 (dynamic range)

### 4.2 Dynamic Range Expansion

To address this problem, we introduce a *expand function* $f(\cdot)$ before quantization to expand the dynamic range of the quantization group and align it with E4M3, which can be formalized as $X_{\text{FP8}},S_{X}=Q(f(X_{\text{FP32}}))$. The expand function we use is defined as

$$ $f(x)=\operatorname{sign}(x)\lvert x\rvert^{k},$ $$

where $k$ is used to control the strength of the expansion.

For a quantization group $X$, after applying the expand function $f$ to $X$, the dynamic range becomes

$$ ${\mathcal{R}}_{f(X)}=\frac{\max(\lvert f(X)\rvert)}{\min(\lvert f(X)\rvert)}= \frac{\max(\lvert\operatorname{sign}(X)X^{k}\rvert)}{\min(\lvert\operatorname{ sign}(X)X^{k}\rvert)}=\Big{(}\frac{\max(\lvert X\rvert)}{\min(\lvert X\rvert)} \Big{)}^{k}=({\mathcal{R}}_{X})^{k}.$ $$

Therefore, when $k>1$, ${\mathcal{R}}_{X}$ will be enlarged and become closer to the ideal ${\mathcal{R}}_{\text{E4M3}}$. The optimal $k$ satisfy that $({\mathcal{R}}_{X})^{k}={\mathcal{R}}_{\text{E4M3}}$, which means that $k=\log_{{\mathcal{R}}_{X}}({\mathcal{R}}_{\text{E4M3}})$. With this optimal $k$, $f(X)$ can fully utilize the representation range of E4M3, while the original $X$ can only utilize a small portion of it. As shown in Figure 2(c), the second-order momentum typically has a larger $k$ value ($5\sim 15$) compared to the first-order momentum’s $k$ ($1\sim 3$). This corresponds to our previous observation, that second-order momentum usually has a smaller dynamic range compared with first-order momentum, so it requires a larger $k$ to align well with E4M3’s dynamic range.

We calculate $k$ on-the-fly for every optimizer step and for every quantization group for accuracy consideration. When dequantize, we apply the inverse of expand function $f^{-1}(x)=x^{\frac{1}{k}}$ after dequantization to recover its original value, which can be expressed as $X^{\text{FP32}}=f^{-1}(DQ(X_{\text{FP8}},S_{X}))$.

We apply regular quantizer and our dynamic range expansion method to both the first-order momentum $m$ and second-order momentum $v$. As visualized in Figure 3(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 1. In AdamW optimizer step, as stated in Eq. 1, $\frac{m}{\sqrt{v}+\epsilon}$ is the actual effective term for weight update, so we report the MSE of $\frac{m}{\sqrt{v}+\epsilon}$ to quantify the performance of a quantization method. We find that E4M3 is more suitable for first-order momentum than E5M2. For second order momentum, although E4M3 better than E5M2, their quantization error is nearly the same after applying our expand function. Our Dynamic Range Expansion can effectively reduce the MSE by $1.63\times$. Appendix A provide more results.

Figure: Table 1: Quantization error of $\frac{m}{\sqrt{v}}$ under different quantization settings. +Expand means applying our Dynamic Range Expansion method.
Refer to caption: x3.png

## 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 to calculate gradients. As illustrated in Table 2, non-linear layers such as LayerNorm/RMSNorm  (Ba, ; Zhang & Sennrich, ) and Activation Functions  (Hendrycks & Gimpel, ; Shazeer, ) typically account for approximately 50% of the memory footprint in the Llama model series Touvron et al. (). In contrast, linear layers contribute less than 25%. Therefore, it is essential to optimize both linear and non-linear layers to reduce activation memory footprint.(^2^22RoPE and FlashAttention is more as a distinct module and should be handled separately.)

One straightforward approach to achieve this is to quantize the activations to FP8 format prior to each non-linear and linear layer and only save the quantized tensors for backward. However, this introduces significant overhead due to the additional quantization step. Furthermore, non-linear layers can be sensitive to quantization. For example, GLU activation will amplify the spike in activations, resulting in significant quantization errors if not carefully handled (Yang et al., ; Fishman et al., ). This necessitates us to develop efficient and accurate quantization techniques.

### 5.2 Mixed granularity FP8 precision flow

**Table 2: Activation memory footprint of different operators. U is a unit to measure memory usage, where $\text{1U}=\text{Batch Size}\times\text{Sequence Length}\times\text{Hidden Size }\times\text{2 bytes (for BF16).}$ For Llama-style model, Act Func refers to SiLU & Multiply, and Linear refers to the summation of QKV/Attn/Up/Gate/Down projection. RMSNorm is upcast to float32 in transformers [implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L70), 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 Func | RoPE | FlashAttn | Linear | Total | Ideal | Achieved |
| Llama-style | BF16 | 4U | 8U | 2U | 3U | 5.66U | 22.66U | 1.00$\times$ | 1.00$\times$ |
| TE | 2U | 8U | 2U | 3U | 3.33U | 18.33U | 1.23$\times$ | 1.20$\times$ |  |
| COAT | 1U | 4U | 2U | 3U | 3.33U | 13.33U | 1.69$\times$ | 1.65$\times$ |  |

To address the inefficiency and inaccurate problem,
we 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 to vary the quantization granularity across different layers to balance precision and efficiency in a *mixed-granularity* manner. For non-linear layers, VS-Quant (Dai et al., ) or Per-Block Quant (Xi et al., ) methods are well-suited due to their fine-grained and precise nature. For linear layers, we apply per-tensor quantization to maximize the performance of Tensor Cores.(^3^33Finer-grained methods can also be applied since they are generally compatible with FP8 precision flow.)

We observe that quantizing the input of layernorm across multiple token axes is detrimental to accuracy. As illustrated in Figure 4(a), when the number of elements that share a scaling factor is fixed, the quantization error increases significantly when quantization is performed across the token axis. Therefore instead of using per-block quantization with block size $B\times B$ as proposed in  (Xi et al., ), we propose to use per-group quantization with group size $1\times G$, where $G=B^{2}$ to keep the granularity the same. This approach enhances the accuracy of non-linear layers while maintaining efficiency.
Our precise FP8 precision flow is visualized in Figure 1(a), where we display the full precision flow for a Llama-style decoder layer, both forward and backward pass.

Figure: Figure 4: (a) Quantization Error in forward pass. (b) Time comparison of various scaling methods.
Refer to caption: x4.png

We also propose Group Scaling, an efficient per-tensor scaling method that balances the performance and precision. To perform per-tensor quantization, the maximum absolute value of the tensor needs to be calculated through max reduction, adding a lot of overhead.

In our Group Scaling, we address these problems by splitting the max reduction into two stages: (1) performing max reduction on each $1\times G$ element and storing the results as intermediate values; (2) applying max reduction on the intermediate tensor to obtain the per-tensor max value. The first stage can be seamlessly fused with the previous operation, adding minimal overhead, while the second stage is more efficient than doing max reduction on the entire tensor, as the intermediate result is $G\times$ smaller than the original tensor. As illustrated in Figure 4(b), Group Scaling successfully reduces the max reduction overhead compared with just-in-time scaling.

In comparison, TransformerEngine proposes delayed scaling to avoid the on-the-fly max reduction required for per-tensor quantization, and can also fuse the division process into the previous operator to optimizes memory accesses.
However, we find that using the current tensor’s statistics to compute the scaling factor is no worse, and could be even better in precision, than using the delayed scaling heuristic. Therefore, we advocate for Group Scaling as a simpler and more flexible alternative that can potentially offer improved numerical stability and precision while not being significantly slower than Delayed Scaling.

## 6 Experiments

### 6.1 Accuracy Experiments

##### Setups

We compare COAT with BF16 training and TransformerEngine  (NVIDIA, ) baselines. For all experiments, we adopt the default hyperparameters in the official training recipe. We validate the effectiveness of our method on multiple tasks, including Large Language Model (LLM) pretraining and fine-tuning, and Vision Language model (VLM) training. For LLM pertaining, we report the perplexity on Wikitext
103 (Merity et al., ), C4 (Raffel et al., ), and Pile (Gao et al., ), and the accuracy on COPA (Gordon et al., ), ARC (Clark et al., ), SciQ (Welbl et al., ), and HellaSwag (Zellers et al., ). For LLM fine-tuning, we conduct experiments in math corpus, and evaluate on Mathmeticas (Davies et al., ), SVAMP (Patel et al., ), NumGLUE (Mishra et al., ),
and GSM8K (Cobbe et al., ).
For VLM training, we report the score on VideoMME (Fu et al., ), POPE (Li et al., ), VizWiz (Gurari et al., ), GQA (Hudson & Manning, ), VQAv2 (Goyal et al., ), TextVQA (Singh et al., ), SEED (Li et al., ), and MMMU Validation Set (Yue et al., ). We use $1\times 128$ per-group quantization for optimizer states and $1\times 16$ per-group quantization for non-linear layer activations.

#### 6.1.1 LLM pretraining

To evaluate our method when pretraining LLMs, We train OLMo-1B (Groeneveld et al., ) and OLMo-7B on Dolma (Soldaini et al., ). Following the official report, we use a global batch size of 4M tokens (2048 macro batch size, with a sequence length of 2048 tokens). We use PyTorch FSDP in our experiments. For OLMo-1B, we pretrain for 5500 steps, which corresponds to approximately 22B tokens. For OLMo-7B, due to the resource limitation, we perform continue pretraining for 1000 steps (about 4B tokens) and resume from the official checkpoint at 5000 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 were consistent with BF16 training and TransformerEngine baseline,
validating the effectiveness of COAT. For OLMo-7B experiment, we report the training curve in Figure 6, where COAT also aligns well with the baseline and is nearly lossless.

**Table 3: OLMo-1B pretraining performance on downstream tasks. TE refers to TransformerEngine.**
|  | Train Loss | WikiText | C4 | Pile | Avg ppl |
| --- | --- | --- | --- | --- | --- |
| BF16 | 2.995 | 30.499 | 27.966 | 17.405 | 25.290 |
| TE | 3.001 | 30.736 | 28.064 | 17.434 | 25.411 |
| COAT | 3.008 | 30.619 | 28.099 | 17.453 | 25.391 |
|  | COPA | ARC(Easy) | SciQ | HellaSwag | Avg Acc |
| BF16 | 60.0% | 45.6% | 67.3% | 33.7% | 51.6% |
| TE | 62.0% | 45.4% | 63.9% | 33.8% | 51.3% |
| COAT | 61.0% | 44.2% | 67.6% | 33.7% | 51.5% |

**Figure 6: OLMo-7B training curve.**
|  | Mathmeticas | SVAMP | NumGLUE | GSM8k | Avg |
| --- | --- | --- | --- | --- | --- |
| Llama-2-7B | 6.0 | 14.6 | 34.5 | 29.9 | 21.3 |
| BF16 | 46.3 | 64.2 | 54.8 | 57.7 | 55.7 |
| TE | 45.3 | 66.1 | 53.5 | 57.7 | 55.6 |
| COAT | 47.8 | 64.4 | 53.3 | 56.6 | 55.5 |

#### 6.1.2 LLM fine-tuning

We further evaluate our method on LLM fine-tuning task. We focus on math corpus and fine-tune a Llama-2-7B model on the MAmmoTH (Yue et al., ) dataset. We train for 3 epochs, and report the downstream tasks performance in  4. After fine-tuning, COAT still performs consistently with the baselines on the downstream task performance, proving the accurateness of our method.

#### 6.1.3 VLM Training

We also evaluate COAT on vision language models. We conduct experiments on VILA  (Lin et al., ) and perform stage-3 SFT of VILA1.5-7B using the same SFT data mixture employed by VILA’s original paper (Chen et al., ; Xu et al., ) and set the global batch size to 1024. We pad the sequence length to a multiple of 4 for efficiency consideration. The training loss curve is visualized in Figure 7. We report the downstream task performance and their average in Table 5. We find that COAT performs on par with BF16 training and is better than the TransformerEngine baseline, which demonstrates the accurateness of our method. We further visualize the VLM captioning experiment in Appendix B to prove the effectiveness of COAT on generation tasks.

Figure: Table 5: VILA1.5-7B Stage-3 SFT performance on downstream tasks. * means it has seen the training data.
Refer to caption: x7.png

### 6.2 Memory Saving and Speedup

We test the memory saving and speedup result of COAT in two settings: The results on a single transformer layer help to accurately analyze the capabilities for memory saving and speedup, while the end-to-end results reflect the practical benefits of our method in real-world applications.

#### 6.2.1 Memory Saving and Speedup for a Single Transformer Layer

Table 6 highlights the speedup and memory reduction achieved for a single transformer layer. We conducted experiments with a batch size of 4, varying the hidden sizes between 2048 and 4096, and sequence lengths of 2048 and 4096.

COAT demonstrates better speedup compared with TE, and significantly better memory reduction ability compared with BF16 and TE. Our approach achieves up to $1.57\times$ speedup over BF16 and achieves a consistent $1.65\times$ memory reduction compared to BF16, which is very close to the theoretically $1.69\times$ reported in Table 2.
The speedup ratio becomes larger with larger hidden sizes and longer sequence lengths.

**Table 6: Memory Saving and Speedup for a single Transformer Layer. Memory refers to Activation Memory. Our method achieves better speedup than TransformerEngine and significantly reduces the activation memory footprint by $1.65\times$.**
| Hidden Size = 2048, Batch Size = 4 |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Sequence Length = 2048 | Sequence Length = 4096 |  |  |  |  |  |  |  |  |  |  |
|  | Forward | Backward | Total | Ratio | Memory | Ratio | Forward | Backward | Total | Ratio | Memory | Ratio |
| BF16 | 3.36 | 8.47 | 11.83 | $1.00\times$ | 1457 MB | $1.00\times$ | 6.88 | 17.24 | 24.12 | $1.00\times$ | 2914 MB | $1.00\times$ |
| TE | 2.96 | 5.32 | 8.28 | $1.42\times$ | 1210 MB | $1.20\times$ | 5.94 | 11.29 | 17.23 | $1.39\times$ | 2420 MB | $1.20\times$ |
| COAT | 2.88 | 5.16 | 8.04 | $1.47\times$ | 883 MB | $1.65\times$ | 5.89 | 10.82 | 16.71 | $1.44\times$ | 1766 MB | $1.65\times$ |
| Hidden Size = 4096, Batch Size = 4 |  |  |  |  |  |  |  |  |  |  |  |  |
|  | Sequence Length = 2048 | Sequence Length = 4096 |  |  |  |  |  |  |  |  |  |  |
|  | Forward | Backward | Total | Ratio | Memory | Ratio | Forward | Backward | Total | Ratio | Memory | Ratio |
| BF16 | 7.77 | 18.78 | 26.55 | $1.00\times$ | 2914 MB | $1.00\times$ | 16.37 | 38.43 | 54.80 | $1.00\times$ | 5828 MB | $1.00\times$ |
| TE | 6.19 | 11.79 | 17.98 | $1.47\times$ | 2420 MB | $1.20\times$ | 12.66 | 24.58 | 37.24 | $1.47\times$ | 4840 MB | $1.20\times$ |
| COAT | 5.89 | 10.96 | 16.85 | $1.57\times$ | 1766 MB | $1.65\times$ | 12.16 | 23.44 | 35.6 | $1.53\times$ | 3533 MB | $1.65\times$ |

#### 6.2.2 Speedup and Memory Saving for end-to-end training

Table 7 presents a detailed comparison of end-to-end memory reduction and speedup results across different configurations for transformer models, specifically Llama-2-7B, Llama-2-13B, and Llama-30B, with variations in the number of GPUs used (1, 2, 4, and 8). It highlights COAT’s effectiveness in reducing end-to-end memory footprint and the speedup compared to standard BF16 and TransformerEngine (TE) setups under varying conditions of batch size and context length.

COAT allows full-parameter training of Llama-2-7B on *a single GPU*, where BF16 and TE both out of memory (OOM). Similarly, for the Llama-2-13B and Llama-30B models, our method enables 2-GPU training for Llama-2-13B and 8-GPU training for Llama-30B when Batch Size = 1.

In all multi-GPU training setting, COAT can double the micro-batch size and therefore lead to even higher speedup. For example, our method can achieve $2.25\times$ speedup when training Llama-2-13B on 4-GPUs since we can effectively increase the batch size to 2.

Overall, COAT significantly reduces end-to-end memory usage by up to $1.55\times$ and speeds up the end-to-end training by nearly $1.44\times$. This facilitates full-parameter training on fewer GPUs, which is particularly beneficial for larger language models.

**Table 7: End-to-end memory reduction and speedup results. BS refers to batch size. CL refers to context length. We report token/s per GPU for speed results. $\ddagger$ means CL=1024.**
| Llama-2-7B | *Context Length = 2048* | *Maximum Batch Size, Context Length = 2048* |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | Optimizer | Activations | Peak | Ratio | Max BS | Speed | Ratio |
| 1 $\text{GPU}_{\text{BS=1}}$ | BF16 | - | - | OOM | - | - | OOM | - |
| TE | - | - | OOM | - | - | OOM | - |  |
| COAT | 13.1 GB | 8.1 GB | 79.3 GB | ✓ | 1 | 5906 token/s | ✓ |  |
| 2 $\text{GPU}_{\text{BS=2}}$ | BF16 | - | - | OOM | - | 1 | 6130 token/s | 1.00$\times$ |
| TE | - | - | OOM | - | 1 | 6842 token/s | 1.11$\times$ |  |
| COAT | 6.5GB | 16.9 GB | 52.8 GB | ✓ | 4 | 11351 token/s | $1.85\times$ |  |
| 4 $\text{GPU}_{\text{BS=2}}$ | BF16 | 13.1 GB | 25.8 GB | 55.1 GB | $1.00\times$ | 2 | 7730 token/s | $1.00\times$ |
| TE | 13.1 GB | 21.9 GB | 51.1 GB | 1.08$\times$ | 2 | 9577 token/s | 1.24$\times$ |  |
| COAT | 3.2 GB | 16.9 GB | 35.6 GB | $1.54\times$ | 4 | 11257 token/s | $1.45\times$ |  |
| 8 $\text{GPU}_{\text{BS=2}}$ | BF16 | 6.5 GB | 25.8 GB | 41.2 GB | $1.00\times$ | 4 | 8238 token/s | $1.00\times$ |
| TE | 6.5 GB | 21.9 GB | 37.2 GB | 1.11$\times$ | 4 | 11704 token/s | 1.42$\times$ |  |
| COAT | 1.6 GB | 16.9 GB | 27.0 GB | $1.52\times$ | 8 | 11241 token/s | $1.36\times$ |  |
| Llama-2-13B | *Context Length = 2048* | *Maximum Batch Size, Context Length = 2048* |  |  |  |  |  |  |
|  |  | Optimizer | Activations | Peak | Ratio | Max BS | Speed | Ratio |
| 2 $\text{GPU}_{\text{BS=1}}^{\ddagger}$ | BF16 | - | - | OOM | - | - | OOM | - |
| TE | - | - | OOM | - | - | OOM | - |  |
| COAT | 12.6 GB | 10.1 GB | 73.2 GB | ✓ | 1 | 2137 token/s | ✓ |  |
| 4 $\text{GPU}_{\text{BS=1}}$ | BF16 | 25.1 GB | 20.1 GB | 76.1 GB | $1.00\times$ | 1 | 2345 token/s | $1.00\times$ |
| TE | 25.1 GB | 17.2 GB | 73.0 GB | 1.04$\times$ | 1 | 2851 token/s | 1.21$\times$ |  |
| COAT | 6.3 GB | 13.2 GB | 49.1 GB | $1.55\times$ | 2 | 5295 token/s | $2.25\times$ |  |
| 8 $\text{GPU}_{\text{BS=1}}$ | BF16 | 12.6 GB | 20.1 GB | 49.4 GB | $1.00\times$ | 2 | 3907 token/s | $1.00\times$ |
| TE | 12.6 GB | 17.2 GB | 46.5 GB | 1.06$\times$ | 2 | 5604 token/s | 1.43$\times$ |  |
| COAT | 3.1 GB | 13.2 GB | 32.5 GB | $1.52\times$ | 4 | 5650 token/s | $1.44\times$ |  |
| Llama-30B | *Context Length = 2048* | *Maximum Batch Size, Context Length = 2048* |  |  |  |  |  |  |
|  |  | Optimizer | Activations | Peak | Ratio | Max BS | Speed | Ratio |
| 8 $\text{GPU}_{\text{BS=1}}$ | BF16 | - | - | OOM | - | - | OOM | - |
| TE | - | - | OOM | - | - | OOM | - |  |
| COAT | 7.8 GB | 24.2 GB | 70.5 GB | ✓ | 1 | 1363 token/s | ✓ |  |

### 6.3 Ablation Studies

#### 6.3.1 Dynamic Range Expansion’s compatibility with other data formats

**Table 8: Dynamic Range Expansion is compatible with DE8 (8-bit dynamic quantization).**
|  | Second 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 |
| DE8 + Expand | 11.57 | 7.47 | 19.69 | 7.65 | 9.91 | 15.81 |

Dynamic Exponent Quantization (DE) was proposed in  (Dettmers et al., ) to quantize the optimizer states since its representation range has a range of 7 orders of magnitude and is very suitable for optimizer states quantization. Therefore, they proposed to quantize both first-order and second-order momentum with 8-bit DE.

We report the quantization error of $\frac{m}{\sqrt{v}}$ in Table 8, and find that applying our dynamic range expansion method to 8-bit DE can further reduce the quantization error by $1.41\times$. Specifically, the lowest quantization error is achieved when first-order momentum $m$ is quantized with 8-bit DE + Dynamic Range Expansion, and second-order momentum $v$ is quantized with E4M3/E5M2 + Dynamic Range Expansion. This proves the effectiveness of our method.

## 7 Conclusion

In this work, we present COAT, a memory-efficient FP8 training framework for foundation models by quantizing both optimizer states and activations into FP8 format.
We observe that the FP8 format’s representation range is not fully utilized when quantizing optimizer states, prompting us to propose Dynamic Range Expansion to align their dynamic range. We then identify the importance of quantizing non-linear layers, and propose mixed-granularity FP8 precision flow to quantize the activations accurately without introducing too much overhead.

Extensive experiments on LLM and VLM training and fine-tuning demonstrate that COAT can achieve nearly lossless performance. In end-to-end training, COAT achieves a $1.54\texttimes$ memory reduction and $1.43\texttimes$ speedup compared to BF16, and is comparable or even faster to TransformerEngine’s training speed. Our method also enables full-parameter training of billion-scale models on fewer GPUs and is able to double the batch size in realistic settings. These results highlight COAT’s capability to enable memory-efficient, large-scale model training without sacrificing accuracy, providing a highly effective solution for memory-constrained environments. Future work could further explore combining our proposed approach with other low-precision gradient compression methods to reduce communication overhead.

## Appendix A Details about optimizer states quantization

When performing optimizer.step(), We first dequantize the optimizer states into FP32, then update the optimizer states and weights in FP32 precision. In the end, we quantize the updated optimizer states back to FP8 and store it in GPU memory. This can be formulated as:

$$ $\displaystyle m_{t-1}$ $\displaystyle=DQ(m_{t-1}^{q},S_{m_{t-1}})$ $\displaystyle(\text{Dequantize to FP32})$ $\displaystyle v_{t-1}$ $\displaystyle=DQ(v_{t-1}^{q},S_{v_{t-1}})$ $\displaystyle(\text{Dequantize to FP32})$ $\displaystyle m_{t}$ $\displaystyle=\beta_{1}\ast m_{t-1}+(1-\beta_{1})\ast g_{t}$ $\displaystyle v_{t}$ $\displaystyle=\beta_{2}\ast v_{t-1}+(1-\beta_{2})\ast g_{t}^{2}$ $\displaystyle\hat{m}_{t}$ $\displaystyle=\frac{m_{t}}{1-\beta_{1}^{t}}$ $\displaystyle\hat{v}_{t}$ $\displaystyle=\frac{v_{t}}{1-\beta_{2}^{t}}$ $\displaystyle w_{t+1}$ $\displaystyle=w_{t}-\eta\left(\frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}}+\epsilon}+ \lambda w_{t}\right)$ $\displaystyle m_{t}^{q},S_{m_{t}}$ $\displaystyle=Q(m_{t})$ $\displaystyle(\text{Quantize to FP8})$ $\displaystyle v_{t}^{q},S_{v_{t}}$ $\displaystyle=Q(v_{t})$ $\displaystyle(\text{Quantize to FP8})$ $$

## Appendix B Qualitative Example - Vision Language Model captioning

Figure: Figure 8: Comparison of BF16 and COAT on VLM captioning. COAT can accurately summarize the figure and identify the key points in the figure.
Refer to caption: x8.png

## Appendix C Visualization of Expand Fuction

We further visualize Figure 3 by flattening it. This helps to understand the effectiveness of our method since floating point numbers can be more easily understood in a binary manner.

Figure: Figure 9: Axis is of base 2.
Refer to caption: x9.png

## Appendix D Detailed Explanation for Table 2

We mainly explain why the reduction ratio of linear layers is not exactly 50% here. 5.66U comes from: QKV Projection - 1U, Attention Projection - 1U, Up&Gate Projection - 1U, Down Projection - 2.66U. Among them, the attention projection’s input is exactly the output of FlashAttention. However, FlashAttention will save its input (QKV) and its output itself, so FlashAttention and Attention Projection actually share the same tensor when saving it. In Python, this tensor will only be saved for 1 time. So we should not further store an FP8 version of the Attention Projection’s input, since this will even increase the memory usage by 0.5U if we do not change the code of FlashAttention. Although we need to quantize the input to FP8 again in backward pass, we can store the scaling factor to reduce this additional overhead.

Therefore, after FP8 quantization, the memory usage of linear layers comes from: QKV Projection - 0.5U, Attention Projection - 1U, Up&Gate Projection - 0.5U, Down Projection - 1.33U. They sum up to 3.33U. Our FP8 precision flow method is also compatible with other quantized attention method Shah et al. (); Zhang et al. () by directly quantize the output of the RoPE function.

## Appendix E Detailed Explanation for Figure 1 (c)

We mainly discuss FP8-LM in this section. FP8-LM reduces the gradient communication precision from FP32 to FP8 and therefore greatly reduces the communication overhead. It also reduces the master weight’s precision from FP32 to FP16/BF16, and reduces the optimizer states precision from FP32 to BF16/FP8. For optimizer states and master weights, the memory footprint reduction result is consistent with our bar. However, for gradient, although FP8-LM can quantize the gradient into FP8, it still needs to preserve the main gradient in FP32 for gradient accumulation. Therefore it can not reduce the memory footprint of the weight gradient.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: See above.
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: None.
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
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.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
soundness: 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.
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.
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.

Meta-review: 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.

Decision: Accept (Poster)