Scaling FP8 training to trillion-token LLMs

## Abstract

Abstract We train, for the first time, large language models using FP8 precision on datasets up to 2 trillion tokens — a 20-fold increase over previous limits. Through these extended training runs, we uncover critical instabilities in FP8 training that were not observable in earlier works with shorter durations. We trace these instabilities to outlier amplification by the SwiGLU activation function. Interestingly, we show, both analytically and empirically, that this amplification happens only over prolonged training periods, and link it to a SwiGLU weight alignment process. To address this newly identified issue, we introduce Smooth-SwiGLU, a novel modification that ensures stable FP8 training without altering function behavior. We also demonstrate, for the first time, FP8 quantization of both Adam optimizer moments. Combining these innovations, we successfully train a 7B parameter model using FP8 precision on 256 Intel Gaudi2 accelerators, achieving on-par results with the BF16 baseline while delivering up to a ∼ 34 % similar-to absent percent 34 \sim 34\% ∼ 34 % throughput improvement.

## 1 Introduction

Large Language Models (LLMs) have revolutionized natural language processing, demonstrating remarkable capabilities across a wide range of tasks. However, the computational demands of training these models have become increasingly challenging, driving the need for more efficient training methods. Low precision formats, particularly FP8, have emerged as a promising solution to reduce memory usage and accelerate training. Recent work by Peng et al. () has demonstrated the potential of FP8 training for LLMs. However, these studies have been limited to datasets of up to 100 billion tokens, leaving open questions about the scalability and stability of FP8 in truly large-scale training scenarios.

In this paper, we present advancements in FP8 training for LLMs, successfully scaling to datasets of up to 2 trillion tokens - a 20-fold increase over previous limits. This leap in scale has revealed critical instabilities in FP8 training that were not observable in earlier, shorter-duration studies. Through rigorous analysis, we trace these instabilities to a previously unidentified phenomenon: the amplification of outliers by the SwiGLU activation function (Shazeer ()), which becomes pronounced only after extended training periods. The severity of this issue is illustrated in Fig. 2(a), where we present the training loss of Llama2 7B using FP8 precision. The graph clearly shows a dramatic divergence caused by outliers after processing 220B tokens - twice the dataset size explored in previous work (Peng et al. ()).

To address this newly discovered challenge, we introduce Smooth-SwiGLU, a novel modification to the standard SwiGLU activation that effectively mitigates outlier amplification without altering the function’s behavior. This innovation ensures stable FP8 training across extended durations, enabling the use of FP8 precision in large-scale LLM training without compromising model performance.

Furthermore, we push the boundaries of low-precision optimization by demonstrating, for the first time, the successful quantization of both Adam optimizer moments to FP8. This advancement reduces memory usage during training, further enhancing the efficiency of large-scale LLM development.

By combining these innovations - Smooth-SwiGLU and FP8 quantization of optimizer moments - we successfully train a 7B parameter model using FP8 precision on 256 Intel Gaudi2 accelerators. Our approach achieves results on par with the BF16 baseline while delivering up to  34% throughput improvement, marking a significant leap in training efficiency for large-scale language models.

Our paper makes several key contributions:

- •
We demonstrate the first successful FP8 training of LLMs on datasets up to 2 trillion tokens, far surpassing previous limits and revealing critical instabilities in extended training regimes.
- •
We identify the root cause of these instabilities: outlier amplification by the SwiGLU activation function over prolonged training periods.
- •
We link, analytically and empirically, the outlier amplification to a weight alignment happening during training with $\ell_{2}$ regularization, in the SwiGLUs with sufficiently large inputs.
- •
We introduce Smooth-SwiGLU, a novel activation function that ensures stable FP8 training without altering model behavior, enabling efficient large-scale LLM training.
- •
We present the first implementation of FP8 quantization for both Adam optimizer moments, further optimizing memory usage in LLM training.
- •
We achieve on-par results with BF16 baselines on downstream tasks while providing throughput improvements on Intel Gaudi2 accelerators, demonstrating the practical viability of our approach for state-of-the-art LLM training.

## 2 Background and Challenges of FP8 Training in LLMs

The computational demands of Large Language Models (LLMs) have driven a shift from traditional FP32 to reduced-precision formats. While FP16 and BF16 have become standard in many training tasks (Micikevicius et al., ; Scao et al., ; Smith et al., ), FP8 represents the next step in this progression towards lower precision. Micikevicius et al. () standardized two FP8 formats for deep learning: E4M3 (4 exponent bits, 3 mantissa bits) optimized for weights and activations, and E5M2 (5 exponent bits, 2 mantissa bits) suitable for gradients.

FP8 shows promise for large-scale training, especially with support from modern hardware like NVIDIA’s H100 and Intel’s Gaudi2. However, its limited range necessitates careful scaling techniques to maintain numerical stability and model performance.

The primary challenge in FP8 training for LLMs stems from its limited dynamic range. To address this, researchers have developed various scaling techniques. Global loss scaling (Micikevicius et al., ) multiplies the entire loss by a constant factor to prevent gradient underflow during backpropagation. Per-tensor scaling (Sun et al., ) takes a more granular approach, scaling each tensor individually based on its specific range of values. These techniques allow for better utilization of the FP8 format’s limited range. However, (Lee et al., ) demonstrates that, even with these techniques, FP8 training can lead to significant instabilities, highlighting the ongoing challenges in reduced-precision LLM training.

