Towards Fully FP8 GEMM LLM Training at Scale

Extracted source text · arXiv 2505.20524

\documentclass{article}

\usepackage[square,numbers]{natbib}
\usepackage[preprint]{./Styles/neurips_2025}

\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc}    % use 8-bit T1 fonts
\usepackage{hyperref}       % hyperlinks
\usepackage{url}            % simple URL typesetting
\usepackage{booktabs}       % professional-quality tables
\usepackage{amsfonts}       % blackboard math symbols
\usepackage{nicefrac}       % compact symbols for 1/2, etc.
\usepackage{microtype}      % microtypography
\usepackage{xcolor}         % colors

\usepackage{amsmath}
\usepackage{graphicx}
\usepackage{wrapfig}
\usepackage{multirow}

\bibliographystyle{newnatbib}

% Commands/macros.
\newcommand{\design}{FOG}
\newcommand{\designXIELU}{\design-max}
\newcommand{\designGELU}{\design-opt}
\newcommand{\designDYT}{\design-flash}

\DeclareMathOperator{\kurt}{kurt}
\DeclareMathOperator{\id}{id}


\title{Towards Fully FP8 GEMM LLM Training at Scale}

\makeatletter
\def\@fnsymbol#1{\ifcase#1\or *\fi} % Only allow * as footnote symbol
\makeatother

\author{%
  Alejandro Hernández-Cano\thanks{Equal contribution.} \\
  EPFL \\  
  \texttt{alejandro.hernandezcano@epfl.ch} \\
  \And
  \phantom{abcdefghijk}Dhia Garbaya\footnotemark[1]\phantom{abcdefghijk} \\
  EPFL \\  
  \texttt{dhia.garbaya@epfl.ch} \\
  \AND
  \phantom{abcdefghi}Imanol Schlag\phantom{abcdefghi} \\
  ETHZ \\
  \texttt{ischlag@ethz.ch} \\
  \And
  \phantom{abcdefghijk}Martin Jaggi\phantom{abcdefghijk} \\
  EPFL \\
  \texttt{martin.jaggi@epfl.ch}
}


\begin{document}

\maketitle

\begin{abstract}
Despite the significant potential of FP8 data formats for large language model (LLM) pre-training, their adoption has been limited due to challenges in maintaining stability at scale.
Existing approaches often rely on suboptimal fine-grained FP8 kernels or fall back to higher-precision matrix multiplications (GEMMs) in sensitive components, such as attention projections, compromising potential throughput gains.
We introduce a new class of LLM architectures that, for the first time, support FP8 computation for all GEMMs within transformer blocks during both forward and backward passes. 
This enables unprecedented throughput gains, particularly at scale, while matching the downstream performance of standard BF16 training.
Our architecture design reduces large outlier activations, promoting stable long-term FP8 training.
In addition, we identify key metrics to monitor low-precision training and predict potential future divergences.
\end{abstract}


\section{Introduction}
Recent progress in the training of transformer-based Large Language Models (LLMs) has significantly advanced the field of language modelling.
Scaling up both model size and training data remains a reliable strategy to enhance their performance~\citep{kaplan2020scalinglawsneurallanguage}. Consequently, state-of-the-art models are typically trained at scale using extensive datasets~\citep{10.5555/3495724.3495883,grattafiori2024llama3herdmodels,qwen2025qwen25technicalreport}, requiring substantial computational resources---often in the order of millions of GPU hours.

Thus, the development of efficient training techniques has become increasingly essential.
One of the main research avenues for efficiency is the use of lower-precision number formats to accelerate training on appropriate hardware accelerators.
Recently, the use of 8-bit floating-point (FP8) formats has shown promising results~\citep{deepseekai2025deepseekv3technicalreport,fishman2025scaling,peng2023fp8lmtrainingfp8large}.
However, the widespread adoption of current approaches is still limited due to suboptimal throughput benefits.
One cause of slowdowns is the use of higher precision in those General Matrix Multiplications (GEMMs) which are most sensitive, such as attention score computation, while another issue is the overhead caused by more granular FP8 scaling strategies.
One of the key challenges in 8-bit LLM training originates from the relatively narrow dynamic range offered by FP8 formats and thus higher risk of underflows and overflows, especially with the prevalence of large outlier features observed in the LLM's neural activations during training~\citep{dettmers2022gptint,sun2024massive,he2024understanding}.
To mitigate this issue, modern FP8 training recipes utilise various scaling techniques before casting from higher-precision formats---typically BF16~\citep{kalamkar2019studybfloat16deeplearning} for activations---to FP8 formats used in matrix multiplications.
These scaling approaches help maximize the effective use of FP8's limited dynamic range, reducing the risk of underflows and overflows. 

Recent work has introduced promising FP8 training recipes by employing multiple scaling factors per single tensor \citep{deepseekai2025deepseekv3technicalreport}, allowing for a finer and more precise casting to lower precision.
Yet, this comes with an efficiency overhead, diminishing the large gains initially expected from using FP8.
Another strategy involves adjusting the standard SwiGLU-based transformer architecture~\citep{shazeer2020gluvariantsimprovetransformer} to prevent emergent outliers from occurring~\citep{fishman2025scaling}.
This area of optimization remains underexplored, as most works focus on FP8 GEMMs within the linear projections within the transformer, while maintaining higher precision for other GEMMs, namely those involved in the dot product attention mechanism. % (FP8DPA).
We refer to such training strategies simply as FP8 training. We label the approach of also including attention computation as FP8DPA.

In this paper, we introduce \design: the fast and outlier-guarded set of LLM architectures specifically designed to mitigate large activation outliers and enable efficient large-scale FP8 training with low-overhead scaling strategies.
For the first time to our knowledge, this approach enables FP8 GEMMs not only in the linear projection, but also within the attention mechanism of each transformer block, achieving unprecedented throughput improvements of up to 40\% in the 8B parameter model scale, while maintaining equivalent downstream performance compared to higher precision baselines.
In addition, we present a comprehensive recipe for monitoring, explaining, and predicting training instabilities that might not surface in the early stages of training.
This approach provides researchers with greater confidence in the long-term stability of FP8 training recipes, reducing the need for costly, full-scale experiments when using new architectures.
Furthermore, we provide an interestingly useful observation about larger models' tendency to diverge later in training with FP8.

Our key contributions are the following:
\begin{itemize}
    \item We introduce \design, a novel set of architectures designed to minimise outlier features during training.
    Our recipe allows stable training with FP8 computation of all GEMMs inside the transformer blocks, surpassing throughput of existing BF16 approaches by \textbf{40\%}.
    It is a significant step forward towards a purely FP8 matrix multiplication pre-training framework.
    
    \item Our design achieves equivalent quality results to BF16 baselines, while providing a significant speed-up.
    We empirically attest both performance and stability on various model sizes (0.4B, 1.5B, 8B) and data regimes, up to 14x the Chinchilla optimal data budget~\cite{10.5555/3600270.3602446}.

    \item Using kurtosis, we provide a recipe to judge an architecture's robustness to FP8 training in long data regimes using diagnostics from shorter runs.
    We use this recipe to explain previously observed divergence behaviour at scale, and offer a wide range of empirical results to demonstrate its usefulness.
    We believe this contribution allows FP8 training insights on future transformer variants developed by the community, without the need for expensive full-scale experiments.
\end{itemize}


\section{Background}

Due to its limited dynamic range, FP8 tensors are particularly prone to overflows and underflows when representing extreme values. 
The FP8 formats come in two standard forms \citep{micikevicius2022fp8formatsdeeplearning}: E4M3 and E5M2, each with different trade-offs.
The first format, with four exponent bits and three mantissa bits, offers higher precision.
In contrast, the E5M2 format, with five exponent bits and two mantissa bits, provides a broader dynamic range at the cost of reduced precision.
Existing large-scale distributed training frameworks such as DeepSpeed or Megatron~\citep{shoeybi2020megatronlmtrainingmultibillionparameter} leverage this distinction by employing E4M3 for tensors in the forward pass to maintain precision and E5M2 for the backward pass to handle the broader dynamic range of gradients effectively.
Nonetheless, both formats have much lower representation capacity than half- or single-precision formats.
Therefore, various scaling strategies are applied when casting tensors down to FP8 in order to make more efficient use of this restricted range. These strategies are mainly tensorwise and fall into two main categories: delayed scaling and just-in-time scaling (JIT).
Delayed scaling uses information from previous training iterations to determine the scaling factor of the tensor for the ongoing iteration, requiring a single pass on the data along with storing a short history of useful metrics observed across an interval of past iterations.
JIT scaling, on the other hand, can hinder the gains from using FP8 because it uses the distribution of the tensor being produced---in higher precision---to compute the scalar, before casting the input and performing the GEMM in FP8, requiring at least two passes through data.
A more recent approach aims to make scaling more robust by using multiple scaling factors per tensor, allowing different tensor blocks to have different scaling factors~\citep{rouhani2023microscalingdataformatsdeep,deepseekai2025deepseekv3technicalreport}.
This leads to a more precise FP8 casting within each block.
Naturally, this finer scaling strategy induces a larger overhead on such GEMM kernels relative to the tensorwise delayed scaling recipe.