Implementation of these scaling techniques typically follows one of two approaches: just-in-time scaling or delayed scaling. Just-in-time scaling dynamically adjusts scaling factors based on the current data distribution. However, it often proves impractical in FP8 training due to the need for multiple data passes, which can negate the performance benefits of using FP8. Delayed scaling, on the other hand, selects scaling factors based on data distributions from preceding iterations. While more practical, it assumes consistent statistical properties across iterations, making it vulnerable to outliers that can disrupt this consistency and potentially destabilize the training process.

Recent work by Peng et al. () demonstrated the first empirical results of training LLMs using FP8 format up to 100 billion tokens, validating the potential of FP8 for large-scale training. Our research extends this work, successfully scaling FP8 training to datasets of up to 2 trillion tokens. We introduce novel techniques to address the challenges of FP8’s limited dynamic range and the instabilities that emerge in extended training scenarios. Our approach not only overcomes these limitations but also achieves substantial improvements in memory usage and training speed, demonstrating the viability of FP8 training for truly large-scale LLM development.

## 3 Outlier Amplification in Large-Scale FP8 Training

The presence of outliers has been observed in numerous studies (Yang et al., ), particularly in the activations during inference. These outliers can significantly impact the stability and performance of the model, as they introduce extreme values that are difficult to manage within the limited dynamic range of reduced-precision formats like FP8. Our work reveals that these outliers become particularly prominent in the later stages of training large language models (LLMs) with large-scale datasets.

Figure: (a)
Refer to caption: x1.png

As shown in Fig. 1, these outliers emerge only after processing approximately 200 billion tokens during training. This phenomenon poses significant challenges to maintaining numerical stability and model performance, especially when using delayed scaling methods that assume consistency across iterations. The sudden appearance of these outliers disrupts the statistical assumptions underlying FP8 training techniques, potentially leading to instability or divergence in the training process.

The emergence of these outliers in the later stages of training is particularly problematic for FP8 formats due to their limited dynamic range. Unlike higher precision formats such as FP32 or even BF16, FP8 has a much narrower range of representable values. When outliers exceed this range, they can cause overflow or underflow, leading to a loss of critical information and potentially destabilizing the entire training process.

Moreover, the sporadic nature of these outliers, as evident in Fig. 1(b), makes them challenging to predict and manage. Traditional scaling techniques, which rely on consistent statistical properties across iterations, struggle to adapt to these sudden, extreme values. This unpredictability further complicates the task of maintaining numerical stability in FP8 training, especially as the scale of the dataset and the duration of training increase.

## 4 SwiGLU and Outlier Amplification

While the previous section highlighted the general problem of outlier emergence in large-scale FP8 training, our investigation reveals that the SwiGLU (Swish Gated Linear Unit) activation function plays a crucial role in amplifying these outliers. This section explores the structure of SwiGLU and demonstrates how its unique properties contribute to the generation and amplification of outliers.

### 4.1 SwiGLU Structure

The Transformer architecture (Vaswani et al., ), which forms the foundation of modern LLMs, has undergone several modifications to enhance performance and efficiency. One notable example is the inclusion of the SwiGLU (Swish Gated Linear Unit) (Shazeer, ) activation function in models like LLaMA (Touvron et al., ) and PaLM (Chowdhery et al., ).

Let $\mathbf{x}\in\mathbb{R}^{d}$ be the input vector from the previous layer. The SwiGLU neuron is defined as follows for two weight vectors $\mathbf{w}_{1},\mathbf{w}_{2}\in\mathbb{R}^{d}$:

$$ $\text{SwiGLU}_{\mathbf{w}_{1},\mathbf{w}_{2}}(\mathbf{x})\triangleq(\mathbf{x} ^{\top}\mathbf{w}_{1})\mathrm{Swish}(\mathbf{x}^{\top}\mathbf{w}_{2}) \triangleq(\mathbf{x}^{\top}\mathbf{w}_{1})(\mathbf{x}^{\top}\mathbf{w}_{2}) \sigma(\mathbf{x}^{\top}\mathbf{w}_{2})$ $$

where $\sigma(z)\triangleq 1/(1+e^{-z})$ is the sigmoid function.

While other standard neuron types, such ReLU, GeLU, and Swish at most linear at large input magnitudes (i.e., $\lim_{u\rightarrow\pm\infty}|f(u)/u|\leq 1$), the SwiGLU neuron is quadratic and can reach much larger values (and cause very strong outliers) if $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ are sufficiently aligned (e.g., if $\mathbf{w}_{1}=\mathbf{w}_{2}$ and $\mathbf{w}^{\top}_{1}\mathbf{x}=1$ then $\lim_{c\rightarrow\infty}\text{SwiGLU}_{\mathbf{w}_{1},\mathbf{w}_{2}}(c
\mathbf{x})/c^{2}=1$). As we show next, precisely such alignment happens during training for neurons with sufficiently large inputs.

### 4.2 Theoretical analysis of weight correlation in SwiGLU

Next, we analyze the behavior of the SwiGLU neuron during training and show its weight vectors tend to align perfectly if the magnitude of its input increases above some threshold. This causes the SwiGLU output magnitude to increase significantly during training, potentially resulting in outliers.

To show this, We assume the SwiGLU neuron is embedded in a neural network with $k$ parameters. The rest of the parameters in the network are denoted by $\theta\in\mathbb{R}^{k-2d}$. We train the neural network with some $\ell_{2}$ regularization,

$$ $\min_{\mathbf{w}_{1},\mathbf{w}_{2},\theta}\sum_{n=1}^{N}\ell_{n}\left(\mathrm {SwiGLU}_{\mathbf{w}_{1},\mathbf{w}_{2}}\left(\mathbf{x}_{n}\left(\theta\right )\right),\theta\right)+\frac{\mu}{2}\left(\left\|\mathbf{w}_{1}\right\|^{2}+ \left\|\mathbf{w}_{2}\right\|^{2}\right)$ (1) $$

where $\mu>0$ is regularization strength, $N$ is the number of training samples, and $\ell_{n}\left(u_{n},\theta\right)$
is the per-sample loss as a function of the SwiGLU output
and the rest of the neural network parameters. We find that

Theorem 1. Suppose we converge to a stationary point $(\mathbf{w}_{1},\mathbf{w}_{2},\theta)$ of the loss function, and for all samples $n$, $\sigma^{\prime}(\mathbf{x}_{n}^{\top}\left(\theta\right)\mathbf{w}_{2})\rightarrow
0$. Then, at this stationary point, $\mathbf{w}_{1}\rightarrow\mathbf{w}_{2}$ or $\mathbf{w}_{1}\rightarrow-\mathbf{w}_{2}$.

Proof. At a stationary point $(\mathbf{w}_{1},\mathbf{w}_{2},\theta)$ we have $\forall i\in\{1,2\}$:

$$ $\sum_{n=1}^{N}\nabla_{\mathbf{w}_{i}}\ell_{n}(\text{SwiGLU}_{\mathbf{w}_{1}, \mathbf{w}_{2}}(\mathbf{x}_{n}(\theta)),\theta)+\mu\mathbf{w}_{i}=0$ $$

Using the chain rule we have:

$$ $0=\sum_{n}\mathbf{x}_{n}\mathbf{x}_{n}^{\top}\mathbf{w}_{2}\sigma(\mathbf{w}_{ 2}^{\top}\mathbf{x}_{n})\delta_{n}+\mu\mathbf{w}_{1}$ $$

$$ $0=\sum_{n}\mathbf{x}_{n}\mathbf{x}_{n}^{\top}\mathbf{w}_{1}\left(\sigma( \mathbf{w}_{2}^{\top}\mathbf{x}_{n})+(\mathbf{x}_{n}^{\top}\mathbf{w}_{2}) \sigma^{\prime}(\mathbf{x}_{n}^{\top}\mathbf{w}_{2})\right)\delta_{n}+\mu \mathbf{w}_{2}$ $$

Note where we defined $\delta_{n}(\mathbf{w}_{1},\mathbf{w}_{2},\theta)\triangleq\left.\frac{\partial
\ell_{n}(u_{n},\theta)}{\partial u_{n}}\right|_{u_{n}=\text{SwiGLU}_{\mathbf{w
}_{1},\mathbf{w}_{2}}(\mathbf{x}_{n}(\theta))}$ and with a slight abuse of notation, we suppressed the dependence of $(\mathbf{w}_{1},\mathbf{w}_{2},\theta)$ and $\mathbf{x}_{n}(\theta)$ on $\theta$. Given the assumption

$$ $\forall n:\sigma^{\prime}(\mathbf{x}_{n}^{\top}\mathbf{w}_{2})\rightarrow 0$ $$

at this limit we obtain

$$ $0=\sum_{n}\mathbf{x}_{n}\mathbf{x}_{n}^{\top}\mathbf{w}_{2}\sigma(\mathbf{w}_{ 2}^{\top}\mathbf{x}_{n})\delta_{n}+\mu\mathbf{w}_{1}\quad;\quad 0=\sum_{n} \mathbf{x}_{n}\mathbf{x}_{n}^{\top}\mathbf{w}_{1}\sigma(\mathbf{w}_{2}^{\top} \mathbf{x}_{n})\delta_{n}+\mu\mathbf{w}_{2}$ $$

Now, defining $\lambda_{n}=-\mu^{-1}\delta_{n}\sigma(\mathbf{w}_{2}^{\top}\mathbf{x}_{n})$, the above equations become:

$$ $\mathbf{w}_{1}=\sum_{n}\lambda_{n}\mathbf{x}_{n}\mathbf{x}_{n}^{\top}\mathbf{w }_{2}=A\mathbf{w}_{2}\quad;\quad\mathbf{w}_{2}=\sum_{n}\lambda_{n}\mathbf{x}_{ n}\mathbf{x}_{n}^{\top}\mathbf{w}_{1}=A\mathbf{w}_{1}$ $$