Ensuring stable FP8 training remains challenging.
It becomes problematic when certain activations produce large outliers during training, making such a low-precision representation unfeasible and leading to rapid divergence.
Prior work introduced the term \textit{massive activations}, a phenomenon similar to outlier features, and showed their crucial role in LLMs' capabilities~\citep{sun2024massive}.
Understanding the dynamics of these outliers is crucial for explaining FP8 divergence and identifying the network components responsible for them.
One notable source of such outliers' amplification has been identified to be the widely adopted SwiGLU (Swish Gated-Linear-Unit) activation function \citep{shazeer2020gluvariantsimprovetransformer}.
Replacing it with a scaled variant, SmoothSwiGLU regulates large outliers and was shown to stabilize previously diverging FP8 training runs and ensure their convergence~\cite{fishman2025scaling}.  

Further examination has shown that not only is SwiGLU an outlier amplifier, but Gated Linear Units (GLUs) in general, as well as pre-normalization layers, suggesting that improper signal propagation is the root cause of outliers \citep{he2024understanding}.
Removing these components and equipping transformers with QK entropy regularization mechanisms such as QK RMS Normalization~\citep{henry-etal-2020-query}, producing the OP architectures, has been shown to diminish late-stage outliers observed by orders of magnitude, while providing equivalent prediction quality.
While OP architectures were shown to be beneficial for post-training quantization, its use for FP8 pre-training remains unexplored.
Finally, an alternative to pre-normalization layers are post-normalization layers~\citep{Liu2021SwinTV}.
Long data regime trainings have confirmed their superiority in terms of training stability with the standard BF16 mixed precision training~\citep{olmo20252olmo2furious}.

\section{\design: Fast and Outlier-Guarded FP8-suited architectures}
\label{sec:fog}

\begin{wrapfigure}{h}{0.3\linewidth}
    \centering
    \includegraphics[width=0.9\linewidth]{Figures/frog.pdf}
    \caption{\textbf{\design{} transformer.}}
    \label{fig:frog}
\end{wrapfigure}

Our architecture base, as illustrated in Figure~\ref{fig:frog}, makes key changes to widely-used transformer networks~\cite{10.5555/3295222.3295349}.
The pre-normalization block before the attention mechanism and FFN is removed.
In addition, a normalization mechanism in the attention is added to prevent entropy collapse, a key training instability in transformers~\citep{10.5555/3618408.3620117}, from occurring.
This mechanism can take the form of a QK RMSNorm block~\citep{henry-etal-2020-query}:
\[\mathcal{N}_{\boldsymbol\gamma}(\boldsymbol{x}) := \frac{1}{\mathrm{rms}(\mathbf{x})} \boldsymbol\gamma \odot \mathbf{x}, \qquad \mathrm{rms}(\mathbf{x}) := \frac{\|\mathbf{x}\|_2}{\sqrt{D}},\]
where $\mathbf{x} \in \mathbb{R}^D$, $\boldsymbol\gamma \in \mathbb{R}^N$ is the learnable \textit{gains} vector, $\odot$ is the Hadamard product and $\|\cdot\|_2$ is the $\ell_2$-norm.
Alternatively, the $\tanh_\alpha(x) := \tanh(\alpha x)$ element-wise activation function, where $\alpha \in \mathbb{R}$ is trainable, can be applied to query and key tensors.
This activation has been shown to have regularization effects akin to RMS normalization blocks~\citep{zhu2025transformersnormalization}, while being computationally more efficient.

Further, the input of the first transformer block is scaled by $\sigma^{-1}$ to maintain unit variance activations at initialization, where $\sigma$ is the chosen standard deviation of the network's random initialization.
Finally, to enhance performance, a learnable normalization block is applied before the residual connections.
This takes the form of a LayerScale~\citep{Touvron_2021_ICCV} block, $\mathrm{LayerScale}_{\boldsymbol\gamma}(\mathbf{x}) := \boldsymbol\gamma \odot \mathbf{x}$, where $\boldsymbol\gamma \in \mathbb{R}^{D}$ is a learnable gain vector, or an RMSNorm block, resulting in a post-normalized architecture~\citep{Liu2021SwinTV}.
In both cases, the learnable gains vector is initialized to $\mathbf{1}/\sqrt{\mathrm{num\_layers}}$ and keeps the residual branch unnormalized, allowing proper signal propagation~\citep{he2024understanding}.
Our architecture suite is specified in Table~\ref{tab:design}.

\begin{table}[h]
\centering
\begin{tabular}{@{}llll@{}}
\toprule
Model & QK-Regularization  & Activation & Normalization \\ \midrule
\designXIELU & RMSNorm* & xIELU~\cite{huang2025derivingactivationfunctionsusing} & Post-RMSNorm     \\
\designGELU & RMSNorm* & GeLU       & Post-RMSNorm     \\
\designDYT & Tanh*    & GeLU       & Post-RMSNorm     \\
OP~\citep{he2024understanding} & RMSNorm & GeLU       & LayerScale \\ \bottomrule
\end{tabular}
\vspace{0.5em}
\caption{\textbf{\design{} architecture suite compared with OP.} Regularizations marked with * indicate that gains are not trainable. Each variant offers different trade-offs, with \designDYT{} having the higher throughput and \designXIELU{} observed to have better downstream quality.}
\label{tab:design}
\end{table}

While the OP architecture already offers several guards to prevent large outliers from occurring, we observed that it remains an impractical choice for FP8DPA training.
In Section~\ref{sec:results-stability} we show that, like all other architectures tested, it suffers a fatal loss divergence early during training.
We isolate the two components responsible for OP's incompatibility with FP8DPA training: the trainable RMSNorm gains vector $\boldsymbol\gamma$, and the lack of any normalization. %, resulting in the architecture \designGELU.
We identify the use of post-normalization as not prone to the outlier tendency pre-normalization networks have.
\begin{figure}[h]
    \centering
    \includegraphics[width=0.9\linewidth]{Figures/stepbystep.pdf}
    \vspace{-0.2em}
    \caption{\textbf{From OP to \designGELU{} step by step.} The first run to diverge is OP, while OP with frozen QK RMSNorm is able to survive through the stable phase of training, but it experiences a significant divergence during learning rate cooldown, which starts around 42B tokens in. The converged run adds post-normalization to the previous recipe, resulting in \designGELU{}.}
    \label{fig:stepbystep}
\end{figure}

Figure~\ref{fig:stepbystep} ablates the components transitioning from OP to \designGELU.
We can see that freezing the trainable QK RMSNorm gains results in a significantly more stable training.
We attribute the early divergence of OP to the fact that uncontrolled QK normalization leads to an explosion of its gains when training in low precision. Note that these gains are generally not weight-decayed. We experimentally observe this explosion, confirm that using $L_2$ regularization helps delay the divergence. We finally opted for freezing the gains to a constant value as it is simpler and sufficient, doesn't compromise performance, and offers a small speedup.
Our ablations highlight that a constant value for the gains slightly greater than 1 improves loss.
Therefore, to retain its benefit after removing the $\boldsymbol\gamma$ gains vector, we increase the standard $s =1/\sqrt{D_{qk}}$ attention softmax temperature--a tiny optimization trick offering equivalent attention score matrix $\mathbf{S}$:
\[\mathbf{S} = \frac{1}{\sqrt{D_{qk}}} \mathcal{N}_{\boldsymbol\gamma_0}(\mathbf{Q}) \mathcal{N}_{\boldsymbol\gamma_0} (\mathbf{K})^\top = \frac{1}{\sqrt{D_{qk}}} \left(\frac{\gamma_0 \mathbf{Q}}{\mathrm{rms}(\mathbf{Q})} \right) \left(\frac{\gamma_0 \mathbf{K}}{\mathrm{rms}(\mathbf{K})}\right)^\top = \frac{\gamma_0^2}{\sqrt{D_{qk}}} \mathcal{N}_{\mathbf{1}}(\mathbf{Q}) \mathcal{N}_{\mathbf{1}}(\mathbf{K})^\top,\]