where $A=\sum_{n}\lambda_{n}\mathbf{x}_{n}\mathbf{x}_{n}^{\top}$ is a symmetric matrix. This implies

$$ $\mathbf{w}_{1}=A\mathbf{w}_{2}=A^{2}\mathbf{w}_{1}\quad;\quad\mathbf{w}_{2}=A \mathbf{w}_{1}=A^{2}\mathbf{w}_{2}$ (2) $$

Since $A$ is symmetric this implies that both $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ are eigenvectors of $A$ with the same eigenvalue: $1$ or $-1$. Plugging this into equation 2 we obtain $\mathbf{w}_{1}=\mathbf{w}_{2}$ or $\mathbf{w}_{1}=-\mathbf{w}_{2}$. $\blacksquare$

Thus, with regularization and sufficient training, the weight vectors $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ must converge to either identical or opposite directions, i.e., $\mathbf{w}_{1}\approx\mathbf{w}_{2}$ or $\mathbf{w}_{1}\approx-\mathbf{w}_{2}$ — if $\sigma^{\prime}(\mathbf{x}_{n}^{\top}\left(\theta\right)\mathbf{w}_{2})\approx
0$ holds for all samples $n$. Since $\sigma^{\prime}(z)$ decays exponentially fast as $|z|$ increases, this simply means that neuron inputs are not too small. This can happen in the case in which $\left\|\mathbf{w}_{2}\right\|$
is sufficiently large, and $\forall n\in S:\,\,\left|\mathbf{w}_{2}^{\top}\mathbf{x}_{n}\left(\theta\right
)\right|>0$. The latter condition is the generic case when we fit a neural network to a large dataset of size $N$, where $N\gg k$ (i.e. we are in an under-parameterized regime). . Interestingly, this alignment phenomenon happens due to $\ell_{2}$ regularization, even if it is very weak. In fact, weak regularization will lead to larger weight norms and will strengthen this effect.

### 4.3 Observing weight correlation growth and training instability

In our experiments, we observed a clear relationship between the increasing correlation of the weight matrices $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ and the eventual divergence in FP8 training loss.

Figure: Figure 2: (a): Training loss of LlaMA2-7b with BF16 and FP8 precision, where a significant loss divergence is seen for FP8 after step $\sim$ 200B tokens. (b): Dynamics of the $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ norms, and their correlation during training, for a specific channel that generates outliers. A drastic increase in correlation and norm is observed at the same point where we start to see loss degradation in (a). (c): Scatter plot of an outlier channel elements in $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$, at an early training stage (8B tokens) and late training stage (330B tokens), demonstrating minimal correlation at start of the training and high correlation in the later stage. (d): Histogram of an outlier channel of $\mathbf{w}_{1}$ at an early training stage (8B tokens) and late training stage (330B tokens).
Refer to caption: x3.png

The process of weight alignment and its impact on FP8 training develops precisely as our theory predicts. As training progresses, input magnitudes in certain channels grow, eventually exceeding a critical threshold. When this occurs, weight vectors $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ in these channels begin to align rapidly. Figure 2b illustrates this evolution: initially, the correlation between $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ is low, but it increases drastically between 125B and 210B tokens. Simultaneously, we observe a significant increase in the norm of these weights. This combination of high correlation and increased norm creates ideal conditions for generating extreme activation values, or “spikes.”

These activation spikes, in turn, lead to the divergence of FP8 training, as shown in Fig. 2a. Importantly, while we primarily observe strong positive correlations in this example, our theory also predicts the possibility of strong negative correlations. We also observe these, as can be seen in Fig. 7 in the Appendix.

This phenomenon highlights the unique challenges posed by SwiGLU in FP8 training of large language models. The gradual alignment of weights, combined with norm growth, creates a scenario where outliers become increasingly likely and severe as training progresses. Consequently, instabilities may not be apparent in shorter training runs but emerge as critical issues in large-scale, long-duration training scenarios. This explains why such problems have not been observed in previous, smaller-scale studies of FP8 training.

### 4.4 Smooth-SwiGlu

As demonstrated earlier, the SwiGLU activation function can lead to outliers in the input of the last linear layer of the MLP component. These outliers pose a significant challenge when using FP8 precision with delayed scaling, which relies on the assumption that statistical properties remain consistent across layers. The sudden spike in value caused by SwiGLU disrupts this continuity, leading to instability in the training process. In Fig. 3 we demonstrate that disabling the quantization of the last linear layer in the MLP component (output of SwiGLU) allows Llama2 FP8 to successfully converge with large datasets, addressing the previously observed divergence issues.

Figure: Figure 3: Training loss of Llama2 FP8 with and quantization of SwiGLU output. As can be seen the cause of the divergence of standard FP8 is the amplification of the SwiGLU (input to $\mathbf{w}_{3}$).
Refer to caption: x4.png