Finally, we empirically show that the addition of post-normalization is important to ensure convergence with FP8DPA during the learning rate decay phase.

Prior works also favored post-normalization over pre-normalization \citep{olmo20252olmo2furious}, providing evidence of their better stability in BF16 training. We extend this observation to our FP8 setting and we confirm that learnable LayerScale blocks alone, even with controlled QK regularization, cannot ensure convergence during this last phase.
We attribute this late divergence of OP to the fact that LayerScale blocks without normalizations are not enough to handle FP8 outliers, potentially due to the considerable changes in model statistics following the learning rate decay, that are summed up in the residual connections resulting in huge activation outliers for last layers, as highlighted by the increasing pattern of outliers on each transformer block's output in Figure~\ref{fig:layers}. 

\begin{figure}[h]
    \centering
    \includegraphics[width=0.45\linewidth]{Figures/layers.pdf}
    \caption{\textbf{Kurtosis of QKV tensors during FP8DPA learning rate cooldown with OP+frozenQK architecture.} Later layers exhibit significantly larger activation outliers.} 
    \label{fig:layers} 
\end{figure}



\section{Long-term outlier dynamics}
\label{sec:kurt}

To analyse the outliers present in neural network activations, we use kurtosis as a metric of the extremity of deviations of activation values (such as by outliers).
We define the kurtosis $\kurt(\mathbf{x})$ of a vector $\mathbf{x} \in \mathbb{R}^D$ as the scalar
\[\kurt(\mathbf{x}) := \frac{\mu[\mathbf{x}^4]}{\sigma^2[{\mathbf{x}^2}]},\]
where $\mu$ and $\sigma^2$ are the sample mean and variance, respectively, and exponentiation is taken element-wise.
Given an activation tensor $\mathbf{X} \in \mathbb{R}^{N \times C \times D}$, where $N$, $C$, and $D$ are the batch size, sequence length, and hidden size respectively, we define its kurtosis as the average $\kurt(\mathbf{X}) := \frac{1}{NC} \sum_{n=1}^N \sum_{c=1}^C \kurt(\mathbf{x}_{nc})$.

Under this definition, $\kurt(\mathbf{x})$ is maximized when few elements of $\mathbf{x}$ reach extremely large values, relative to the variance across the entire vector, i.e., when large outlier features are present.
This definition has been used to analyse outliers in BF16 training in previous work~\cite{he2024understanding} and, unlike the standard definition of kurtosis~\citep{10.1093/biomet/4.1-2.169} in the probability theory literature, this definition does not center $\mathbf{x}$ to have zero-mean.
For our use, this is consistent with the fact that FP8 kernels do not shift their inputs before scaling and casting down.
We track the dynamics of kurtosis in key activations.
Namely, the inputs of the second projection in FFNs, the QKV matrix, and the output of each transformer block.
Unless explicitly stated, we report the average activation kurtosis across all layers.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.875\linewidth]{Figures/kurtosis.pdf}
    \caption{\textbf{Loss and kurtosis training dynamics} of 1.5B \designXIELU{} and a pre-normalized GLU-based model trained for a total of 100B tokens. Note the logarithmic kurtosis axis in the transformer block output and FFN activations.}
    \label{fig:kurt}
\end{figure}

Using these activations, we can analyse the emergence of large outlier features at different stages during training.
Figure~\ref{fig:kurt} demonstrates an equivalent loss progression to the baseline while offering up to orders of magnitude lower kurtosis in some activations.
Note that, unlike previous FP8 approaches, \design{} architectures are trained with FP8 attention computations, introducing more quantization errors.
As a result, the kurtosis of key, query, and value projections becomes particularly relevant.

First, these architectures exhibit a sub-linear to logarithmic trend in the long-term growth of QKV outliers, as consistently shown by kurtosis in Figure \ref{fig:kurt}.
This behavior supports their robustness to FP8DPA, as it suggests that prohibitively longer training would be required to see a substantial increase in kurtosis.
Our extended run is consistent with the hypothesis as it does not exhibit any sign of divergence.

Second, baseline Llama exhibits late divergence during FP8 training (with attention in BF16), which has been attributed solely to the quadratic behavior of its gated activation function—emerging when weights become sufficiently aligned late in training \citep{fishman2025scaling}. In our extended 420-billion-token run using the \designXIELU{} architecture, we employ the inherently quadratic xIELU activation function \citep{huang2025derivingactivationfunctionsusing} and observe smooth training with kurtosis levels orders of magnitude lower than those of baseline Llama. This challenges the completeness of the prior explanation. From another perspective, maintaining very low kurtosis throughout training enables the stable use of quadratic activations. This is particularly interesting given that such activations are known to produce linear gradients, which benefit the backward pass—likely contributing to \designXIELU{}'s superior performance over GeLU-based variants.

Finally, we show in Figure~\ref{fig:divergence} an example of a diverging FP8DPA run, comparing it with the successful \designXIELU{} training.
This emphasises the importance of tracking tensor-level metrics such as kurtosis to potentially predict later divergences, before common global metrics like the loss and gradient norms show any symptoms of divergence.
In this example, while loss irrefutably diverged around the 15B token mark and the gradient norm consistently spiked no earlier than 12B tokens, the QKV kurtosis was already diverging from the expected sub-linear growth consistently seen across different architectures as early as the 3B mark, giving a potential early divergence sign.

\begin{figure}[h]
    \centering
    \includegraphics[width=\linewidth]{Figures/predict.pdf}
    \caption{\textbf{Training dynamics of a failed and a successful FP8DPA run.} Kurtosis exhibits atypical behaviour much earlier than when the loss diverged.}
    \label{fig:divergence}
\end{figure}


\section{Experimental Results}
\label{sec:experiments}

We perform extensive experiments to verify our architecture across several scales.
We use the FineWeb-Edu~\cite{penedo2024finewebdatasetsdecantingweb} text corpus, filtering out any web opt-out domains with \texttt{robots.txt}, resulting in a rigorous data-compliant corpus~\citep{fan2025performantllmsethicalquantifying}.
The data is tokenised using a 131K vocabulary BPE tokenizer.
We keep a consistent context length of 4096 during all experiments.
In terms of the optimisation algorithm, we use AdamW~\citep{loshchilov2018decoupled} with default hyperparameters.
Our learning rate schedule is comprised of three phases: Warm-up, Steady, and Decay phases (WSD), as it has been shown to provide equivalent performance to the cosine schedule~\citep{hagele2024scaling}, while allowing to train beyond fixed training durations.
For the models, we train 390M, 1.5B and 8B parameter models for different token counts, specified at each experiment.
Our baseline architecture follows the Llama3 8B model design~\citep{grattafiori2024llama3herdmodels}, with the 390M and 1.5B being adapted to their respective sizes.
Since Llama3 uses a gated linear unit, unlike the OP and \design{} variants, we increase the FFN sizes of OP and \design{} to maintain an equal parameter count.
Further details regarding architectures and hyperparameters are available in the Appendix.

Our hardware infrastructure consists of nodes with 4 Nvidia Grace Hopper GPUs each.
Our distributed training framework is adapted from Megatron-LM~\citep{shoeybi2020megatronlmtrainingmultibillionparameter}, which uses Transformer Engine~\citep{noauthor_nvidiatransformerengine_2025} FP8 recipes.
With 390M parameters, our experiments reach 50B tokens.
We scaled 1.5B experiments to 125B tokens to obtain more meaningful evaluations. In addition to the absence of late-in-training outlier amplification from \design{}'s non-gated activation functions and our kurtosis progression guarantees, we further validate our method's stability on long data regimes by continuing pretraining \designXIELU{} up to 420B tokens.
Finally, we scale the model size to 8B and train for 20B tokens.
We show the divergence of other architectures with FP8DPA while \design{} variants converge and match the baseline Llama3 BF16 loss, while being 35-40\% faster.
During all experiments, we use the delayed scaling strategy, with a margin of zero and a history length of 1024 steps.

We make our implementation, along with reproducibility steps for our experiments, public under the anonymized repository \url{https://github.com/anonymous4375934/FOG}.

\subsection{FP8 stability}
\label{sec:results-stability}

We compare our approach with different architectures proposed in the literature.
Namely, the OP architecture, OLMo2, Llama3, and Llama3 with the SmoothSwiGLU activation following the previous work~\cite{fishman2025scaling}, adapting each network to 390M and 1.5B parameter count.
In the case of the Llama3 baseline, we also provide results on the 8B scale.
Results are shown in Figure~\ref{fig:stability}.
This experiment displays the unsuitability of existing architectures for FP8DPA training, as all of them diverge.
For the case of the OP and OLMo2 architecture, despite having an attention outlier-mitigation strategy---the QK RMSNorm---divergence is still observed, as discussed in Section~\ref{sec:fog}.

\begin{figure}[h]
    \centering
    \includegraphics[width=\linewidth]{Figures/diverge.pdf}
    \caption{\textbf{Cross-entropy loss plots across different architectures with FP8DPA training.} No other tested architecture was able to surpass the 20B token mark without diverging.}
    \label{fig:stability}
\end{figure}

Another interesting observation from these experiments is the tendency of larger models to diverge in later stages of training compared with similar but smaller models.
We validate its consistency across architectures, as presented in Table~\ref{tab:divergence-clean}.
This observation has not been raised before, possibly due to the longer time needed for FP8-only models to diverge, in contrast to FP8DPA training.
While this trend could have many practical implications, exploring it fully falls outside the scope of this work, and we encourage future research in this direction.
%This might be explained by the common use of lower learning rates with larger models, requiring more training iterations to reach sufficiently high levels of outliers.


\begin{table}[h]
\centering
\begin{tabular}{@{}llc@{}}
\toprule
Architecture & Model Size & \begin{tabular}[c]{@{}c@{}}Divergence Mark \\ (in billions of tokens)\end{tabular} \\ \midrule
Llama3       & 390M       & 0.7 \\
Llama3       & 1.5B       & 1.1 \\
Llama3       & 8B         & 6.6 \\ \midrule
OLMo2        & 390M       & 3.3 \\
OLMo2        & 1.5B       & 15.9 \\ \bottomrule
\end{tabular}
\vspace{0.5em}
\caption{\textbf{Token mark when loss was observed to diverge.}}
\label{tab:divergence-clean}
\end{table}

\subsection{Downstream performance}
\label{sec:results-evals}
We compare our proposals with the higher-precision Llama3 baseline across a wide range of standard benchmarks to measure their downstream performance.
In Table~\ref{tab:benchmarks}, we report some of the most relevant scores along with an average* across a larger set of tasks, detailed in the Appendix. All \design{} variants offer comparable downstream performance with the higher precision Llama3 baseline with \designXIELU{} architecture, even outperforming it.
The 1.5B models are trained on 125B tokens, whereas smaller models are trained on 50B tokens. We note that we initially tested the idea of cooling down the previously constant weight decay during the learning cooldown phase, aiming to conserve model weights' norm~\citep{kosson2024rotational}. We show later in the Appendix that such intervention has no noticeable effect on stability neither performance. We had to conserve it for consistency and fair comparison across ablations. Finally, we use identical hyper-parameters such as batch size, learning rate and weight decay.
\begin{table}[h]
\centering
\begin{tabular}{@{}lllll@{}}
\toprule
Model               & Hellaswag             & ARC                  & PIQA                & Average* \\ \midrule
\small{Llama 390M}           & {\small 33.5 $\mid$ -}     & {\small 47.9 $\mid$ -}    & {\small 65.0 $\mid$ -}   & {\small 39.8 $\mid$ -} \\
\small{\designXIELU}        & {\small 36.5 $\mid$ 36.3}  & {\small 62.9 $\mid$ 62.5} & {\small 68.0 $\mid$ 68.2} & {\small \textbf{41.2} $\mid$ 40.8} \\
\small{\designGELU }        & {\small 36.1 $\mid$ 35.6}  & {\small 61.5 $\mid$ 61.3} & {\small 68.0 $\mid$ 67.8} & {\small 40.9 $\mid$ 40.4} \\
\small{\designDYT }         & {\small 35.9 $\mid$ 35.2}  & {\small 61.5 $\mid$ 60.4} & {\small 68.1 $\mid$ 68.2} & {\small 40.5 $\mid$ 40.3} \\ \midrule
\small{Llama 1.5B}          & {\small 43.7 $\mid$ -}     & {\small 71.8 $\mid$ -}    & {\small 72.5 $\mid$ -}   & {\small 46.1 $\mid$ -} \\
\small{\designXIELU}        & {\small 43.3 $\mid$ 43.4}  & {\small 71.6 $\mid$ 73.0} & {\small 72.6 $\mid$ 73.3} & {\small 46.0 $\mid$ \textbf{47.1}} \\
\small{\designGELU}         & {\small 43.3 $\mid$ 42.7}  & {\small 71.3 $\mid$ 70.8} & {\small 72.6 $\mid$ 72.0} & {\small 45.7 $\mid$ 46.0} \\
\small{\designDYT}          & {\small 42.8 $\mid$ 41.9}    & {\small 70.9 $\mid$ 69.4}   & {\small 72.2 $\mid$ 72.0}  & {\small 45.7 $\mid$ 44.9} \\ \bottomrule
\end{tabular}
\vspace{0.5em}
\caption{\textbf{Performance across various tasks.} For each task and model size, the first score results from the BF16 ablation and the second from the FP8DPA one.}
\label{tab:benchmarks}
\end{table}

\subsection{Efficiency}
Table~\ref{tab:throughput} explores the efficiency of our architectures in the 1.5B and 8B model scales for FP8 training.
We compare our FP8DPA solution with the BF16 baseline and the FP8 training enabled by the use of SmoothSwiGLU, which is, to the best of our knowledge, the only dense architecture proposal demonstrated to work at scale with FP8.
Note that the SmoothSwiGLU cannot benefit from enabling FP8 GEMMs in the attention mechanism, as it was shown to suffer a big loss divergence in Figure~\ref{fig:stability}.
Furthermore, our throughput gains increase as the model size increases.  The GEMM input tensors increase in size and consume significantly more time during the overall forward-backward pass, compared with other operators.

\begin{table}[h]
\centering
\begin{tabular}{@{}cllc@{}}
\toprule
Size & Model & Precision & \begin{tabular}[c]{@{}c@{}}Throughput\\ (tokens/second/GPU)\end{tabular} \\ \midrule
\multirow{5}{*}{8B}   & Llama              & BF16   & 9105                     \\
                      & \designXIELU       & FP8DPA & 12344 (+35.5\%)           \\
                      & \designGELU        & FP8DPA & \underline{12414} (+36.3\%)          \\
                      & \designDYT         & FP8DPA & \textbf{12764 (+40.2\%)} \\
                      & Llama+SmoothSwiGLU & FP8    & 12228 (+34.3\%)          \\ \midrule
\multirow{5}{*}{1.5B} & Llama              & BF16    & 46470                    \\
                      & \designXIELU       & FP8DPA & 53551 (+15.2\%)          \\
                      & \designGELU        & FP8DPA & 53877 (+15.9\%)          \\
                      & \designDYT         & FP8DPA & \underline{54848} (+18.0\%) \\
                      & Llama+SmoothSwiGLU & FP8    & \textbf{54903 (+18.1\%)} \\ \bottomrule
\end{tabular}
\vspace{0.7em}
\caption{\textbf{Throughput measures with \design{} versus other baselines.} Using eight 4xGH200 nodes with Zero-1 sharding~\citep{DBLP:conf/sc/RajbhandariRRH20}, without model parallelism, for 8B models and a single 4xGH200 node for 1.5B models. Notably, in the 8B scale, all \design{} variants outperform other architectures.}
\label{tab:throughput}
\end{table}

\subsection{Long-data regimes and FP16 optimizer states}

To further justify the viability of FP8DPA long training with \design, we train our 1.5B \designXIELU{} to 420B tokens, way beyond the previously identified 200B tokens divergence mark of Llama2-7B~\citep{fishman2025scaling}.
Note that our observation of smaller models' tendancy to diverge earlier with FP8DPA, and the long-term outlier analysis in Section~\ref{sec:kurt} further proves the sufficiency of such a training duration.
We also switch to use FP16 optimizer states and BF16 gradients after 130B tokens, saving up memory previously used by full precision states, gradients, and model parameters master copy.
We display the learning dynamics of our approach in Figure~\ref{fig:hero}.
The language modeling loss exhibits equal to better smoothness compared to the corresponding Llama baseline.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.6\linewidth]{Figures/hero.pdf}
    \caption{\textbf{Long-data training regimes.} \designXIELU{} 1.5B FP8DPA is trained on 420B tokens. The higher precision Llama3 shorter experiment is also included as reference.} 
    \label{fig:hero}