While disabling quantization of the SwiGLU output effectively prevents divergence, it reduces the potential acceleration benefits of FP8. To maintain full FP8 acceleration while addressing the outlier problem, we propose a novel modification called Smooth-SwiGLU. Figure 4 illustrates the key idea behind Smooth-SwiGLU: applying a scaling factor to the linear branch of the SwiGLU function and rescaling it back after the last linear layer. This approach prevents outliers in the quantization of the input to the last linear layer while preserving the overall function of the SwiGLU activation, enabling us to fully leverage FP8 precision throughout the network.

Figure: (a)
Refer to caption: x5.png

Mathematically, we express the quantized Smooth-SwiGLU function for each channel $i$ as:

$$ $\text{Smooth-SwiGLU}_{\hat{\mathrm{w}}_{1,i},\hat{\mathrm{w}}_{2,i}}(\mathbf{x })=s^{-1}_{i}\cdot Q(s_{i}\cdot(\hat{\mathbf{w}}_{1,i}^{\top}Q(\mathbf{x})) \mathrm{Swish}(\hat{\mathbf{w}}_{2,i}^{\top}Q(\mathbf{x}))))$ (3) $$

where $s_{i}$ is the per-channel scaling factor, $\hat{\mathbf{w}}=Q(\mathbf{w})$ denotes the quantized version of a weight vector $\mathbf{w}$, and $Q$ is a quantization function (in a slight abuse of notation, we suppress the dependence of $Q$ on the tensor).

To minimize computational overhead, we compute the scaling factors $s_{i}$ using an efficient parallel approach:

- 1.
Split the tensor into chunks, where each chunk corresponds to a channel.
- 2.
For each chunk (channel), compute its maximum value in parallel.
- 3.
Use these per-channel maximum values to determine the individual scaling factors $s_{i}$ for each channel.

This method allows for efficient per-channel scaling, as each channel’s scaling factor is computed independently and in parallel. The computational cost of this approach is moderate compared to the matrix multiplications in the linear layers, especially given the parallelization, even with our non optimized implementation. During inference, the scaling factors can be merged into the weights of the first and third linear layers in the MLP layer that includes the SwiGLU layer followed by a linear layer (see Figure 4), i.e.

$$ $\displaystyle\sum_{i}\hat{\mathbf{w}}_{3,i}\text{Smooth-SwiGLU}_{\hat{\mathrm{ w}}_{1,i},\hat{\mathrm{w}}_{2,i}}(\mathbf{x})$ $\displaystyle=$ $\displaystyle\sum_{i}s^{-1}_{i}\cdot\hat{\mathbf{w}}_{3,i}Q(s_{i}\cdot(\hat{ \mathbf{w}}_{1,i}^{\top}Q(\mathbf{x}))\mathrm{Swish}(\hat{\mathbf{w}}_{2,i}^{ \top}Q(\mathbf{x}))))$ $$

So we can absorb the scalar by re-defining $\tilde{{\mathbf{w}}}_{1,i}\triangleq Q(s_{i}\cdot{\mathbf{w}}_{1,i})$ and $\tilde{{\mathbf{w}}}_{3,i}\triangleq Q(s^{-1}_{i}\cdot{\mathbf{w}}_{3,i})$. Thus, this procedure results in zero additional cost at inference.

## 5 FP8 Optimizer

The Adam optimizer and its variants are widely used in deep learning due to their effectiveness in handling various training challenges. A key characteristic of the Adam optimizer is its storage of two moments, traditionally in high precision (FP32). This significantly increases memory usage, particularly for large-scale models. While previous research Peng et al. () has shown the feasibility of reducing the first moment to FP8 precision, they retained the second moment in FP16. Our work pushes the boundaries further by successfully quantizing both moments to FP8, significantly improving the optimizer efficiency for large language models.

### Challenges

The Adam optimizer uses two moments to adapt learning rates for each parameter:

- 1.
The first moment is an estimate of the mean of the gradients.
- 2.
The second moment is an estimate of the uncentered variance of the gradients.

A critical aspect of the Adam optimizer is the use of the inverse square root of the second moment in the parameter update step. This operation has important implications for precision requirements.

Due to this inverse square root operation, the smallest values in the second moment become the most significant in determining parameter updates. This characteristic creates a unique challenge when considering precision reduction for the second moment.

### Methodology

We conducted extensive experiments to determine the optimal FP8 formats for both moments. Our investigation revealed that different precision requirements exist for each moment:

- 1.
First Moment: The E4M3 format (4 exponent bits, 3 mantissa bits) provides sufficient precision. This format offers a good balance between range and accuracy for representing the mean of the gradients.
- 2.
Second Moment: The E5M2 format (5 exponent bits, 2 mantissa bits) is necessary. This format provides a higher dynamic range, crucial for preserving information about the smallest values in the second moment. The additional exponent bit ensures that we can accurately represent both very small and moderately large values, which is critical given the inverse square root operation applied to this moment.

In our experiments, presented in Fig. 5 we show that while the first moment is able to converge with E4M3, the second moment requires a wider dynamic range and is able to converge only with E5M2 format. In Table 1 we compare the proposed quantization scheme for the optimizer moments with the one presented in Peng et al. (). Our scheme shows, for the first time, the ability to quantize both moments with standard FP8 formats.