\end{figure}

\section{Limitations}
\label{sec:limitations}

Despite its robustness, outstanding throughput  and relative flexibility, the \design{} set of architectures remains mainly limited by two aspects.
Since the introduction of transformers~\citep{10.5555/3295222.3295349}, various iterations have followed for more stability and higher performance.
Future better variants might be incompatible with our design and thus require different recipes to ensure FP8DPA training.
Second and more importantly, \design{} is trained with BF16 GEMM for the final projection. %, also called the language modeling head.
This operator is known to be very sensitive to outliers and has always been used with half-precision in forward-backward FP8 training approaches, including ours.

Moreover, this work focuses on maximising throughput gains while ensuring both strong stability and downstream performance.
We also show it to be robust to half-precision optimizer states and gradients, allowing considerable memory savings.
Yet, we do not discuss further casting down optimizer states to FP8.
Although this will allow for comparatively less memory savings, especially as the optimizer states are usually sharded~\citep{DBLP:conf/sc/RajbhandariRRH20}, it may help to avoid some memory saturation issues in compute-constrained settings.
Finally, even though we managed to scale training to a long data regime at 1.5B scale, showed consistent observations across model sizes, and provided further guarantees, our experiments were limited to a short regime of 20B tokens at the 8B parameter scale.
Due to computational constraints, we decided to keep the study of these limitations for future work.

\section{Conclusion}

In this paper we demonstrate, for the first time, stable LLM training with fully FP8 computations within the transformer blocks--including the attention mechanism--without sacrificing performance.
We tested FP8DPA training across a wide set of previously proposed architectures and show that they consistently diverge early during training, highlighting the difficulty of FP8DPA training and novelty in our results.
Moreover, in contrast with other granular scaling recipes, we use the low-overhead delayed scaling FP8 strategy.
Our design provides on-par downstream quality with the higher precision baseline, while offering up to \textbf{40\%} faster training.
We scale our 1.5B model to 420B tokens, 14x the Chinchilla-optimal data budget for its size.
Our work brings the community one step closer to fully FP8 GEMM training at scale i.e including the language modeling head.

We further justify the long-term stability of our architecture by observing the outlier training dynamics across key activations by using kurtosis.
The use of kurtosis to track outliers present during training was shown to provide meaningful insights to favour certain architectural components or to predict future instabilities, as it is a quantitative metric that measures outliers.

\section*{Acknowledgement}
This work was initiated during the master’s thesis of Alejandro Hernández Cano at EPFL.
It was supported as part of the Swiss AI Initiative by a grant from the Swiss National Supercomputing Centre (CSCS) under project ID a06 on Alps.

\bibliography{newref}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\newpage
\appendix

\section{Hyperparameters}
We detail the selection of hyperparameters used in Table~\ref{tab:hyperparams}.
For the case of \designDYT{}, the $\alpha_0$ initialization value of $\tanh_\alpha$ entropy-regularization is 0.5 for all model sizes.
All models use a linear warmup schedule, and 1-sqrt cooldown schedule.
The long-data 1.5B \designXIELU{} experiment was trained for a total of 400,000 steps, consuming approximately 419.4B tokens, using the same hyperparameters as the shorter run--except for the absence of a cooldown phase.

\begin{table}[h]
\centering
\begin{tabular}{@{}lccc@{}}
\toprule
Hyperparameter              & 390M      & 1.5B                 & 8B                   \\ \midrule
Layers ($L$)                & 16        & 16                   & 32                   \\
Hidden size ($D$)           & 1024      & 2048                 & 4096                 \\
FFN hidden size             & 4096      & 8192                 & 14336                \\
Attention heads             & 8         & 16                   & 32                   \\
QK groups                   & 4         & 8                    & 8                    \\
Softmax scale* ($s$)        & 0.17678   & \multicolumn{2}{c}{0.125}                   \\
Tied embeddings             & Yes       & \multicolumn{2}{c}{No}                      \\ \midrule
Weight decay ($\lambda$)    & \multicolumn{3}{c}{0.1}                                 \\
AdamW $\beta_1$             & \multicolumn{3}{c}{0.9}                                 \\
AdamW $\beta_2$             & \multicolumn{3}{c}{0.95}                                \\
Gradient clip value         & \multicolumn{3}{c}{1.0}                                 \\ \midrule
Context length $T$          & \multicolumn{3}{c}{4096}                                \\
Global batch size           & 128       & 256                  & 512                  \\
Total training steps        & 100,000   & 125,000              & 10,000               \\
Peak learning rate ($\eta$) & $10^{-3}$ & $2.5 \times 10^{-4}$ & $1.5 \times 10^{-4}$ \\
Warmup $\eta$ steps         & 5,000     & 2,500                & 1,250                \\
Cooldown $\eta$ steps       & 20,000    & 25,000               & N/A                  \\
Minimum $\eta$              & \multicolumn{3}{c}{$10^{-8}$}                           \\ \bottomrule
\end{tabular}
\vspace{0.5em}
\caption{\textbf{Hyperparameters used in experiments.} Note that FFN hidden size indicates the dimensionality of each linear projection in gated activation functions; networks without GLUs use $1.5\times$ this value to match the parameter count. Softmax scale specified only applies to \design{} models, all other models follow the standard $s = 1/\sqrt{D_{QK}}$.}
\label{tab:hyperparams}
\end{table}


\section{Architectures}
We provide detailed formulations for all architectures presented in this paper. Our transformer architecture consists of the following components in sequence:
\begin{enumerate}
    \item Input token embeddings
    \item An input scaling factor $u \in (0,\infty)$, which may equal $1$
    \item A series of $L$ transformer blocks as described below
    \item A final normalization function $N^{(\mathrm{final})}$, which may be the identity
    \item A linear output layer
\end{enumerate}
The transformer block is defined as
\[\mathrm{block}(\mathbf{X}) := \hat{\mathbf{X}} + \left(N_2^{(\mathrm{post})} \circ \mathrm{FFN} \circ N_2^{(\mathrm{pre})}\right)(\hat{\mathbf{X}}), \qquad \hat{\mathbf{X}} := \mathbf{X} + \left(N_1^{(\mathrm{post})} \circ \mathrm{GQA} \circ N_1^{(\mathrm{pre})}\right)(\mathbf{X}).\]
The $N_i^{(*)}$ are normalization layers that may be the identity, and $\mathrm{FFN}(\mathbf{X})$ is a two-layer FFN with a nonlinear activation function $\varphi$ and no bias.
The $\mathrm{GQA}$ follows the standard grouped-query self-attention definition with softmax scaling factor $s$ and Rotary Position Embeddings.
Each attention head uses the definition
\[\mathrm{attnhead}(\mathbf{X}) := \mathrm{selfattn}\left(N_Q^{(\text{QK})}(\mathbf{XW}^{(Q)}), N_K^{(\text{QK})}(\mathbf{XW}^{(K)}), \mathbf{XW}^{(V)}\right),\]
where $N^{(\text{QK})}$ is the entropy-regularization mechanism, and $\mathrm{selfattn} = \mathbf{PV}$.
The $\mathbf{P}$ matrix is the attention probabilities matrix
\begin{equation}
   \mathbf{P} :=\mathrm{Softmax}\left( s \mathbf{QK}^\top + \mathbf{M}\right) \label{eq:p}
\end{equation}
With this notation, Table~\ref{tab:architectures} details the architecture families used in the project.

\begin{table}[h]
\centering
\begin{tabular}{@{}lcccccl@{}}
\toprule
Parameter                       & $u$        & $N^{(\mathrm{final})}$            & $N^{(\mathrm{pre})}$              & $N^{(\mathrm{post})}$                     & $N^{(\mathrm{QK})}$               & $\varphi$               \\ \midrule
Llama3                          & 1               & $\mathcal{N}_{\boldsymbol\gamma}$ & $\mathcal{N}_{\boldsymbol\gamma}$ & $\id$                                     & $\id$                             & $\mathrm{SwiGLU}$       \\
Llama3+SmoothSwiGLU             & 1               & $\mathcal{N}_{\boldsymbol\gamma}$ & $\mathcal{N}_{\boldsymbol\gamma}$ & $\id$                                     & $\id$                             & $\mathrm{SmoothSwiGLU}$ \\
OLMo2                           & 1               & $\mathcal{N}_{\boldsymbol\gamma}$ & $\id$                             & $\mathcal{N}_{\boldsymbol\gamma}$         & $\mathcal{N}_{\boldsymbol\gamma}$ & $\mathrm{SwiGLU}$       \\
OP$^{(\text{a})}$               & $\sigma_0^{-1}$ & $\id$                             & $\id$                             & $\mathrm{LayerScale}_{\boldsymbol\gamma}$ & $\mathcal{N}_{\boldsymbol\gamma}$ & $\mathrm{GeLU}$         \\
\designXIELU{}$^{(\text{a,b})}$ & $\sigma_0^{-1}$ & $\id$                             & $\id$                             & $\mathcal{N}_{\boldsymbol\gamma}$         & $\mathcal{N}$                     & $\mathrm{xIELU}$        \\
\designGELU{}$^{(\text{a,b})}$  & $\sigma_0^{-1}$ & $\id$                             & $\id$                             & $\mathcal{N}_{\boldsymbol\gamma}$         & $\mathcal{N}$                     & $\mathrm{GeLU}$         \\
\designDYT{}$^{(\text{a,b})}$   & $\sigma_0^{-1}$ & $\id$                             & $\id$                             & $\mathcal{N}_{\boldsymbol\gamma}$         & $\tanh_\alpha$                    & $\mathrm{GeLU}$         \\ \bottomrule
\end{tabular}
\vspace{0.5em}
\caption{\textbf{Architecture details for the used models.} Models with (a) initialize the post-normalization gains with $\boldsymbol{\gamma}_0 = \mathbf{1}/\sqrt{L}$. Models with (b) have frozen gains in the QK entropy regularization $N^{(\mathrm{QK})}$. The $\id$ is the identity function, $\sigma_0$ is the chosen initialization standard deviation, $\mathcal{N}$ is the RMS normalization. The $u$ input scaling is not trainable.}
\label{tab:architectures}
\end{table}

\section{FP8 training}
For all our experiments, we used Transformer Engine's delayed scaling implementation with history length $\ell = 1024$ and margin $m = 0$.
Mathematically, given a history of abs-max values, denoted $H = \{h_t\}_{t=1}^\ell \subseteq [0, \infty)$, of a tensor $\mathbf{X}$, we define its scaling factor as:
\[\rho(\mathbf{X}) := \frac{\mathrm{FP8MaxValue}}{2^m\, \max H}\]
where $\mathrm{FP8MaxValue} \in (0, \infty)$ is the maximum value representable with the FP8 format used.
We update the history using $H \gets \{\max_{x \in \mathbf{X}} |x|\} \cup \{h_t\}_{t=2}^\ell$ to use for this activation in the next iteration.
The end-to-end FP8 matrix multiplication is
\[\mathrm{GEMM}(\mathbf{X}, \mathbf{Y}) := \frac{1}{\rho(\mathbf{X}) \rho(\mathbf{Y})} \mathrm{FP8GEMM}(\mathrm{FP8cast}(\rho(\mathbf{X}) \mathbf{X}), \mathrm{FP8cast}(\rho(\mathbf{Y}) \mathbf{Y})),\]
where $\mathrm{FP8GEMM}$ receives FP8 tensors and returns the BF16 result.
We further detail the precision used for every matrix multiplication during our FP8 and FP8DPA experiments in Table \ref{tab:fp8}.

\begin{table}[h]
\centering
\begin{tabular}{@{}lcccc@{}}
\toprule
Method & \begin{tabular}[c]{@{}c@{}}Linear\\ operators\end{tabular} & \begin{tabular}[c]{@{}c@{}}Attention scores\\ $\mathbf{QK}^\top$\end{tabular} & \begin{tabular}[c]{@{}c@{}}Attention-value\\ GEMM $\mathbf{PV}$\end{tabular} & \begin{tabular}[c]{@{}c@{}}Output\\ layer\end{tabular} \\ \midrule
FP8    & FP8                                                        & BF16                                                                          & BF16                                                                         & BF16                                                   \\
FP8DPA & FP8                                                        & FP8                                                                           & FP8                                                                          & BF16                                                   \\ \bottomrule
\end{tabular}
\vspace{0.5em}
\caption{\textbf{Comparison between FP8 methods.} The FP8DPA method allows for all---excluding the output head---GEMMs computations to be done with FP8 precision. In contrast, FP8 training uses higher precision for all GEMMs in the core attention. The linear operators are linear layers of the form $\mathrm{Linear}_{\mathbf{W}}(\mathbf{X}) = \mathbf{XW}$: namely the FFN linear layers, QKV projections and attention output projection. See Equation~\eqref{eq:p} for the definition of the attention probability matrix $\mathbf{P}$. }
\label{tab:fp8}
\end{table}