Figure: Figure 5: All combinations for quantization the Adam moments with standard FP8 formats in Llama2 100m. The only combination that is able to converge to baseline is first moment E4M3 format and second moment E5M2 format.
Refer to caption: x7.png

**Table 1: Comparison of the different datatypes for the two Adam optimizer moments.**
| Model | Mom 1 | Mom 2 |
| --- | --- | --- |
| BF16 (Baseline) | FP32 | FP32 |
| FP8 (Peng et al., ) | FP8 | FP16 |
| FP8 (Ours) | FP8 | FP8 |

## 6 Experiments

We conducted extensive experiments to evaluate the effectiveness of our proposed FP8 training method for Large Language Models (LLMs) across various scales.

### 6.1 Experimental Setup

#### Model and Dataset.

We used the Llama2 model (Touvron et al., ) as our baseline. This model is a decoder-only Transformer (Brown et al., ) with pre-normalization RMSNorm (Zhang & Sennrich, ), SwiGLU activation function (Shazeer, ), and rotary positional embeddings (Su et al., ). We trained the models on the open-source Red Pajama dataset (Computer, ) for 2 trillion tokens, maintaining hyperparameters consistent with Touvron et al. ().

#### Hardware.

All training was conducted on 256 Intel Gaudi2 devices.

### 6.2 Results

#### Training Stability.

In Fig. 6 we show the training loss of Llama2 over 300B tokens with the proposed scheme, which include the use of Smooth SwiGLU (Section 4.4) + FP8 quantization of both Adam moments (Section 5). Notice we are able to overcome the divergence point of standard FP8 training . Additional training steps will be added in the next version of the manuscript. The FP8 model was trained using the standard format (Micikevicius et al., ) which include E4M3 for the forward phase and E5M2 for the backward phase.

Figure: Figure 6: Training loss of Llama2 7B using the proposed Smooth SwiGLU as the activation function (Section 4.4) and with FP8 quantization of both moments of Adam optimizer (Section 5). As can be seen, the proposed scheme is able to converge similar to the BF16 baseline while standard FP8 with SwiGLU and FP32 Adam moments diverge after 200B tokens.
Refer to caption: x8.png

#### Zero-shot Performance.

Table 2 compares the zero-shot performance (accuracy and perplexity) on downstream tasks between the BF16 baseline and our FP8 model with Smooth-SwiGLU and FP8 optimizer. The results demonstrate that our FP8 approach achieves on-par performance with the BF16 baseline across all tested metrics.

**Table 2: Zero shot accuracy and perplexity comparison between the BF16 baseline and the proposed FP8. Notice both models achieve on-par results across all tests.**
| Precision | Accuracy $\uparrow$ | Perplexity $\downarrow$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Lambada | HellaSwag | Winogrande | Arc-C | Wikitext | Lambada |  |
| BF16 | 61.98 | 68.3 | 64.25 | 37.37 | 5.59 | 5.75 |
| FP8 | 61.73 | 68.03 | 64.4 | 37.2 | 5.56 | 5.89 |

#### Performance gains.

Table 3 presents the performance of different configurations on Intel Gaudi2 hardware. While full FP8 quantization achieves the highest acceleration ($\sim$37%), it leads to training divergence (as shown in Figure 2a). Disabling quantization for the $\mathbf{w}_{3}$ layer enables convergence (Figure 3) with a $\sim$27% speedup. Our proposed Smooth-SwiGLU scheme not only converges with results on par with the BF16 baseline (Figure 6) but also delivers a substantial $\sim$34% acceleration.

**Table 3: Performance acceleration with different configurations in our non optimized implementation in Llama2 7B model. The measurement were done on 8 Intel Gaudi2 devices.**
| Configuration | Micro BS | Status | Throughput |
| --- | --- | --- | --- |
|  |  |  | (Samples/sec) |
| BF16 | 1 |  | 12.65 |
| FP8 + SwiGLU output in BF16 | 1 | Converge | 16.07 (+ 27.04 %) |
| FP8 + Smooth SwiGLU | 1 | Converge | 16.89 (+33.52 %) |
| FP8 | 1 | Diverge | 17.34 (+37.08 %) |

#### Memory reduction.

In Table 4 we present the memory reduction achieved by changing the optimizer moments from standard FP32 to FP8. As can be seen, we can reduce the memory consumption by $\sim 30\%$,

**Table 4: Memory reduction when applying the proposed FP8 optimizer (Section 5). The measurements were done on 8 Intel Gaudi2 devices, using Deepspeed Zero-1.**
| Configuration | Status | Memory | FP8 Optimizer |
| --- | --- | --- | --- |
|  |  | (GB/HPU) |  |
| BF16 |  | 63.25 | ✗ |
| FP8 + SwiGLU output in BF16 | Converge | 63.26 | ✗ |
| FP8 + Smooth SwiGLU | Converge | 63.26 | ✗ |
| FP8 | Diverge | 63.24 | ✗ |
| FP8 + SwiGLU output in BF16 | Converge | 44.08 | ✓ |
| FP8 + Smooth SwiGLU | Converge | 44.08 | ✓ |
| FP8 | Diverge | 44.09 | ✓ |