\section{Evaluations}
We selected the following set of benchmarks: \texttt{ARC-Easy}, \texttt{CommonsenseQA}, \texttt{HellaSwag}, \texttt{LAMBADA-OpenAI}, \texttt{LAMBADA-standard}, \texttt{OpenBookQA}, \texttt{PIQA}, \texttt{SocialIQA}, and \texttt{WinoGrande}. \\
We used a standard open-source LLM evaluation package for conducting these evaluations, as cited in the code repository \url{https://github.com/anonymous4375934/FOG}.

In Table \ref{tab:benchmarks}, we report raw accuracy scores as percentages on three key benchmarks as well as the average over the full set of tasks mentioned above. 
In the following table \ref{tab:ext_evals}, we provide all scores along with their estimation errors for the 1.5B model size, demonstrating that the slight differences observed across many values are statistically insignificant.

\begin{table}[h]
\centering
\begin{tabular}{@{}lccccl@{}}
\toprule
Architecture                       & Llama3     & \designXIELU{}    & \designGELU{}  & \designDYT{}                                 \\ \midrule
Hellaswag              & \small $43.7 \mid -$    & \small $43.3 \mid 43.4$  & \small $43.3 \mid 42.7$       &\small $42.8 \mid 41.9$      &  $\pm 0.5$  \\
ARC-easy               & \small $71.8 \mid -$    & \small $71.6 \mid 73.0$  & \small $71.3 \mid 70.8$       &\small $70.9 \mid 69.4$      &  $\pm 0.9$  \\
PIQA                  & \small $72.5 \mid -$    & \small $72.6 \mid 73.3$  & \small $72.5 \mid 72.0$       &\small $72.2 \mid 72.0$      &  $\pm 1.0$  \\
Commonsense-qa        & \small $19.6 \mid -$    & \small $20.2 \mid 22.2$  & \small $19.3 \mid 21.2$       &\small $21.1 \mid 20.8$      &  $\pm 1.2$  \\
Lambada-openai        & \small $44.5 \mid -$    & \small $43.7 \mid 44.6$  & \small $44.3 \mid 44.5$       &\small $42.4 \mid 41.2$      &  $\pm 0.7$  \\
Lambada-standard      & \small $38.9 \mid -$    & \small $37.0 \mid 39.5$  & \small $37.7 \mid 37.9$       &\small $35.9 \mid 33.8$      &  $\pm 0.7$  \\
Openbook-qa           & \small $26.2 \mid -$    & \small $28.0 \mid 27.0$  & \small $26.8 \mid 27.0$       &\small $28.4 \mid 27.6$      &  $\pm 2.0$  \\
Social-iqa            & \small $41.3 \mid -$    & \small $41.8 \mid 42.0$  & \small $41.7 \mid 40.8$       &\small $41.0 \mid 40.8$      &  $\pm 1.1$  \\
Winogrande            & \small $56.5 \mid -$    & \small $55.6 \mid 58.7$  & \small $54.6 \mid 57.4$       &\small $56.9 \mid 56.8$      &  $\pm 1.4$  \\  \midrule
Average                 & \small $46.1 \mid -$    & \small $46.0 \mid \textbf{47.1}$  & \small $45.7 \mid 46.0 $       &\small $45.7 \mid 44.9$      &  $\pm 0.3$  \\
\bottomrule
\end{tabular}
\vspace{0.6em}
\caption{\textbf{More detailed results at 1.5B scale.} For each model and each task, the first score results from BF16 training and the second from FP8DPA training. }
\label{tab:ext_evals}
\end{table}

\textbf{Weight decay cooldown}
As mentioned in Section \ref{sec:results-evals}, we experimented with cooling down the weight decay, often used as a constant value equal to $0.1$ that is coupled with the learning rate, to see if it solves the OP+frozenQK architecture's consistent divergence during the learning rate decay phase. We also tested it on other architectures and, to optimize the use of resources, we had to keep it later for the final experiments.
This trick helped stabilize the weights' norm indeed, but couldn't solve the divergence issue. However, it had no effect on final performance nor on stability. The following table highlights this no-effect claim at 1.5B scale.
\begin{table}[h]
\centering
\begin{tabular}{@{}lccc@{}}
\toprule
Setting                  & WD               &  Loss           & Average score   \\ \midrule
OP+FrozenQK              & cooldown         & diverges        &   - \\
OP+FrozenQK              & constant         & diverges       &   -   \\
\designGELU{}           & cooldown          & converges      &  46.0 $\pm0.3$  \\ 
\designGELU{}           & constant          & converges       &   46.3 $\pm0.3$   \\   \bottomrule
\end{tabular}
\vspace{0.6em}
\caption{\textbf{Weight Decay (WD) during the LR decay phase}. If constant, it equals $0.1$. Else, it starts from $0.1$ and is proportional to LR.}
\label{tab:coolwd}
\end{table}

\section{Computational Resources}
Our experiments were conducted on nodes equipped with 4 Grace Hopper GPUs each. We typically used 4, 8, and 16 nodes for our 390M, 1.5B, and 8B parameter experiments, respectively, with minor variations across different runs. Importantly, all throughput measurements were taken under identical hardware configurations.
Table~\ref{tab:shame} details the computational resources in GPU hours (GPUh) required for our main experimental results. This includes the computational cost of training all architectures that diverged during FP8DPA training, the FP8DPA and BF16 stable training runs for our three architectures, and the BF16 Llama3 baseline. The aggregation includes node start-up times, computation lost due to node failures, and overhead from calculating and logging kurtosis metrics. 
The complete research project required additional computational resources beyond those specified in the table, as we conducted numerous preliminary experiments and explored ideas that did not appear in the final paper.

\newpage
\begin{table}[h]
\centering
\begin{tabular}{@{}lc@{}}
\toprule
Group                 & GPUh     \\ \midrule
Divergent runs (FP8DPA)     & 886      \\
Llama3 baselines (BF16) & 1,395  \\
\design{} experiments & 11,162 \\ \bottomrule
\end{tabular}
\vspace{0.5em}
\caption{\textbf{GPU hours used for the main experiments.}}
\label{tab:shame}
\end{table}

\end{document}
\begin{thebibliography}{}
\makeatletter
\relax
\def\mn@urlcharsother{\let\do\@makeother \do\$\do\&\do\#\do\^\do\_\do\%\do\~}
\def\mn@doi{\begingroup\mn@urlcharsother \@ifnextchar [ {\mn@doi@} {\mn@doi@[]}}
\def\mn@doi@[#1]#2{\def\@tempa{#1}\ifx\@tempa\@empty \href {http://dx.doi.org/#2} {doi:#2}\else \href {http://dx.doi.org/#2} {#1}\fi \endgroup}
\def\mn@eprint#1#2{\mn@eprint@#1:#2::\@nil}
\def\mn@eprint@arXiv#1{\href {http://arxiv.org/abs/#1} {{\tt arXiv:#1}}}
\def\mn@eprint@dblp#1{\href {http://dblp.uni-trier.de/rec/bibtex/#1.xml} {dblp:#1}}
\def\mn@eprint@#1:#2:#3:#4\@nil{\def\@tempa {#1}\def\@tempb {#2}\def\@tempc {#3}\ifx \@tempc \@empty \let \@tempc \@tempb \let \@tempb \@tempa \fi \ifx \@tempb \@empty \def\@tempb {arXiv}\fi \@ifundefined {mn@eprint@\@tempb}{\@tempb:\@tempc}{\expandafter \expandafter \csname mn@eprint@\@tempb\endcsname \expandafter{\@tempc}}}

\bibitem[\protect\citeauthoryear{Brown et~al.,}{Brown et~al.}{2020}]{10.5555/3495724.3495883}
Tom~B. Brown, et~al., 2020, in Proceedings of the 34th International Conference on Neural Information Processing Systems. NIPS '20.
Curran Associates Inc., Red Hook, NY, USA

\bibitem[\protect\citeauthoryear{Corporation}{Corporation}{}]{noauthor_nvidiatransformerengine_2025}
{NVIDIA} Corporation,, {NVIDIA}/{TransformerEngine}, \url {https://github.com/NVIDIA/TransformerEngine}

\bibitem[\protect\citeauthoryear{DeepSeek-AI et~al.,}{DeepSeek-AI et~al.}{2025}]{deepseekai2025deepseekv3technicalreport}
DeepSeek-AI, et~al., 2025, DeepSeek-V3 Technical Report (\mn@eprint {arXiv} {2412.19437}), \url {https://arxiv.org/abs/2412.19437}

\bibitem[\protect\citeauthoryear{Dettmers, Lewis, Belkada  \& Zettlemoyer}{Dettmers et~al.}{2022}]{dettmers2022gptint}
Tim Dettmers, Mike Lewis, Younes Belkada,  Luke Zettlemoyer, 2022, in Alice~H. Oh, Alekh Agarwal, Danielle Belgrave,  Kyunghyun Cho, eds, Advances in Neural Information Processing Systems. \url {https://openreview.net/forum?id=dXiGWqBoxaD}

\bibitem[\protect\citeauthoryear{Fan, Sabolčec, Ansaripour, Tarun, Jaggi, Bosselut  \& Schlag}{Fan et~al.}{2025}]{fan2025performantllmsethicalquantifying}
Dongyang Fan, Vinko Sabolčec, Matin Ansaripour, Ayush~Kumar Tarun, Martin Jaggi, Antoine Bosselut,  Imanol Schlag, 2025, Can Performant LLMs Be Ethical? Quantifying the Impact of Web Crawling Opt-Outs (\mn@eprint {arXiv} {2504.06219}), \url {https://arxiv.org/abs/2504.06219}

\bibitem[\protect\citeauthoryear{Fishman, Chmiel, Banner  \& Soudry}{Fishman et~al.}{2025}]{fishman2025scaling}
Maxim Fishman, Brian Chmiel, Ron Banner,  Daniel Soudry, 2025, in The Thirteenth International Conference on Learning Representations. \url {https://openreview.net/forum?id=E1EHO0imOb}

\bibitem[\protect\citeauthoryear{Grattafiori et~al.,}{Grattafiori et~al.}{2024}]{grattafiori2024llama3herdmodels}
Aaron Grattafiori, et~al., 2024, The Llama 3 Herd of Models (\mn@eprint {arXiv} {2407.21783}), \url {https://arxiv.org/abs/2407.21783}

\bibitem[\protect\citeauthoryear{H{\"a}gele, Bakouch, Kosson, allal, Werra  \& Jaggi}{H{\"a}gele et~al.}{2024}]{hagele2024scaling}
Alexander H{\"a}gele, Elie Bakouch, Atli Kosson, Loubna~Ben allal, Leandro~Von Werra,  Martin Jaggi, 2024, in The Thirty-eighth Annual Conference on Neural Information Processing Systems. \url {https://openreview.net/forum?id=Y13gSfTjGr}

\bibitem[\protect\citeauthoryear{He, Noci, Paliotta, Schlag  \& Hofmann}{He et~al.}{2024}]{he2024understanding}
Bobby He, Lorenzo Noci, Daniele Paliotta, Imanol Schlag,  Thomas Hofmann, 2024, in The Thirty-eighth Annual Conference on Neural Information Processing Systems. \url {https://openreview.net/forum?id=npJQ6qS4bg}

\bibitem[\protect\citeauthoryear{Henry, Dachapally, Pawar  \& Chen}{Henry et~al.}{2020}]{henry-etal-2020-query}
Alex Henry, Prudhvi~Raj Dachapally, Shubham~Shantaram Pawar,  Yuxuan Chen, 2020, in Trevor Cohn, Yulan He,  Yang Liu, eds, Findings of the Association for Computational Linguistics: EMNLP 2020. Association for Computational Linguistics, Online, pp 4246--4253, \mn@doi{10.18653/v1/2020.findings-emnlp.379}, \url {https://aclanthology.org/2020.findings-emnlp.379/}

\bibitem[\protect\citeauthoryear{Hoffmann et~al.,}{Hoffmann et~al.}{2022}]{10.5555/3600270.3602446}
Jordan Hoffmann, et~al., 2022, in Proceedings of the 36th International Conference on Neural Information Processing Systems. NIPS '22.
Curran Associates Inc., Red Hook, NY, USA

\bibitem[\protect\citeauthoryear{Huang \& Schlag}{Huang \& Schlag}{2025}]{huang2025derivingactivationfunctionsusing}
Allen~Hao Huang, Imanol Schlag, 2025, Deriving Activation Functions Using Integration (\mn@eprint {arXiv} {2411.13010}), \url {https://arxiv.org/abs/2411.13010}

\bibitem[\protect\citeauthoryear{Kalamkar et~al.,}{Kalamkar et~al.}{2019}]{kalamkar2019studybfloat16deeplearning}
Dhiraj Kalamkar, et~al., 2019, A Study of BFLOAT16 for Deep Learning Training (\mn@eprint {arXiv} {1905.12322}), \url {https://arxiv.org/abs/1905.12322}

\bibitem[\protect\citeauthoryear{Kaplan et~al.,}{Kaplan et~al.}{2020}]{kaplan2020scalinglawsneurallanguage}
Jared Kaplan, et~al., 2020, Scaling Laws for Neural Language Models (\mn@eprint {arXiv} {2001.08361}), \url {https://arxiv.org/abs/2001.08361}

\bibitem[\protect\citeauthoryear{Kosson, Messmer  \& Jaggi}{Kosson et~al.}{2024}]{kosson2024rotational}
Atli Kosson, Bettina Messmer,  Martin Jaggi, 2024, Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks, \url {https://openreview.net/forum?id=Kr7KpDm8MO}

\bibitem[\protect\citeauthoryear{Liu et~al.,}{Liu et~al.}{2021}]{Liu2021SwinTV}
Ze~Liu, et~al., 2021, 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp 11999--12009

\bibitem[\protect\citeauthoryear{Loshchilov \& Hutter}{Loshchilov \& Hutter}{2019}]{loshchilov2018decoupled}
Ilya Loshchilov, Frank Hutter, 2019, in International Conference on Learning Representations. \url {https://openreview.net/forum?id=Bkg6RiCqY7}

\bibitem[\protect\citeauthoryear{Micikevicius et~al.,}{Micikevicius et~al.}{2022}]{micikevicius2022fp8formatsdeeplearning}
Paulius Micikevicius, et~al., 2022, FP8 Formats for Deep Learning (\mn@eprint {arXiv} {2209.05433}), \url {https://arxiv.org/abs/2209.05433}

\bibitem[\protect\citeauthoryear{OLMo et~al.,}{OLMo et~al.}{2025}]{olmo20252olmo2furious}
Team OLMo, et~al., 2025, 2 OLMo 2 Furious (\mn@eprint {arXiv} {2501.00656}), \url {https://arxiv.org/abs/2501.00656}

\bibitem[\protect\citeauthoryear{PEARSON}{PEARSON}{1905}]{10.1093/biomet/4.1-2.169}
KARL PEARSON, 1905, \mn@doi [Biometrika] {10.1093/biomet/4.1-2.169}, 4, 169

\bibitem[\protect\citeauthoryear{Penedo, Kydlíček, allal, Lozhkov, Mitchell, Raffel, Werra  \& Wolf}{Penedo et~al.}{2024}]{penedo2024finewebdatasetsdecantingweb}
Guilherme Penedo, Hynek Kydlíček, Loubna~Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro~Von Werra,  Thomas Wolf, 2024, The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale (\mn@eprint {arXiv} {2406.17557}), \url {https://arxiv.org/abs/2406.17557}

\bibitem[\protect\citeauthoryear{Peng et~al.,}{Peng et~al.}{2023}]{peng2023fp8lmtrainingfp8large}
Houwen Peng, et~al., 2023, FP8-LM: Training FP8 Large Language Models (\mn@eprint {arXiv} {2310.18313}), \url {https://arxiv.org/abs/2310.18313}

\bibitem[\protect\citeauthoryear{Qwen et~al.,}{Qwen et~al.}{2025}]{qwen2025qwen25technicalreport}
Qwen, et~al., 2025, Qwen2.5 Technical Report (\mn@eprint {arXiv} {2412.15115}), \url {https://arxiv.org/abs/2412.15115}

\bibitem[\protect\citeauthoryear{Rajbhandari, Rasley, Ruwase  \& He}{Rajbhandari et~al.}{2020}]{DBLP:conf/sc/RajbhandariRRH20}
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase,  Yuxiong He, 2020, in SC. p.~20, \url {https://doi.org/10.1109/SC41405.2020.00024}

\bibitem[\protect\citeauthoryear{Rouhani et~al.,}{Rouhani et~al.}{2023}]{rouhani2023microscalingdataformatsdeep}
Bita~Darvish Rouhani, et~al., 2023, Microscaling Data Formats for Deep Learning (\mn@eprint {arXiv} {2310.10537}), \url {https://arxiv.org/abs/2310.10537}

\bibitem[\protect\citeauthoryear{Shazeer}{Shazeer}{2020}]{shazeer2020gluvariantsimprovetransformer}
Noam Shazeer, 2020, GLU Variants Improve Transformer (\mn@eprint {arXiv} {2002.05202}), \url {https://arxiv.org/abs/2002.05202}

\bibitem[\protect\citeauthoryear{Shoeybi, Patwary, Puri, LeGresley, Casper  \& Catanzaro}{Shoeybi et~al.}{2020}]{shoeybi2020megatronlmtrainingmultibillionparameter}
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper,  Bryan Catanzaro, 2020, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (\mn@eprint {arXiv} {1909.08053}), \url {https://arxiv.org/abs/1909.08053}

\bibitem[\protect\citeauthoryear{Sun, Chen, Kolter  \& Liu}{Sun et~al.}{2024}]{sun2024massive}
Mingjie Sun, Xinlei Chen, J~Zico Kolter,  Zhuang Liu, 2024, in First Conference on Language Modeling. \url {https://openreview.net/forum?id=F7aAhfitX6}

\bibitem[\protect\citeauthoryear{Touvron, Cord, Sablayrolles, Synnaeve  \& J\'egou}{Touvron et~al.}{2021}]{Touvron_2021_ICCV}
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve,  Herv\'e J\'egou, 2021, in Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV). pp 32--42

\bibitem[\protect\citeauthoryear{Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser  \& Polosukhin}{Vaswani et~al.}{2017}]{10.5555/3295222.3295349}
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan~N. Gomez, \L{}ukasz Kaiser,  Illia Polosukhin, 2017, in Proceedings of the 31st International Conference on Neural Information Processing Systems. NIPS'17.
Curran Associates Inc., Red Hook, NY, USA, p. 6000–6010

\bibitem[\protect\citeauthoryear{Zhai, Likhomanenko, Littwin, Busbridge, Ramapuram, Zhang, Gu  \& Susskind}{Zhai et~al.}{2023}]{10.5555/3618408.3620117}
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu,  Josh Susskind, 2023, in Proceedings of the 40th International Conference on Machine Learning. ICML'23.
JMLR.org

\bibitem[\protect\citeauthoryear{Zhu, Chen, He, LeCun  \& Liu}{Zhu et~al.}{2025}]{zhu2025transformersnormalization}
Jiachen Zhu, Xinlei Chen, Kaiming He, Yann LeCun,  Zhuang Liu, 2025, Transformers without Normalization (\mn@eprint {arXiv} {2503.10622}), \url {https://arxiv.org/abs/2503.10622}

\makeatother
\end{thebibliography}