## 7 Conclusions

In this paper, we successfully demonstrated FP8 training on datasets up to 2 trillion tokens, significantly exceeding the previous limit of 100 billion tokens Peng et al. (), with on-par results with the BF16 baseline. Importantly, we discovered that earlier FP8 training attempts were not long enough to reveal critical instabilities caused by outliers. Through both analytical methods and simulations, we showed that these outliers emerge over time, particularly in extended training runs. Our investigation revealed that the SwiGLU activation function amplifies these outliers, destabilizing FP8 training in large-scale scenarios.

To address this issue, we applied per-channel quantization to the SwiGLU activation function, a technique we refer to as Smooth-SwiGLU. Although identical to SwiGLU in function, this method effectively reduces outlier amplification, ensuring stable FP8 training with a moderate effect on model performance during training, and without any effect on the inference. Additionally, we introduced the first implementation of FP8 quantization for both Adam optimizer moments, further optimizing memory usage.

Our proposed method, combining Smooth-SwiGLU and FP8 optimizer moments, achieved comparable performance to BF16 baselines on downstream tasks while providing significant throughput improvementsd. This approach successfully overcome the divergence challenges typically encountered in standard FP8 training on large datasets.

## Appendix A Appendix

### A.1 Training instability - additional data

Figure: Figure 7: (a): Scatter plot of outlier channel elements in $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$, at an early training stage (8B tokens) and late training stage (330B tokens), demonstrating minimal correlation at start of the training and high negative correlation in the later stage. (b): Histogram of an outlier channel with negative correlation of $\mathbf{w}_{1}$ at an early training stage (8B tokens) and late training stage (330B tokens).
Refer to caption: x9.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: I would like to ask the authors for several clarifications:

1. If I understand correctly, smooth-SwiGLU does not directly address the weight alignment problem. Rather, it is an alternative implementation specifically designed for FP8 training, which circumvents the overflow issue by preventing outliers in the inputs to the last linear layer. Is this accurate?

2. Figure 2 (b) on the dynamics of $\mathbf w_1$ and $\mathbf w_2$ norm correlation is insightful and interesting. Is the dynamics based on llama2's trianing hypermaraters? Do the authors have additional empirical results on how these dynamics change with different training hyperparameters?
soundness: 3
strengths: 1. The paper demonstrates opportunities and challenges in scaling FP8 training to trillion-token scale using LLaMA2 architecture. In particular, the authors identify that with SwiGLU (which is an important to contemporary LLMs), weight alignment issues can induce large-magnitude outliers and result in training stability issue that poses significant challenge for FP8's limited dynamic range.

2. Extensive experiment results are provided, showing the proopsed Smooth-SwiGLU training with FP8 optimizer (moment 1: E4M3,moment 2: E4M3) is able to converge to baseline FP16 training. 

3. Experiments show around 30% improvement in training throughput, which provides significant energy saving for large-scale training.

4. Paper writing and presentation are clear with well identified bottleneck and detailed solutions.
summary: The paper presents new findings and addresses key challenges in large-scale FP8 training for LLMs on modern hardware that supports FP8 operations. The authors identified that one major challenge in scaling FP8 training is the occurrence of sporadic, high-magnitude activations. They pinpoint SwiGLU as the main source of these extreme outliers, providing the novel insight that a weight alignment process during training leads to substantial SwiGLU activations (as discussed in Section 4.1). To address this issue and stabilize large-scale FP8 training, the authors propose a smooth-SwiGLU activation function. This approach prevents outliers in the quantization of inputs to the last linear layer with an efficient per-channel scaling for better parallelism.
weaknesses: 1. The paper has an extensive discussion about how feedforward SwiGLU affect FP8 training stability. However, there is no mention about how other model components like RMSNorm, MHA/GQA affect training stability. Could the authors also discuss whether each of the other model components affect FP8 training stability and provide quantitative results? 

2. It seems the SwiGLU weight alignment issue explains the occurrence of huge-magnitude outliers. However, can the authors comment more on the sporadic nature of these outliers? Is SwiGLU held accountable for that as well?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: How do the results generalize to other activation functions beyond SwiGLU?
soundness: 3
strengths: Very interesting findings!

The introduction of Smooth-SwiGLU and the quantization of Adam optimizer moments are significant advancements in the FP8 training methodology.

The paper provides thorough experimental results demonstrating the effectiveness of the proposed techniques, achieving comparable performance to BF16 with improved training throughput.

Successfully training models on datasets up to 2 trillion tokens sets a new benchmark for FP8 training, addressing scalability issues in LLMs.
summary: The paper presents a novel approach to training large language models (LLMs) using FP8 precision on datasets of up to 2 trillion tokens, revealing critical instabilities associated with FP8 training that were not observable in prior studies. It identifies that the SwiGLU activation function amplifies outliers over prolonged training, leading to instability. To mitigate this, the authors propose Smooth-SwiGLU, a modified activation function that maintains stability without altering performance. The paper also introduces FP8 quantization for both Adam optimizer moments, significantly improving memory usage and training throughput.
weaknesses: While the findings are significant, their applicability to other model architectures beyond those tested (like LLaMA2) could be explored further.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: How does this paper compare with nvidia transformer-engine's fp8 training with automatic mixed precision?

In table 3, why micro batch is 1? The reason for doing quantization is to support larger batch training (higher throughput), if limiting to micro-batch as 1, it is almost impossible to get good throughput/token per sec numbers.

In table 4, deepspeed zero-1 itself would reducing some GPU memory footprint compared with pure pytorch DDP. Why use zero1 not 2 or 3?

In figure 6, why bf16 loss curve has more spikes compared with FP8 + Smooth-SwiGLU + FP8 Optimizer? to me, I think bf16 should have better and more smooth loss curve compared with any FP8 methods?
soundness: 3
strengths: Training LLM on trillion-token big dataset makes the paper result solid.

Adopting a pure FP8 adam optimizer (both two momentum are fp8) and make training converge is good contribution. 

Identify swiglu cause model training instability and optimized as a smooth-swiglu is another contribution.
summary: This paper proposed a practical solution for quantized (FP8) training over trillion tokens.

The paper conduct LLM training over a real big 2 trillion dataset and found swiglu is the main cause of instability when training using FP8, thus proposed a optimized and smoothed method called smooth-swiglu.

Smooth-swiglu in together with FP8, the authors achieve good model convergence on 2 trillion token dataset for llm training
weaknesses: 1 ) Paper novelty is limited. The only thing new is adding a scaling factor on top of swiglu may not be a big novelty. Previous work already study similar smooth method on activation functions. e.g., Swish, SMU as smooth function for relu.

[1] swish: https://arxiv.org/pdf/1710.05941

[2] SMU: https://arxiv.org/pdf/2111.04682

2 ) FP8 optimizer contribution is just extending previous 1 momentum using FP8 to both momentum using FP8 may lack a bit novelty here. For example, Figure 5 shows empirical study for both FP8 momentum on llama2 100m, it shows second momentum using E5M2 format achieve similar model loss curve as bf16. But how does it generalize to bigger models (more practical models like llama 8b 70b)? Is second momentum always need larger dynamic range rather than higher precision?

3 ) The only baseline is from this paper[3], how does this approach compared with nvidia transformer engine fp8. There is no comparison in either design or evaluation results. Please include a comparison with NVIDIA's Transformer Engine FP8 implementation, both in terms of methodology and empirical results. This comparison would provide valuable context for the paper's contributions.

[3]FP8-LM: Training FP8 Large Language Models

4 ) Billion level token training is almost sufficient for most LLM downstream tasks (e.g. Apple [4]). This trillion-token improvement may not have wide application scenarios. Please discuss more on potential use cases where trillion-token training might be beneficial, and how this paper's method scales compared to existing approaches at different dataset sizes.

[4] Apple Intelligence Foundation Language Models https://arxiv.org/pdf/2407.21075

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 4
questions: When training larger models, do similar issues arise earlier or later?
soundness: 3
strengths: 1. The paper uses comprehensive experiments and analysis to link the instabilities in FP8 training to SwiGLU weight alignment.
2. The paper provides an efficient implementation of Smooth-SwiGLU.
3. The paper is well-written.
summary: This paper 1) identifies instability issues in FP8 training for large language models (LLMs) on trillion-token-scale datasets, 2) links the root cause to the alignment of two weight matrices of the SwiGLU neuron, and 3) proposes Smooth SwiGLU to mitigate outliers in the activations. Smooth SwiGLU achieves this by applying scaling factors before and after the last linear layer of the MLP components. Furthermore, they compared different FP8 formats to obtain the best configuration for the FP8 optimizer. The method successfully stabilizes FP8 LLM training on trillion-token-scale datasets and delivers a 5.1% throughput improvement over the baseline on 256 Intel Gaudi2 accelerators.
weaknesses: 1. The motivation of Smooth SwiGLU is that "While disabling quantization of the SwiGLU output effectively prevents divergence, it reduces the potential acceleration benefits of FP8". However, "FP8 + Smooth SwiGLU" is only 5.1% faster than the "FP8 + SwiGLU output BF16". Given that the performance gap is relatively small and the experiments were conducted only on Intel Gaudi2, it would be helpful to compare the two variants on more hardware configurations (i.e., NVIDIA and AMD GPUs) to verify whether the performance gap persists. 
2. Table 2 currently only compares the accuracy and perplexity between the proposed FP8 and the standard BF16 configurations. To provide a more comprehensive evaluation, it would be beneficial to include additional comparisons, such as 'FP8 + Smooth SwiGLU,' 'FP8 + SwiGLU output BF16,' and 'FP8 + SwiGLU output BF16 + FP8 Optimizer'. If evaluating these variants is challenging or unnecessary, please provide the corresponding discussion.

Meta-review: All reviewers agreed this paper should be accepted: it addresses an important problem, the method is thoughtfully-designed, and the paper is clearly written. A clear accept. Authors: you've already indicated that you've updated the submission to respond to reviewer changes, if you could double check their comments for any recommendation you may have missed on accident that would be great! The paper will make a great contribution to the conference!

Decision: Accept (Spotlight)