µnit Scaling: Simple and Scalable FP8 LLM Training
Authors: Saaketh Narayan
Abhay Gupta
Mansheej Paul
Davis Blalock
Venue: ICML 2025 poster (2025)
Keywords: LLM
FP8
Transformer
Model Training
Attention
Primary area: deep_learning->large_language_models
TL;DR: μnit Scaling combines stable, highly efficient training via FP8 with significant cost savings through hyperparameter transfer.

Abstract: Large language model training with 8-bit floating point (FP8) formats promises significant efficiency improvements, but reduced numerical precision makes training challenging. It is currently possible to train in FP8 only if one is willing to tune various hyperparameters, reduce model scale, or accept the overhead of computing dynamic scale factors. We demonstrate simple, scalable FP8 training that requires no dynamic scaling factors or special hyperparameters, even at large model sizes. Our method, \textit{µnit Scaling (µS)}, also enables simple hyperparameter transfer across model widths, matched numerics across training and inference, and other desirable properties. µnit Scaling is straightforward to implement, consisting of a set of minimal interventions based on a first-principles analysis of transformer operations. We validate our method by training models with parameters ranging from 1B to 13B, performing all hidden linear layer computations in FP8. We achieve quality equal to higher-precision baselines while also training up to 33% faster.

pnit Scaling: Simple and Scalable FP8 LLM Training
Saaketh Narayan
Abhay Gupta? ? Mansheej Paul
Davis Blalock?
Abstract
strate simple, scalable FP8 training method with straight- forward hyperparameter transfer large LLMs, called "umit Scaling" (uS).
Large language model training with 8-bit float- ing point (FP8) formats promises significant effi- ciency improvements, but reduced numerical pre- cision makes training challenging. It is currently possible to train in FP8 only if one is willing to tune various hyperparameters, reduce model scale, or accept the overhead of computing dy- namic scale factors. We demonstrate simple, scal- able FP8 training that requires no dynamic scaling factors or special hyperparameters, even at large model sizes. Our method, fuit Scaling (uS), also enables simple hyperparameter transfer across model widths, matched numerics across training and inference. and other desirable properties. jnit Scaling straightforward to implement, consist- ing of set of minimal interventions based on a first-principles analysis of transformer operations. We validate our method by training models with parameters ranging from 1B to 13B, performing all hidden linear layer computations in FP8. We achieve quality equal to higher-precision base- lines while also training up to 33% faster.
Our unit Scaling method builds on Unit Scaling (Blake et al., 2023), which aims to maintain unit variance in weights, ac- tivations, and gradients. To ensure this, it scales neural network operations with static constants and initializes net- work parameters to have unit variance. If all tensors used in training can maintain unit variance, they representable with sufficient range and resolution by low-precision for- mats like FP16 and FP8. However, preserving high-quality tensor representations in low-precision formats is challeng- ing for large models.
Besides faster training, several other properties are desir- able in a low-precision training scheme. Examples include minimizing extra hyperparameters, avoiding dynamic scale factor overhead, and allowing optimal hyperparameters from small models to transfer to large models. As summarized in Fig. 1, μu is the only method that provides these benefits. We elaborate on each of these properties below.
Straightforward hyperparameter transfer Tuning hy- perparameters for large LLMs is expensive. A promising way to reduce this cost is to tune the hyperparameters for smaller LLMs and "transfer" them to large ones, either by using them directly or by applying a model-size-based formula as explored in p-Parametrization (uPP)(Yng et al., 2021; 2024a;b). However, applying hyperparameter trans- fer techniques practice to low oow-precision training can be challenging: frequent divergences due numerical issues may require training in higher precisions like FP32 (Yang et al., 2021). To address this, Blake et al. (2024) introduced u-uP, which combines Unit Scaling (Blake et al., 2023) and μP to enable hyperparameter transfer in low precision. Unfortunately, compared to conventional BF16 mixed pre- cision training (henceforth termed "standard parametrized" (SP) models), both μP and u-uP have many more hyperpa- rameters to sweep over (see Table 3), diminishing realized compute savings and increasing complexity. Specific im- plementation intricacies, such as zero-initialized queries in μP or LR scaling for embeddings by fan-out in u-uP, make these schemes harder to use in practice than SP. In contrast, our pnit Scaling (uS) scheme combines μP and Unit Scaling in a greatly simplified way, making it easier to use and more
1. Introduction
Because LLM training is computationally expensive, low- precision training provides large compute savings. Mod- 6
em LLMs are typically trained in (BF16), where most computation occurs in BF16, but some components requiring higher precision (such as accumu- lators and master weights) use FP32 (Micikevicius et al., 2018). Thanks to increased hardware support for FP8 formats, mixed precision training using FP8 computation promises even greater training efficiency (Micikevicius et al., 2022). However, the reduced range and resolution of FP8 make L.MM training challenging. In this work, we demon-
Work done while at Databricks Mosaic Research Databricks Mosaic Research, San Francisco, CA. Correspondence to: Saaketh Narayan <narayan. waaketh@gmail.coom Davis Blalock <davis. @ databricks.comm>
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

pnit Scaling: Simple and Scalable FP8 LLM Training
Method
Uses FP8
Hparam transfer
Number of Hparams
No dynamic scaling factors
Scales stably to large models
Training- Inference precision match
Efficient
distributed training
BF16 mixed precision (SP)
No
No
3
Yes
Yes
Maximal Update Parametrization (IP)
No
Yes
Yes
Yes
Unit Scaling u-ui
Partially
Yes (u-uP)
7
Yes
Partially
Partially
Partially
Dynamically Scaled FP8 (SP), e.g. TE
Yes
No
3
No
Partially
μnit Scaling (ours)
Yes
Yes
3 Yes
Yes
Yes
Yes
Figure 1. Comparison of low-precision training methods. Our proposed method, junit Scaling (uS, bottom row), enables FP8 training and hyperparameter transfer at scale. Unlike existing methods. it does not use dynamic scaling, requires only small set of hyperparameters, permits FP8 computation for all hidden layers, and makes the model more easily quantizable for inference.
cost-effective. We demonstrate hyperparameter transfer of learning rate (1) and weight decay (A) to models of up to 20x larger widths.
Identifying root causes for poor numerics in conven- tional transformer blocks-for example, explaining diminishing variance self-attention outputs with in- creasing sequence position.
No Dynamic Scaling With dynamic scaling, calcu- lates per-tensor scaling factors for each weight, activation, and gradient tensor in training. These scales shift BF16 tensors into the representable ranges FP8 formats in each forward and backward pass. Typically, one also decouples the forward and backward formats, using e4m3 for weights and activations and e5m2 for gradients (Sun et al., 2019). notable example of FP8 training library that uses dynamic scaling (NVIDIA, 2023). Calculating scaling factors dynamically adds train- ing and inference overhead and complicates large-scale dis- tributed training and checkpointing.
Introducing simple method for fixing these issues that enables FP8 training in all hidden linear layers and with less overhead than existing methods. It also achieves desirable properties such as improved train- ing efficiency and matched numerics training and inference time.
NVIDIA's TransformerEngine
2. Methods
In this section, we detail the components of our proposed method, pnit Scaling (uS). The modifications to the standard transformer training scheme that μf requires are summa- rized in Table 1. We elaborate on novel components such as our handling of self-attention numerics, residual modifica- tions, and hyperparameter transfer below.
Apply to All Linear Layers Existing work on applying Unit Scaling at larger scales requires certain "critical mat- muls" (attention out projection, FFN down projection) to stay in BF16 (Blake et al., 2024). Assuming transformer model with conventional multiheaded attention and an MLP with an expansion ratio of 4, this means 41.7% all hidden linear layer FLOPs are not in FP8. In contrast, μS ensures that, regardless of scale, all hidden layers use FP8.
2.1. Self-attention Numerics
The causal self-attention mechanism at the core of decoder layers in LLMs is not variance- preserving,,,, making low- precision training challenging.
Match Inference-Time Quantization For efficient infer- ence, LLMs are often quantized to FP8 or INT8 for faster computation and reduced memory footprints (Khudia et al., 2021; Dettmers et al., 2022). Since training typically occurs in higher bitwidths (e.g.. BF16), a mismatch in precisions at training time and inference time means that some level of quantization error unavoidable, degrading model quality. Training with μS avoids this mismatch- -ince the LLM has already been trained in FP8, it is immediately ready for inference in FP8 for both weights and activations (W8A8).
Recall that standard self-attention is defined as:
QK" Vd
Attention Q, K, V) softmax
V
(I)
Proposition 2.1. Suppose have E B* and V YERR Rx
s STV, and o & Var(a]. Define - softmax(x), sid Assume that each element x; ~ N(O,, 1), and that entries Vi are independent and distributed with pv & E[]] = O, oo A Var[V]] 1. Then, up to first-order Taylor approximation, o² x 1 & for >
1.1. Contributions
Our work makes the following contributions:
Proof. Recall that by the definition of the softmax function,
2

---

juit Scaling: Simple and Scalable FP8 LLM Training
Table Components of the μS training scheme. μ$ makes the following modifications to standard decoder-only transformer training practices, A deeper explanation of these modifications is provided in Appendix A.
Modification
Description
static scaling factor applied in both forward and backward pass.
Linear layer scaling factors
fan
The final LM head uses a multiplier of fan_in instead, in line with pi.
Les-Post-LayerNorm
LayerNorm is the last operation in each residual branch instead of the first.
"Fixed" residual modification
Use a fixed constant 7 to make residuals prriance-preserving according to Eq. 11.
Unit variance initialization
All linear layer weights initialized with variance 1.
Use FP8E4M3 for weights and activations, FP8E5M2 for gradients. Before casting, clip BF16 values to FP8 dtype max. Keep embedding table and LM head in BF16.
FP8 hidden layers
Optimal 1 stays constant for input and output layers, but is scaled by Vdnus for all
Learning rate (7) scaling
Vdnuaa
hidden layers, when transferring from a base model with width diase
With fully decoupled weight decay, optimal stays constant for all layers with increasing width.
Weight decay (A) scaling
c 8: softmax Σj e5s Denote the vector of ele- ments' numerators e and the vector of denominators E
In the causal self-attention operation shown in Eq. 1, the attention logits matrix QK" row of logits for token at sequence position k has length k. For a given token, by Prop. 2.1, the output of the self- attention operation will therefore have variance inversely related to that token's sequence position k. This causes tokens that appear later in the sequence to have much smaller variance than those that appear earlier, as shown in Fig. 2.
vdd is causally masked such that the
5 as d, such that d. Since 2; (N, 1), n is log-normally distributed and d is sum of log-normals. This implies that!:
In e//2
o c(c- 1)
Ma = 1/2
a = ke(e - 1)
(2)
Covin, d o e(e 1)
To address this issue, we make use of basic property of the variance of linear combinations of independent random variables. With a(k) denoting the outputs of self-attention applied over a sequence of length k, the variance of a(k) (denoted 2(()) is the variance of a sum of k random vari- ables {X....., Na3 with coefficients E R
We can then use first-order Taylor approximations to esti- mate the moments of the ratio d shown in Casella & Berger (2002), to obtain:
n
Pn 1 Md
h/ =
E
(3)
d
k
22
2n Cov(n, d]
Var Σ cxx;
- Var[X.] = c.v. v,
0.- = Var
2
P'
Pd
(4)
1
e-1 1
12
k3
where Us A Var[X,], and the equality holds all X, are independent. If Vi: v, 1, we further have a2 (a()) ||c||2-
Note that Eq. 3 holds exactly from the fact that all & entries in s are positive and must sum to 1. Now, because each element a; = Σ s; V.y, with independent entries Vij. and with the fact that μv = 0 and o. = 1, the mean and variance of can be determined as:
Now recall that the softmax operation outputs positive co- efficients s that sum to 1. This means that if we simply set coefficients = Si, we obtain:
Pa
UaRV= 0
(5)
)²(0) = |||l2 = V E
Σ
8 =1
(8)
-
- + J2 +o2
2
-1
(6)
That is, by taking the square root of attention scores, at- tention can be made variance-preserving for independent value tokens. This modification, which we term "Square- Root Softmax attention", is shown in Eq. 9. Square-Rov Softmax attention also easily implemented via modern
k
k22
The first term dominates for large k and so 0²
'See Appendix A.2 for the derivation Cov (n, d]
3

---

pnit Scaling: Simple and Scalable FP8 LLM Training
attention kernels like Flex-Attention (Dong et al., 2024).
Cosine Similarities of Value Tokens
Observed Value Tokens
Random, Tokens
QK vdz.
Attention (Q; K, V) softmax
V
1024
1024
11:38
1530
In practice, standard self-attention does have diminishing as sequence position increases; however, the observed variance is consistently higher than predicted by the above analysis of independent elements. This same effect is ob- served even when using Square-Root Softmax attention, causing observed a to increase over sequence position in- stead (Fig. 2).
2360
2560
-10"
2672
3072
5864
4000
4000
1106
-10
Sequence Position
Sequence Position
Figure 3. Value tokens in text are highly correlated. Comparison of cosine similarity between observed value tokens a text data distribution versus value tokens sid (O, 1). Repeated tokens in the value matrix, an unavoidable result of token frequency in real text data, lead to higher-than teaatthhnn-nnnnnn sequence position increases (cf. Fig. 2).
Attention Output a vs. Sequence Position (at Initialization)
2a
Standard Deviation (α)
2
29
Actention
Standard (predicted)
2-1
Standard (obeerved) Seyt Softmax (predicted) Sogt Softmax (obsarved)
shown in Fig. 4(b). All μu models we train use Res-Post- LayerNorm.
2:2
2-1
2-4
2.2. Residual Modification Schemes
2-3
500
1.0k
1.5k
2.0k
2.5k
3.0k
3.5k
4.0k
Every skip connection in neural network adds another tensor to the residual stream. Summing all these tensors tends to increase the variance of the residual stream deeper in the network. To make residual connections variance- preserving instead, Blake et al. (2023) proposed replacing simple summation with weighted summation, where the weights and of the skip connection and residual branch satisfy a² b² = 1. They proposed two methods for set- ting these coefficients: fixed and running-ment mean, which are shown in Eq. 11 and Eq. 12, respectively. The former uses a constant coefficient т, while the latter uses coefficients that are a function of the layer index The standard residual layer modification is shown Eq. 10.
Sequence Position
Figure 2. Attention output variance changes over sequence length. For standard attention, a decreases over sequence po- sition both when simulated with iid value tokens (light red) and when observed in training (red). Taking the square root of atten- tion scores keeps constant when simulated with iid value tokens (light blue), but during training (bluc), causes a to increase with sequence position. In practice, neither attention variant provides consistent scale across outputs.
We provide mechanistic explanation for this phenomenon: this increase in attention variance is an unavoidable conse- quence of the statistics of natural data. If all value tokens are truly independent, then Square-Rood Root Softmax attention keeps da constant. However, due to a high number of re- peated tokens in real text data, value tokens are often highly correlated (Fig. 3). Due to this correlation, da will be higher than predicted, and in the case of standard self-attention, diminish more slowly with respect to the token position.
standard : 1141 = 21 + fmm)
(10)
fixed(T) 2141 VI-T 21 + VT f(())
(11)
running-mean : :141 V VII
T f(m) (12)
To address this inconsistency in attention output vari- ance, we use Res-Post-LLayerNom placement, as shown in Fig. 4(a). This architecture change consists of moving the normalization operation from the start of each resid- ual branch the end, and was first proposed in Liu al. (2022) for training stability. Res-Post- Lasteosot ensures consistent for all tokens in the residual stream, regard- less of sequence position, correlation with other tokens, or the distribution of attention scores. A convergence test on 10yer models validating the Res-Post-LayerNorm trans- former against the standard Pre-LayerNorm transformer is
As shown in Fig. 5, we found that using either modifica- tion is better than the standard approach, with the fixed scheme providing better convergence than the running-mean scheme. All μS models we train therefore use the fixed scheme. We set the coefficient 7 based on the depth using the results in Appendix A.3.
2.3. Hyperparameter Transfer with pnit Scaling
Zero-shot hyperparameter transfer allows hyperparameters to be tuned on a small proxy network, then directly used

---

pnit Scaling: Simple and Scalable FP8 LLM Training
100 layer model convergence, Residual modification
4.0
2.9
Train Loss
2.8
8.7
(a)
3.0-
9900
9950
Pre-LayerNorm
Fixed (t 0.1)
2.5
Bunning-Mean
Standard
Layer norm
Attention
Layer norm
0
2000
4000
6000
8000
10000
Batch
Res-Post-Layerker
Figure 5. Residual modification schemes affect jnit Scaled model convergence. The fixed residual modification (green, Eq. 11) achieves better training convergence for deep transformers than the running-mes mean residual modification (blue, Eq. 12). The fixed residual coefficient for this model is r = 0.1. Both of these settings outperform the standard residual layer modification (red, Eq. 10).
Attention
Layer norm
FFN
Layer norm
(b)
100 layer model convergence, LN placement
4.0
2.7075
2.7050
Train Loss
2,7025
2,7000
on much larger networks without any further tuning (Yang et al., 2021). The width of the small proxy network is typically referred to as the "base width", or doase. Because it eliminates the need to sweep hyperparameters at large scale, such hyperparameter transfer yields massive compute savings.
3.0
SP, pre-LN
2.5
145, res-post-LN
2000
4000
6000
8000
10000
Batch
Hyperparameter transfer with unit Scaling follows from neural network equivalencies set forth in Yang et al. (2021, Appendix J.2.1), reproduced below for convenience. As detailed in Blake et al. (2024), Equations 13, 14, and 15 define the hidden layer in a model undergoing training. All hidden layers are initialized with weights Wo drawn from a normal distribution with variance b², use learning rate of c, and have an output multiplier a. X and Y denote input and output activation matrices respectively; is the training time step: and 4, & Lo, ... Vl,) denotes the weight update for time step t using prior loss gradients.
(c)
100 layer model convergence, LN placement
4,0
2.675
Train Loss
3.5
3.0
2523
2060
2913
res-post-LN
2.5
us, pre-LN
2000
4000
6000
8000
10000
Batch
Figure 4. Les-PoostLayyerNrm.. (a) Pre- Pre-LayeNNomm transformer architecture versus Post-os-LLayyroorm architecture. Res-Post- LayerNorm moves the LayerNorm operation from the start of each residual branch to the end (Liu et al., 2022). This ensures con- sistent variance across tokens when added to the residual stream. In contrast, Pre- Pre-LayerNorm networks permit unnormalized repre- sentations with inconsistent variance to be added the residual stream, as shown with self-attention outputs in Fig. 2. (b) Conver- gence test loss curves with 100-layer models show that μ$ with Dos-Post-LaycrNon achieves nearly identical convergence versus SP with Tre-LayerNorm. (c) Additional convergence tests with 100-layer models show that Pes-Post--ayerNNe achieves better convergence over Pre-LayerNorm with pS.
Wo ~ (00,83
(13)
Y XW.
(14)
W = W.+C (V VL)
(15)
Under Adam-like optimizers, the output of this hidden layer is invariant to any scale factor 0 > 0 that changes b, c as:
t ad, b b b/0, c c/0
(16)
Under μP, 1, b- Vianjn and set 0 Vian we obtain:
fan in If we instead
1
a
b b= C
(17)
Tan in
fan
Notice that Mam and b 1 are exactly the output multiplier and unit initialization that Unit Scaling requires.
5

---

jnit Scaling: Simple and Scalable FP8 LLM Training
Therefore, the learning rate for hidden layers should scale as Vinn in for Unit Scaled models. This leads to the μf hyperparameter transfer scheme in Table 2.
global batch size of 64 and sequence length of 1024 (i.e.. 655M total tokens). SP models use Pre-LayerNNom place- ment and are trained in both BFI6 and FP8 (using Trans- formerEngine). μu models were trained in both BF16 and FP8 and use LessPost-LLayeNorre placement (Fig. 4). μu used base models of width 256. For all models described in this and subsequent sections, we used the Lion optimizer (Chen al., 2023) with fully decoupled weight decay and a cosine learning rate schedule decaying 10% of the maxi- mum learning rate. For details on why Lion is an Adam-like optimizer for hyperparameter transfer, please refer to Ap- pendix A.4. All models were trained on Nvidia H100 GPUs using the Databricks MosaicML. LLMFoundry (MosaicML, 2022a), Composer (MosaicMI. 2021), and Streaming (Mo- saicML, 2022b) libraries.
In practice, given base model with width duase a new model with width dnew and optimal base model learning rate Tasse μm keeps Nhew constant for the embedding table, all LayerNorm parameters, and the LM head. The learning rate only changes for hidden layers, with Neww
dihaas Taase v Vagen
Table μS scaling rules. To transfer hyperparameters across model widths with μS, initialize layers, scale their outputs, and modify their learning rates as shown here.
Weight Type
Input Layer
Final I Layer Hidden I Layers
Hy Hyperparameters: We evaluate hyperparameter transfer over learning rate (1)) and weight decay (A). While μP Yang et al. (2021) does not give a theoretical basis for À transfer over width, we evaluate its transfer empirically because of its practical importance. Prior work by Lingle (2024) has shown that μP does not admit transfer of X with Adam W. However, Wang & Aitchison (2024) found that optimal À should scale with model size. To elucidate how À scales with model width, we jointly sweep over both and A. We use fully decoupled weight decay, motivated by findings from Wortsman et al. (2024) that doing so results in more stable training. 17 and À are swept over powers of 2. Based on the relationship between the residual coefficient T and depth in Appendix A.3, the residual coefficient 7 is 0.4 for these four-layer models.
Init. Var. 1
Output Mult. 1
1/fan in
1/ M/yfan 1/yfan_in
Adam-like LR
1
In addition to enabling hyperparameter transfer, μn also re- quires sweeping over much smaller set of hyperparameters than existing schemes (Table 3).
Table Required hyperparameters in transfer schemes. Hy- perparameters used in practice to train transformer models under various schemes. While yP and related schemes provide better hyperparameter transfer than SP, they require sweeping over more hyperparameters to get reasonable model quality. In contrast, μf provides hyperparameter transfer and model quality with much smaller set of hyperparameters. This makes the implementation simple and makes hyperparameter sweeps less expensive.
As shown in Fig. 6, μu models have stable optimal learning rate (17") and weight decay (X*) from width 256 up to width 8192. Mirroring previous findings, 17" for SP models de- creases as the inverse of the width. A transfer across widths is relatively stable for both model types, with μn showing the most consistency.
Scheme
# Hparams
Hparams A,
uS (ours)
3
SP
3
17, X, Finit
17, X, Fiitt
μP
6
3.2. FP8 Training at Scale
Ores: Oattn; Opput
1): X, Offinnact, sofuaaa
The previous section demonstrated hyperparameter transfer for small, shallow models. However, the real test of utility is scaling up titiiiiilliin-paramettt models. This sec- tion demonstrates that μu allows to train in FP8 while transferring hyperparameters for realistic model sizes. We also validate that our method is compatible with efficient distributed training.
u-uP
7
Ores: Ores-sttn-ratio Ojoss--oftmax
3. Results
3.1. Successful Hyperparameter Transfer
Setup: To evaluate hyperparameter transfer, we first train four-layer decoder-only LLMs with widths of 256 through 8192 using Standard Parametrization (SP) and unit Scal- ing (uS). We begin with these small models since doing so allows us to collect ground truth optimal hyperparame- ters. All models use multi-headed attention (Vaswani et al., 2017) and were trained for 10,000 training steps with a
Setup: We train 1B, 3B, 7B, and 13B parameter LLMs on approximately compute-optimet token budgets (-~20x token- to- parameteer ratio) using SP and μS, and in both BF16 and FP8, resulting in 4 individual models for each model size. The training configurations are detailed in Table 4. Based on the previous sections' hyperparameter transfer results (Fig. 6), we sweep 1) and À on small models with a base
6

---

pnit Scaling: Simple and Scalable FP8 LLM Training
Table Large model training configurations, Model training configurations for IB, 3B, 7B, and 13B models. Only μs models use the residual coefficient which is dictated by model depth using results in Appendix A.3.
Model
Params
Tokens
TPR
Steps
Batch Sz.
Seq. Len.
Width
Depth
# Heads
1B
1.6B
31.5B
19.4
7.5k
1024
4096
2048
24
16
0.3
3B
3.0B
62.9B
20.8
15k
1024
4096
2560
32
20
0.3
7B
7.3B
140.0B
19.3
16.7k
2048
4096
4096
32
32
0.3
13B
13.6B
260.1B
19.1
31k
2048
4096
5120
40
40
0.2
Optimal Learning Rate and Weight Decay for SP, Sm
responding 10 steps for IB and 3B models and 5 steps for 7B and 13B models). Training loss curves are shown in Fig. 7.
Model Width
SP, n (fixed A)
SP, A
(fixed n)
As shown in Fig. 7, μS models train stably with FP8 even as the model size increases. We successfully transfer hyperpa- rameters from a narrow base model with a width of 256 to models with widths up to 5120, demonstrating 20x width transfer (~400x fewer FLOPs per run) in realistic, practical LLM training scenarios. This validates zero-shot hyper- parameter transfer using psS. Evaluation results in Table 5 show that μm models achieve equal or better quality than SP models. These models demonstrate that μu successfully combines FP8 training with zero-shot hyperparameter trans- fer. To emphasize, all hidden layers use FP8 computation, and there are no dynamic scaling factors.
Findl CE Loss
6.0
Final GE Loos
4.0
(A)
IAS. /) (fixed A)
us, (fixed m))
6
Final CE Loos
3.8
3.6
-10
-0
-5
94
14
12 -10
We also note that at the 13B scale, attempted rem- edy the divergence of the SP FP8 model by using multiple different values of X, but this did not mitigate the frequent loss spikes and eventual divergence. μS models, by con- trast, train stably. We also show the instability in training with Unit Scaling (US) at larger scales in Appendix A.5, motivating runs only with SP and μS for our final results.
logo 0n)
logo (A)
Figure 6. With pS, optimal learning rate (17") and weight de- cay (X") are stable across widths. Optimal v (left column) and A (right column) are shown across range of model widths for models trained with SP (top row) and μS (bottom row). For each curve, the other hyperparameter is fixed at its optimal value. The base model width is 256. μU models have stable optimal and X, even when width increases 32x to 8192. As expected, 1)" for SP models decreases with width. X is relatively stable as the width increases across both model types.
3.3. FP8 Training Efficiency
To achieve state-of-the- tart FP8 distributed training efficiency with pmit Scaling, we make use of operator fusion and static scaling. As shown in Fig. 8, FP8 training with μ$ is 25-33% faster than in BF16, and 1-6% faster than FP8 training with TransformerEngine (TE) (NVIDIA, 2023). All models were benchmarked on 64 NVIDIA H100 GPUs, and characteris- tics such as batch size and distributed training configuration were held constant. While TransformerEngine has fused modules such as LayerNorr- or LayerNorm-Mi did not those modules in order to make equal comparison between μS and TE.
width of chass 256, then transfer optimal hyperparameters to large models with width dnews shown below.
SP: all layers: nhew Mosse duess
dosse
0.5
new
μS: hidden layers: Thew Viruse other layers: Vhew Thases
Vilaase dileww X
Nnew
Xisse
Abase
By relying on dynamic scaling, FP8 training with libraries like TE imposes additional overhead that is eliminated in μ.S. Calculating the absolute max both the weight and activation tensors (or storing and reading past absolute max values in delayed scaling approach) are operations that can be completely discarded in μs. Weights, activations, and gradients can be directly cast to FP8 formats, with
We the Databricks Model Gauntlet to evaluate the quality of all models on specific tasks (Dohmann, 2023; Barton, 2024). These results are shown in Table 5.
Evaluation:
We also compare model convergence via the final training cross-entropy loss averaged over the last 41. 11.9 tokens (cor-

---

pnit Scaling: Simple and Scalable FP8 LLM Training
SP-FP8
SP-BF16
US-FPP8
IS-BF16
1b
3b
7b
13b
2.24
2.20
3.5
2.575
3.5
3.5
3.5
Train Loss
2.39
2.15
2.570
2.23
2:38
7440
7480
14880
14960
16560 16640
30750 30900
2.5
2.5
2.5
2.5
0
2000
4000
6000
5000
10000
15000
0
5000
10000
15000
0
10000
20000 30000
Batch
Batch
Batch
Batch
Figure? 7. SuS models successfully train in FP8 at scale. Comparison of training loss curves for standard parametrized (SP) and pnit scaled (uS) models in both FP8 and BF16, across 1B, 3B, 7B, and 13B parameter models. yu models successfully train FP8 and converge to similar train loss values their BF16 and SP counterparts. SP FP8 models trained with TransformerEnginne (TE). In our experiments at the 13B scale, SP models trained FP8 with TE experienced frequent loss spikes and did not properly converge. We achieve of FP8 training efficiency via yS, with further details in Appendix 3.3.
Large model evaluation results. We evaluate SP and μS models in FP8 and BF16 on a varicty of tasks, with best results per eval and model size in bold. Final train loss (avg. over last <40M tokens) also shown. μu models have equal or better quality than SP models, and maintain this quality even when training in FP8 as model size increases. Note that 13B SP FP8 models failed to properly converge, denoted by an asterisk.
Table 5.
SP
JAS
SP'
SP
148
SP
BF16
FP8
BF16
FPS
BF16
FP%
BF16
FP8
BF16
FP8
BF16
FF8
BF16
FPW
BP16
FP8
Triin
2,590
2.590
2.316
2.112
2,211
2.108
ARC Easy (3 (3-shot)
52.1%
52.4%
53.4%
53.31
60.9%
60.8%
61.9%
60.81%
67.2
65.6% 616
67.1%%
68.0%
72.3%
35.7%
11.8%
69.7%
Joogundy (3-hot)
4.1%
4.3%
4.5%
3.5%
13.4%
11. 3/0
16.8%
16. 5%
27 3/6
27
32.7%
30.6%
40 25
0.2%%
43.1%
41. 1.77%
SQuAD (3-shot)
32.6%
33.2%
30.9%
31.3%
45.
47.9%
53
50.09%
57.1%
55.1%
HallaSwag
25
47.5
47.4%
57.77
66
68
74. 74.3%
BIG-bench Wikidata QA (3-bbot)
47.3%
48.6%
49.3%
50.2%
53.05
55.095
56.2%
57.5%
60.4%
60.0%6
60.0%
59.9%
66.9%
4,0%
66. 1.6
62 9%
WinoGrande (5 (5-shot)
55.0%
52.6%
51.1%
52.0%
58.8%
54.9%
59.5%
58.6%
62.
61.1%
65.7%
65.3%
70.3%%
57.8%
71.1%
70.5%
OpenBookQA (10-shor)
58.058
43.0%
44.0%
7.5
6/
8/
46.6%
PIQA (0-sho)
70.7%
71.1%
71.5%
71.2%
74.5%
75.2%
74.3%
74.3%
77.2%
77. 0%%
76.7%
76.5%
78.7%
$4.5%
80.1%
79.4%
TriviQA (3-4hor)
9.7%
10.5%
10.8%
9.7%
17.8%
17.7%
20.4%
18.7%
50.25
29.1%
12.55
33.8%
43.4%%
0.5%
3%
44.8%
Winogradd (3-shor)
64.5%
67.0%
68.9%
73.3%
75.8%
76.6%
50.6%
83.9%
86.1%
LAMBADA (0-shox)
44.8%
$1.5%
13.6%
41.3%
52.8%
54.2%
55.9%
57.4%
60.3%
60.7%
63.0%
64.6%
65.7%
34.8%
61.6%
64.3%
CoQA (0 sbot)
19.3%
21.3%
20.8%
20.0%
26.2%
25.4%
27.91
28.6%
28.2%
32.0%
33.3%
38.0%
39.8%
13.2%
44.4%
44.6%
ARC Challengu (3-hhot
25.4%
26.07 0%.00%
27.8%
25.0%
30.3%%
30.1%%
31.8%
30.9%
56 1%
35.7%
38.3%
39.0%
42.0%
27.6%
42.2%
41.5%
COPA (0-4hot)
65.0%
71.0%
76.0%
78.0%
62.0%
84.0%
BIG beach Operatoes (3-shot)
12.4%
12.9%
13.8%
14.3%
19.5%
17.1%
17.1%
18.61
21.4%
20.0%
20.0%
23.3%
31.4%
24.3%
37.6%
37.1%
GSMHK (Il-shor)
2.4%
2.6%
2.4%
2.4%
3.7%
1.7%
2.3%
2.0%%
3.9%
5.0%
4.0%
3.9%
8.7%
0.00
9.38
10.9%
constant a = Vfaninn scaling factor used in the hidden linear layers' GEMM calls, where a GEMM is defined as:
4. Conclusion
This work presents pnit Scaling (uS), an LLM training method enabling both statically-scaled FP8 computation and zero-shot hyperparameter transfer at scale. ynit Scal- ing consists of set of principled model and optimization modifications, including Ges-Post-LayerNorm, variance- preserving skip connections, unit-variance initialization. and straightforward scaling of optimization hyperparameters with model width. Compared to alternatives, juit Scaling is simpler, faster, more stable across model scales, and has fewer hyperparameters. We demonstrate successful FP8 training with hyperparameter transfer at scale with high- quality junt Scaled LLMs at IB, 3B. 7B, and 13B sizes.
C- CA-AB- 3C
(18)
NVIDIA'S H100 GPUs support FP8 GEMMs through the Mablal))Matmul))) operation (NVIDIA Corporation, 2024).
To maximize training speed and mirror TransiormerEngine NVIDIA (2023), fuse clipping to the FP8 range, casting to FP8, and transposing into a single Triton (Tillet et al., 2019) kernel. A transpose is necessary because H100s only support one layout ("TN") with FP8, but the forward and backward passes use different layouts (thanks to using W vs WT).
8

---

pnit Scaling: Simple and Scalable FP8 LLM Training
FP8 Training Speedup over BF16
ity, and Resource Optimization (WANT 2024),
Speeclup over BF16
Transoomeenggreecc (TE)
1.32
pet Scoling (15)) ours
1.53
2024. URL https: /wwpnreeeeeeeeeeee//e//eee net/ forum?
id= 4NKK221n5.
.22
1.24
1.28
1.10
Casella, G. and Berger, R. L. Statistical Inference. Duxbury, Pacific Grove, CA, 2nd edition, 2002.
ISBN 978-0-534-24312-. URL https://pages stat
18
38
136
Model Size
Sarger_Statsttion Intaala__nfeeeeeee pdf.
Figure 8. Training in FP8 with μS achieves state-of-the-ar ef- ficiency. FP8 training with umit Scaling provides 25-33% higher throughput than BF16 training and 1-6% higher throughput than FP8 training with TransformerEngine (TE), over IB, 3B, 7B, and 13B model sizes. Models are configured as specified in Table and benchmarked on NVIDIA H100 GPUs. Static scaling, operator fusion, and simplifications to Unit Scaling make this efficiency possible.
Chen, X., Liang, C., Huang. D., Real, E., Wang, K., Pham, H., Dong, X., Luong, T., Hsieh, C.-J., Y., Le, Q. V. Symbolic discovery of optimization algorithms. In Thirty- seventh Conference on Neural Information Processing Systems, 2023, URL https: opwprevieneeeeeeeeeeeeeeee hnt.sieiwwee.eeeeee............e forum? dd=ne wew.dddeeeeeeeeeee 6zeqLFCZ.
/
Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. LLM.int80: 8-bit matrix multiplication for transformers at scale. In Oh, A. H., Agarwal, A.. Belgrave, D., and Cho, (eds.), Advances Neural Information Processing Systems, 2022. URL https: openreview net./ forum ferrm?id=dxicooe dw
Impact Statement
This paper introduces unit Scaling (uS), method designed to enhance the efficiency of Large Language Model (LLM) training through scalable FP8 computation and straightfor- ward hyperparameter transfer. The advancements provided by IS could reduce both the computational and environ- mental costs associated with training large-scale models, potentially democratizing access high-peeforrance ma- chine learning by lowering resource requirements. While this work's 's primary goal is advancing training efficiency, we acknowledge that, as with all machine learning technologies, continued attention to ethical considcrations and societal implications remains important.
Dohmann, J. Blazingly fast LLM evaluation for in-context learning, 2 2023. URL https: wwww dat. ceatabricks. com/blog/ Laa.llgg/nn-eeaaltiiooo------- for-iicc..
Dong, J., Feng, B., Guessous, D., Liang, Y., and He, H. Flex attention: programming model for gener- ating optimized attention kernels, 2024. URL https: //arxiv. org/ abs 22412.05493
Khudia, D., Huang, J., Basu, P., Deng, S., Liu, H., Park, J., and Smelyanskiy, M. Fbgemm: Enabling high- performance low-precision deep learning inference, 2021. URL https: //tp:i/aaiiiviv//a orgg abs 2511..0555...
References
Anonymous. Scaling FP8 training to trillion-token LLMs.
In Submitted to The Thirteenth International Conference Learning Representations, 2024. URL https: // der review.
Kingma, D. P. and Ba, J. Adam: A method for stochastic op- timization, 2017. URL ht http:///xxxxxx.../.aaa/.o 1412.6980.
un-
Lingle, L. large-scale exploration of μ-transfer, 2024. URL https: http//aavvvv///// abs / 40404 05728.
Barton, T. ation Gauntlet. 2024. www..databricke. com/blog/ callbraciigg-oooal--eeaaaau
Calibrating the Mosaic evalu-
URL https:
Liu, Z., Hu, H., Lin, Y., Yao, Z., Xie, Z., Wei, Y., Ning, J., Cao, Y., Zhang, Z., Dong, L., Wei, F., and Guo, B. Swin transformer v2: Scaling up capacity and resolution. In CVPR, pp. 11999-12009, 2022. URL https: 01170.
Blake, C., Orr, D., and Luschi, C. Unit scaling: Out-of-
the-box low-precision training. 2023. URL https: //arxiv.orge wwtp://rg/////oeeeoooooooosss//s /2303 11257.
org/10. 10. 1101 1119/
Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., and Wu, H. Mixed precision training. In International Conference on Learning Representations, 2018. URL https: https::/openzee////// /ppprevveneeeeeeeeeeeeeeeeeeeeee net/ forum? id=r1gs9ugRa 9JgRz.
Blake, C., Eichenberg, C., Dean, J., Balles, I.,, Prince, L. Y., Deiseroth, B., Cruz-Salinas, A. F., Luschi, C., Weinbach, S., and Orr, D. u-up: The unit-scaled maximal update parametrization. In 2nd Workshop on Advancing Neural Network Training: Computational Efficiency, Scalabil-

---

pnit Scaling: Simple and Scalable FP8 LLM Training
Micikevicius, P., Stosic, D., Burgess, N., Cornea, M., Dubey, P., Grisenthwaite, R., Ha, S., Heinecke, A., Judd, P.,
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polo- sukhin, I. Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL Ittptttss::/ neurips, ccw.ippecc/ppppr.ccccc http://a.....aeee...eeeee/////// 3 (55 pdf.
Kamalu, J., Mellempudi, N., Oberman, S., Shoeybi, M., Siu, M., and Wu, H. Fp8 formats for deep learning. 2022.
URL htsps::/p/..........m..tpm..t / a aarxiii. v. org/ abs/ 9.09...00. 0.4.3..
Mirzadeh, S. I., Alizadeh- Vahid, K., Mehta, S., del Mundo, C. C., Tuzel, O., Samei, G., Rastegari, M., and Fara- jtabar, M. ReLU strikes back: Exploiting activation sparsity in large language models. In The Twelfth In- ternational Conference on Learning Representations, 2024. URL https://openrer//// wttps/vwwwweee/fouuuu.... id=osoWxYBo22.
Wang, X. and Aitchison, L. How to set AdamW's weight decay as you scale model and dataset size, 2024. URL https: wwxiv.vvoo .org/ aaa//b///// 13698.
MosaicML. Composer. https://github.com mosaicml mosaicm//composse /composer///eooccoocooooocoeeeet 2021.
Wortsman, M., Liu, P. J., Xiao. L., Everett, K. E., Alemi, A. A., Adlam, B., Co-Reyes, J. D., Gur, I., Kumar, A., Novak, R., Pennington, J., Sohl-Dickstein, J., Xu, K., Lee, J., Gilmer, J., and Kornblith, S. Small-scale prox- ies for large-scale transformer training instabilities. In The Twelfth International Conference on Learning Repre- sentations, 2024. URL https: wwwpeeeeee..ne.eeeeeeeeeeeeeeeee forum? (tt::ddddwww00000000mmmmm
MosaicML. LLM Foundry, <htttp::/ com/
mosaicml 1aamn//m1mmmmouudm /1mmm//odddo/d////22222222
MosaicML. Streaming. <https:// com/ coosaicm//stremiingg
NVIDIA. Asynchronous mullinly-and-accumlltt instruction:
wgmna.mma, UL. URL.
Yang. G., Hu, E. J.. Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. Tuning large neural networks via zero-shot hyperparameter transfer. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https: openreview net// Forum?id=Bx6
NVIDIA. TransformerEnggine, 2023. URL https:// /// wwt.ub.c.m./W/aa... Com./NVIDDA//Taasffooooetttt
NVIDIA Corporation. cuBLAS: cublastaMaumul(.
NVIDIA, 2024. URL httphh::///css.e/eeeees wwwiiiiii......... com/ ouda/cublas hatallass/cublasssttmmmt.
OLMo, T., Walsh, P., Soldaini, L., Groeneveld, D., Lo, K., Arora, S., Bhagia, A., Gu, Y., Huang, S., Jordan, M., al. 2 olmo 2 furious. arXiv preprint arXiv:2501.006 00656,,000666,,.00011 2024.
Yang, G., Simon, J. B., and Bernstein, J. A spectral condition for feature learning. 2024a. URL https:
1/arxiv.ore abs 2310.17811111 17813.
Sun, X., Choi, J., Chen, C.-Y,, Wang, N., Venkatara- mani S., Srinivasan, V., Cui, X., Zhang. W. and Gopalakrishnan, K. Hybrid 8-bit floating point (hfp8) training and inference for deep neural networks. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F,, Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL htpptttt///proeeeingg. neuri heurips.cc//apeeeeee http:///ecccceppeeeeeta////////r 355o91849978897733222222cco Paper. pdf.
Yang, G., Yu, D., Zhu, C., and Hayou, S. Tensor programs VI: Feature learning in infinite depth neural networks. In The Twelfth International Conference on Learning Repre- sentations, 2024b. URL https: / openreview. net/ forum?id- 1eddd11ppDppww....
Tillet, P. Kung, H. T., and Cox, D. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop Machine Learning and Progrumming Languages, MAPL 2019, pp. 10-19, New York, NY, USA. 2019. Association for Computing Machinery. 02.1445/311508.3329999 URL https://////..... /ww.ooooooooooooooooooooo org/ 10.1145/331 3329973.
ISBN 9781450367196. 196.
10

---

juit Scaling: Simple and Scalable FP8 LLM Training
A. Appendix
A.I. Why these modifications?
Table I contains number of modifications to standard bfl6 training setups. Where did these come from? Are they simply a result of trying ideas until something worked? Or are they the result of more principled analysis and ablations?
While we do explain the basis for each modification over the course of the main text, this section summarizes how arrived at each of them. We can group the origins of these changes into three categories: simple math, adhering to prior art, and ablation experiments.
A.1.I. SIMPLE MATH
Recall that, in order to ensure stable training and consistent hyperparameter meanings, we wish to ensure that all weight and activation tensors have unit variance. Enforcing unit variance difficult because the weights are constantly being modified throughout training. To enforce exact unit variance everywhere would require significant overhead in the form of added normalization operations. We therefore relax the constraint the following:
1. Each residual branch must have exactly unit variance
2. Weight tensors must have unit variance at initialization
3. Linear layer outputs have unit variance at initialization, assuming the inputs are iid with unit variance.
4. Weight updates should attempt to preserve the weight and activation variances to the extent that this is possible without significant overhead.
The last three requirements mirror Blake et al. (2023) while the first is stronger.
Our core modifications follow immediately from these requirements and a of math.
Unit variance initialization, linear layer scaling factors. Suppose initialize our weights with unit variance to achieve requirement (2). Given iid standard normal input elements, outputs will be X2 random variables with & degrees of freedom, where k is the contraction dimension. This has mean and variance of fan in and variance of 2 fan in, which are nowhere near and so violate requirement (3). The typical solution to this is scaling down the initialization by a factor of vfan_in in, but this violates requirement (2). As observed in (Blake et al., 2023), we can reconcile both by scaling down the outputs by / ffnn_nn in al runtime as part of the GEMM call. This one extra multiply per output element essentially free, and in fact fused into instructions such as the NVIDIA Hopper architecture's wgmma (NVIDIA). See Blake et al. (2023) for further discussion.
Learning rate scaling. Recall from (Yang al., 2021) and Section 2.3 that one scale weight initialization variance, learning rate, and linear layer output arbitrarily as long as all three are scaled according to precise relationship. Since we have fixed the weight initialization variance to 1 and the output scaling fan_in 1 our learning rate scale of fan_in is uniquely determined. Further, when changing fan in from doase to dnew, this implies scaling the learning rate by
viduo
V
1.1.2. ADHERING TO BEST PRACTICES.
Some aspects of our training recipe are crucial but already common (though not universal) practices. These include:
Weight decay (A) scaling, Recall that decoupled weight decay amounts to multiplying weights by a constant 1 X, <= during each update. This operation already has the same semantics across model widths.
FP8 hidden layers. Using e4m3 weights and activations along with e5m2 gradients is common practice (NVIDIA, 2023: Micikevicius et al., 2022) Clipping instead of overflowing prevents NaN/Inf values. Keeping the first and last layers in higher precision is also common.
1.1.3. ABLATION EXPERIMENTS.
Two modifications in our recipe can be implemented in multiple ways, so we chose the details based on smaller-scale experimental results.

---

jnit Scaling: Simple and Scalable FP8 LLM Training
Table Comparing μS with other schemes μr components have commonalities differences with existing training schemes. the only one which combines scalable, complete FP8 LLM training with hyperparameter transfer; see Figure I for a comparison of features of low- pwe-priciion training methods.
uS Component
μP
Unit Scaling
u-uP
Used, but can be different in forward and backward pass.
Linear layer scaling factors
Not used
Used
Res- Res-Post-Layee Post-LayerNorm
Not used
Not used
Not used
"Fixed" residual modification
Not used
Proposed
Not used
Unit variance initialization
Not used
Used
Used
FP8 hidden layers
Not used
Used, but not at scale
Used, but restricted only some layers
Learning rate (n) scaling
Used
Not used
Used
Weight decay (A) scaling
Not used
Not used
Used
Fixed residual modification. In order satisfy our design goal of having fixed-variance residual stream, we need to combine the previous residual stream tensor and the latest residual branch output in some manner that preserves variance. As discussed in the paper, this can be done by replacing summation with weighted summation. However, we are left with a degree of freedom in setting the weighting coefficient. To keep the search space small, we consider only the two schemes from (Blake et al., 2023) and decide between them based on the experiments in Section A.3.
Res-Post-LayerNorm. As we show in Section 2.1, the variance of token representations tends to collapse later in the sequence. If closed-form correction could exactly undo this effect, we could apply such a correction and avoid modifying the architecture. However, shown in Figures 2 and 3, the pattern of variance collapse is input-dependent and deviates greatly from what iid assumptions would lead one to expect. In order to satisfy our requirement that residual streams have unit variance, we therefore must resort to blunt instrument: imposing normalization at runtime. We could normalize the residual stream itself, add a normalization op at the end of each residual branch, or move the normalization in a Pre-LN transformer from the start of the branch to the end. We decided to go with the last option because it adds no extra operations, normalizes both the residual stream token embeddings and their updates, consistent with previous work (Liu et al., 2022; OLMo et al., 2024), and worked well our ablation experiments (Fig 4b).
A. 1.1.4. COMPARISON TO EXISTING SCHEMES
As supplement to Table I which enumerates the components of uS compared to standard practice (SP), Table 6 compares these components with yP, Unit Scaling, and u-uP.
A.2. Covariance of softmax numerator and denominator
In the proof for Prop. 2.1, we state that Cov[n,d] d] o Here we derive this result. Just as in Sec. 2.1, define as the output of the softmax function applied to vector of k independent elements x. The softmax function is defined as 8 softmax(x) [x);
As shown previously, we denote the vector of elements containing numerators of elements of s as n and denominators of elements of s as d, such that s = a. By the definition of covariance:
Cov(n, d] E[(n; - pn)) - -aa)]
(19)
By the definition of softmax, d, = Ek= nss and by linearity of expectation, Ha = kpmn Using this, we obtain:
Cov(n, d) Fl(n;; (m 12 + .... + m4 + +... n -kpn)]
(20)
Expanding this expression:
Cov jn, d = E((n (n)((n1 - pm) + (n2 - mnn +... + (n+ - Un) + .... + (nx - -m)))]
(21)
By linearity of expectation:
Cov(n, d] El(n mn) +
Σ
El(n, (m) - -Pa)]
(22)
12

---

pnit Scaling: Simple and Scalable FP8 LLM Training
Because elements of the softmax input x are independent, and nt et, elements of also independent. Therefore El[nn Mn)(nn [n]] 0 for j * i. Then by the definition of variance as Var[n] =l((n; [nn))) we obtain:
Cov(n, ( dJ Var[n]
(23)
A.3. Modifying Residual Connections with
To make skip connections variance- preseriig,,,,,,, we use the fixed residual modification scheme, as shown in Eq. 11, with coefficients based the hyperparameter (Blake al., 2023). To understand the relationship of the optimal residual coefficient 7* with network depth, we swept over various values of T for models of different widths (256, 512, 1024, 2048) and depths (20, 40, 60, 80, 100). In order to assess potential confounding effects between r" and 17" and X", we tuned those two hyperparameters as well. We trained each model for 10.5B tokens with a global batch size of 256 and sequence length of 4096. We define the optimal subset of models as those which had a final cross-entropy loss within 0.25% of the optimum (with loss averaged over the last 10 steps, i.e. 10.5M tokens). As shown in Fig. 9, T" (for the optimal subset of models) decreases network depth increases. Since the contribution of each residual branch exponentially decays with depth, lower 7 ensures a lower rate of decay, likely useful as networks get deeper. This relationship between and depth is consistent model width increases. In experiments, T can be coarsely swept. We use the results shown in Fig. 9, to directly choose r* for all yS model training.
Optimal Residual Coefficient (T vs. Network Depth
1.0
Width: 256
0.8
Width: 512
Width: 1024
0.6
Width: 2048
0.4
0.2
0.0
20
40
60
80
100
Network Depth
Figure 9. Optimal residual coefficient r" decreases with depth. The 3 hyperparameters of n. and À swept for models of varying widths (256, 512, 1024, 2048) and depths (20, 40, 60, 80, 100). The mean and standard error of is shown for the optimal subset of models from each hyperparameter sweep, where model is included in the optimal subset if it had final cross coos-cntropy loss within 0.25% of the sweep optimum. 7", which controls the decay rate of residual branch contributions in the residual stream, decreases as network depth increases.
A.4. Lion Optimizer and Hyperparameter Transfer
Here, we show why Lion Chen et al. (2023) is an "Adam-like optimizer, so the μP rules for hyperparameter transfer with Adam (Kingma& & Ba, 2017) are applicable to Lion as well. Because Adam and Lion are both adaptive optimizers that normalize gradients coordinatewise before updating parameters, the nonlinear tensor product matrix results obtained in Yang et al. (2021, Appendix J.1.3) apply to both optimizers. One can that Lion differs from Adam only in that it has different second moment estimate. Under both optimizers, with gradient gt, parameter 0 is updated as:
0c+1 =0 0, 8ime + (1- (1-31)g:
(24)
v
For Lion, this follows by expressing seg)(c;) as a/c2. Then, the second moment estimate s; for Adam (Eq. 25) and Lion (Eq. 26) are below.
Adam = B2 + (1 B2)9i +
(25)
13

---

jnit Scaling: Simple and Scalable FP8 LLM Training
Lion
=2 23 - 11))m.9. + (1 - 3
(26)
This justifies why Lion is an Adam-like optimizer for the purposes of hyperparameter transfer. We use Lion for its reduced memory footprint in all our experiments.
A.5. μnit Scaling Unit Scaling for larger model training
We test the unit scaling (US) and unit scaling (uS) methods at the 7B model scale with FP8 training, Figure 10 shows that unit scaling models diverge very early in training, while jnit scaling runs converge smoothly. Based on this experiment, we did not conduct final model runs at different model scales with unit scaling (IB-13B).
Unit Scaling vs unit Scaling at 7B model size 3.0
2.8
Train Loss
2.6
2.4
US
μsS
2.2
5000
10000
15000
Batch
Figure 10. Unit Scaling (US) vs pnit Scaling (uS) for 7B models. Convergence test loss curves at 7B model scale show that μS converges smoothly while US training diverges carly training.
A.6. Activation Outliers
We analyze activation distributions taken over 32,768 tokens at every 10 layers for all FP8 models trained according to Table 4, with results shown Fig. 13. These figures show the distribution activation values for attention and FFN block inputs and outputs in the final IB. 3B, 7B, and 13B FP8 models. While SP models consistently have outliers in the attention block and FFN block inputs at all model scales, Sf models do not have these outliers in block inputs. This may make μS models more easily quantizable. It is important note, however, that in SP models, the Pre-LayerNorm placement means that activations from the residual stream are first normalized before subsequent operations.
While we do not identify the exact mechanism by which these outliers arise in the residual stream in SP models, we show their absence in μm models here, with activation distributions that may more conducive to quantization. An activation distribution with fewer outliess requires fewer bits to represent it.
A.7. Activation Function Choice
The choice of activation function can have a significant impact on activation underflow when training in FP8. For example, recent work by (Anonymous, 2024) identifies outlier amplification from SwiGLU as a challenge for FP8 LLM training. Nearly all state-of-the-ar LLMs today use either SiLU GELU as their activation function, but when training in FP8, this may lead to underflow in activations during training. This is because these functions asymptotically approach zero inputs + oos We define the FP8 underflow fraction, or the fraction of elements flushed to 0 from BF16 to FP8 cast, as metric to evaluate various activation functions. As shown in Fig. 11, this can cause many activations to underflow.
To better understand how activation function choice influences FP8 underflow when training with unit scaling, we train small& 4 layer models with GELU, SiLU, and ReLU, Our findings, detailed in Fig. 12 that during unit scaled model training, the choice of activation function drastically impacts the FPS underflow rate for activation outputs. GELU greatly degrades
14

---

pnit Scaling: Simple and Scalable FP8 LLM Training
FP8 Underflow for select activation functions
Normally Distributed Input
Uniformly Distributed Input
1.0
1.0
2-14
2-14
FP8 Undertlow traction
0.8
0.8
2-32
2-12
Activaaion Iivvaiio Value]
GELU Underflow
0.6
ReLU FPB Underflow
0.6
Underflow
2-50
[SELU]
(PelU]
2-46
2.
[SALU]
FPB E4M3 abs min
2-204
-01
0.2
0.2
0.0
-25
-21
2
29
0.0
-25
-27
-21
2
Input Value
Input Value
Figure 11. Different activation functions different amounts of FP8 underflow. When casting (0, 1) Unif( 128,, 128) values from BF16 to FP8 (e4m3), GELU, SiLU, and ReLU (green) erroneously round to zero (underflow) with different probabilties. GELU and SiLU experience significant FP8 underflow because they slowly approach 0 for increasingly negative inputs. SiLU approaches O more slowly than GELU and so underflows for a wider range of inputs. ReLU simply maps all nogative values to 0, regardless of the numerical format.
FP8 Underflow in Activation Function Outputs
Low Precision
FPB Unnerrfow Fraction
Block
Block
Block
Block
Convergence Error
30a
10P
10
10
0.4
Error (%)
10-2.
10-2
10
10-2
0.2
10000 Training Step
10d00 Training Step
10000
10000 Training Step
0.0
GELU SLU RalLi
Training Step
FP8 underflow of activation function outputs for each block a 4 layer transformer model during training is shown for GELU, SiLU, and ReLUU. Low precision convergence error, defined the percent difference final cross entropy loss hetween FP8 model and its BF16 counterpart, shown in the rightmost chart. GELU and SiLU cause significant underflow over the course of training, and models trained with these activation functions have twice as much low precision convergence error as with ReLU, ReLU greatly reduces this FP8 underflow by multiple orders of magnitude.
Figure 12. Activation function choice impacts FP8 underflow and low-precision convergence
the representation of FFN down projection inputs, reaching up 30% underflow during training. SiLU causes similar degradation, but at lower rate, reaching up 7% during training. In contrast, Rel.U does not suffer from this problem, with a maximum of 0.04% FP8 underflow during training. As a result, FP8 unit scaled models trained with ReLU have smaller low-precision convergence error (defined as the percent difference between the final cross entropy loss an FP8 model and its BF16 counterpart). Based on these observations and results, ReLU minimizes FP8 underflow and low- www-precision convergence error. ReLU also has the added benefit of sparsifying activations, a property which enables significant inference-time optimizations (Mirzadeh al., 2024). However, using GELU results in models with lower final training loss. For this reason, use GELU when training all yS models. Additional investigations into activation functions more suitable for FP8 training can help mitigate underflow while also providing improved convergence.
15

---

unit Scaling: Simple and Scalable FP8 LLM Training
(a)
FP8 model activation distributions.
Distribution of Activation Values, 1B SP-FP8
Layer Index
Attention Block Inputs
10
20
Attention Block Outputs
Count
10a
10a
10
20 18
14-14-12
11
-8
-6
-2
2
4
6
8
10
12
-20
18
14-11-1-111
-8
-6
-4
-2
6
8
10
12
FFN Block Inputs
FFN Block Outputs
Count
10a
10
10a
10
-20
18
-16
16--1--12-1
-8
-2 (log2 ([valuel)
-6
-4
2
4
6
8
10
12
20
188
16
14
12
10
-8 -6
-4 -2 (logz ([value))
6
8
10
12
μ FP8 model activation distributions.
Distribution Activation Values, μsS-PP8
Layer Index
Attention Block Inputs
10
20
Attention Block Outputs
Count
10a
10a
20-16
16
12-
-1
8
-4
-2
2
8
10
12
20
18-11-14--2210
-8
2
4
10
12
FFN Block Inputs
FFN Block Outputs
Count
10a
10a
10
20
1B
14
11
8
0
10
12
20
18
16
12
11
10
12
1log2 ([value])
[log2 ([vauee))
(c) 3B SP FPS model activation distributions.
Distribution of Activation Values, SP-FP8
Layer Index
Attention Block Inputs
10
20
30
Attention Block Outputs
Count
10
10a
10
10"
20
14-12
-10
88
-2 0
10
12
-20
18
11-11-1--2221
-8 46
2
10
12
FFN Block Inputs
FFN Block Outputs
Count
10a
10a
10a
20
-0-18-10
12
-10
0
10
12
10"
220
18
16
-8
2
12
llog> ([value|)
llog> (|valuel)]
(d) 3B μ FP8 model activation distributions.
Distribution of Activation Values, 3B μS-FP8
Layer Index
Attention Block Inputs
10
20
30
Attention Block Outputs
Count
10a
10a
100
20
-2 0
10
100
-20
18
-8
2
12
FFN Block Inputs
FFN Block Outputs
Count
10a
10a
100
10
20
16.16
-14-12-11111
16
2
12
(logz (|value])
[loga (|value))
16

---

pnit Scaling: Simple and Scalable FP8 LLM Training
(c) 7B SP FP8 model activation distributions.
Distribution of Activation Values, 7B SP-FP8
Layer Index
Attention Block Inputs
10
20
30
Attention Block Outputs
Count
105
10a
10
10"
20-18-16-14-12-10
-2 0
10
12
-0-1-116-14 -8 -6
2
10
12
FFN Block Inputs
FFN Block Outputs
Count
103
10a
10a
10"
20-18-16
10
0 1log> (|valuel)]
18
16
-10 8 6
2
12
llogo ([value))]
(1) 7B μf FP8 model activation distributions.
Distribution of Activation Values, 7B μS-FP8
Layer Index
Attention Block Inputs
Attention Block Outputs
Count
103
10a
10
20-
16-166-8
14
-8
-2 0
2
8
10
12
-20
16-1-614414
12
10
-8
-6
-2
2
4
10
12
FFN Block Inputs
FFN Block Outputs
Count
10a
10
20-188-1
18-16-14-14
14-12-10
8
-2 llog (|value|)
6
-4
0
2
8
10
12
20
18
16-
14
-12 -10
-8
-6
-2
0
2
4
10
12
llog> (|value))
(g) 13B SP FP8 model activation distributions.
Distribution of Activation Values, 13B SP-FP8
Layer Index
Attention Block Inputs
30
Attention Block Outputs
Count
105
105
109
18-18-18
16
14-12-10
8
6 44 - 0
2
4
8
10
12
20
18--18-16--00 11--1------2-100
18
-6 4
2
2
4
6
10
12
FFN Block Inputs
FFN Block Outputs
Count
10a
10
109
20
18-
16
14-12--11111
8 5 2
2
4
8
10
12
20
18 16 12 10
4
6
10
12
1lgz ([value|]
(log2 (|value))
(h) μf model activation distributions.
Distribution of Activation Values, 13B
JuSFPP8
Layer Index
Attention Block Inputs
10
20
30
Attention Block Outputs
Count
10a
10a
20-
15-15-17
16
14
14-12-10
-8
-6
-4
2
0
2
4
8
10
12
-20
18
14
12
10
88
4
44
-2
0
2
4
6
10
12
FFN Block Inputs
FFN Block Outputs
Count
109
10a
20
18
16
14
-2
10
88
-6 4
2 0 2
8
10
12
-20
18
-6
14
12
-1
88
5
44
2 0
2
4
6
10
12
llog> ([value|]
llog> ([value]]
Figure 13. Activation distributions μS and SP models. Activation distributions for attention and FFN block inputs and outputs shown for 1B, 3B, 7B, and 13B FP8 models at every 10th layer. μS models lack the notable right tail of activation outliers block inputs that SP models suffer from. This may make them easier to quantize.
17


Reviews:
[{"review_id": "sjFHhav7aU", "reviewer": "Reviewer_F349", "summary": "This paper introduces µnit Scaling (µS), a method for efficient FP8 training of large language models without requiring dynamic scaling factors or extensive hyperparameter tuning. µS builds on Unit Scaling to maintain unit variance in weights, activations, and gradients, ensuring stable low-precision training. It enables hyperparameter transfer across model sizes and eliminates the need for mixed-precision layers, allowing all hidden linear layers to compute in FP8. The method achieves training speedups of up to 33% while maintaining quality comparable to higher-precision baselines.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer’s helpful feedback and questions.\n\n## Novelty of µS\n> The novelty is limited. The proposed µnit Scaling (µS) scheme combines previously published µP and Unit Scaling.\n\nWhile µS does build on ideas from both of these methods, µS involves modifications that are not present in either and obtains desirable properties that neither of these methods achieve (c.f. Figure 1: https://imgur.com/a/v1oA2cH). In particular, µS achieves hparam transfer with fewer extra hparams than µP and has better FP8 numerics at scale than Unit Scaling.\n\nWe elaborate on specific differences below.\n\n### µP\n\nµP enables hparam transfer from smaller to larger models. **However, µP suffers from numerical instabilities even with 16-bit formats**. As described in Section 7.4 of their paper, numerical issues caused frequent divergences that required them to train their GPT-3 in FP32. In contrast, µS provides hparam transfer even in FP8. µP also requires tuning many more hparams than µS (6 vs. 3, see Table 3), as well as changes such as $\\frac{1}{d}$ attention and zero-initialization for some layers, which µS does not impose. **In summary, µS is a simpler solution than µP for hparam transfer on top of enabling FP8 training.**\n\n### Unit scaling\n\nUnit Scaling facilitates low precision training by maintaining near-unit variance of weights, activations, and gradients through unit variance weight initialization and per-operation static scaling factors. However, as we show in Section 2.1, the masked self-attention operation central to LLMs has diminishing output variance over sequence position, **which Unit Scaling doesn’t address**. Our work is the first to identify this issue and uses post-branch-norm to address it. We demonstrate LLM training up to 13B parameters in FP8, while the largest model trained in the Unit Scaling work was 340 million params (BERT Large). µS also enables hparam transfer, which Unit Scaling does not. **In summary, µS fixes key numerical issues that Unit Scaling does not, enables hparam transfer, and scales FP8 training to much larger model sizes**.\n\n### u-µP\n\nConcurrent work on u-µP also builds on both µP and Unit Scaling. However, **unlike µS, u-µP cannot keep all hidden layers in FP8**, requiring “critical matmuls” in transformer blocks to stay in BF16. u-µP also requires tuning many more hparams than µS (7 vs. 3, see Table 3). Key architectural modifications such as post-branch-norm and fixed residual coefficients permit µS to scale better. **In summary, unlike u-µP, µS provides full FP8 training for large LLMs and does so with fewer hparams.**\n\n## Ablating components of µS\n> Lack of Ablation studies.\n\nPlease refer to the “Ablating components of µS” section of our response to Reviewer NZBg.\n\n## Additional baselines\n> Lack of Baselines. µP, u-µP, and Unit Scaling are strongly related techniques that should be included in the main results\n\nWe have attempted comparison to Unit Scaling (US), but on models 7B and larger, US models were unable to converge in FP8 (see this representative figure at 7B model size: https://imgur.com/a/lJAMKmU). Further, our comparison of pre-norm vs. post-norm in FP8 showed that pre-norm (which US models use) had worse convergence (see here: https://imgur.com/a/qhq9CTP). In light of these findings, and because SP was more important to compare with, we allocated resources towards SP baselines at large scales instead of US.\n\nWe did not compare to a µP baseline because of the difficulty of training µP models in 16 bit formats, let alone with FP8 (see “µP” subsection above). Being forced to train in FP32 means that 1) we can’t obtain an apples-to-apples comparison, and 2) collecting results for this baseline would be extremely expensive.\n\nWe did not compare to a u-µP baseline since this work was done concurrently with our own. Given more time/compute resources, we agree that this comparison would be interesting, but unfortunately is not feasible to complete. Similar to our work, the u-µP work also only uses SP as a baseline.\n\nWhile assessing these additional baselines could be valuable, we note that **our ***demonstrated*** results are stronger than any existing method’s ***claimed*** results**. Even if we were to tune the hparams of these methods and achieve good hparam transfer and FP8 training quality, these methods would not surpass our µS results. This is because the focus of all methods, including ours, is on *qualitative* rather than quantitative properties–existing methods cannot fully preserve BF16 accuracy *even more* or keep the optimal hparams *more* unchanged with width. µS already matches BF16 quality with FP8 and achieves near-perfect hparam stability, leaving room for only tiny improvements on these axes. Such small improvements would not outweigh our method’s **reduced hyperparameter count**, **faster training**, and **alignment between training and inference precisions**.", "claims_and_evidence": "The claims made in this submission are generally supported by Fig 2, Fig 7, Table 5.", "methods_and_evaluation_criteria": "The evaluation criteria make sense for the problem and are aligned with previous papers.", "theoretical_claims": "Did not completely check the correctness of the proofs.", "experimental_designs_or_analyses": "The proposed µnit Scaling (µS) combines µP and Unit Scaling. However, the main results did not include these closely related baselines: µP, u-µP, and Unit Scaling.", "relation_to_broader_scientific_literature": "https://arxiv.org/pdf/2407.17465", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n\n(1) It is interesting to see the theoretical analysis of the attention output variance.\n\n(2) This paper is well-organized and easy to follow.\n\nWeakness:\n\n(1) The novelty is limited.  The proposed µnit Scaling (µS) scheme combines previously published µP and Unit Scaling.\n\n(2) Lack of Baselines.  µP, u-µP, and Unit Scaling are strongly related techniques that should be included in the main results like Table 5 and Figure 7.\n\n(3) Lack of Ablation studies. The proposed µnit Scaling (µS) scheme contains several modifications as shown in Table 1. However, there are no ablations to show each contribution of these modifications.", "other_comments_or_suggestions": "None", "questions_for_authors": "(1)  Could you include µP, u-µP, and Unit Scaling baselines in your main results?\n\n(2) Could you include ablation studies for the proposed µnit Scaling (µS) ?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IjRvTkpMrl", "reviewer": "Reviewer_NZBg", "summary": "This paper presents µnit Scaling(µS), a straightforward and scalable FP8 training method. It addresses the root causes of numerical instability in conventional transformer blocks and proposes effective solutions to mitigate these issues. µS approach incorporates Square-root Softmax Attention and Post-Branch-LayerNorm within transformer blocks, along with zero-shot hyperparameter transfer, enabling hyperparameters tuned on smaller networks to be applied directly to larger models. µS demonstrates stable training for LLMs up to 13B parameters in FP8 precision without dynamic scaling and achieves a 25-33% throughput increase compared to NVIDIA’s Transformer Engine.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s comments on our work, and are glad that the ideas we presented are clear.\n\n## Ablating components of µS\n> Are there ablation experiments on the role of different components of µS?\n\nMost interventions in µS are **uniquely determined by simple math and the design goals** of:\n1. Enforcing near-unit variance in all tensors\n2. With negligible overhead\n3. While enabling hparam transfer\n\nIn the few cases where there are degrees of freedom, we either adhere to common best practices or perform ablations.\n\nWe’ve added a new section at the start of the Appendix that spells out the origins of each component and we are modifying the text to make this clearer. An abbreviated version of this section is as follows:\n- **Unit variance init** - Necessary to ensure that weight tensors have unit variance.\n- **Linear layer static scales** - Since we don’t scale down the weights by $\\frac{1}{\\sqrt{\\text{fan in}}}$, we have to scale linear layer outputs by this factor to maintain unit variance outputs.\n- **Learning rate scaling** - Based on µP, this is the unique way of scaling LR with model width that enables hparam transfer with the above weight initialization and static scaling factors.\n- **Weight decay scaling** - We adhere to the best practice of using decoupled weight decay from prior work, since coupled weight decay complicates hparam transfer. This is noted so our µS “recipe” is fully self-contained.\n- **FP8 hidden layers** - Standard practices for FP8 training, again included for completeness.\n- **Post-Branch-Norm** - As shown in Section 2.1, masked self-attention has diminishing variance with sequence position. This cannot be corrected with per-tensor scaling factors. Norm placement *can* correct this, so this degree of freedom has ablation results in Fig. 4b.\n- **“Fixed” residual modification** - Maintaining unit variance in the residual stream requires a weighted sum, but how to weight the branch and stream is a degree of freedom. We use the same weights across all residuals based on the results in Fig. 5.\n\nOmitting one or more of these modifications would prevent unit variance and/or hparam transfer, as shown by the mathematical analysis. While it may be interesting to explore scenarios where training doesn’t completely fail without some components, these results would not be useful enough to warrant inclusion–especially since these components are easy to implement and have minimal overhead.\n\nAnother subtlety is that variance preservation and hparam transfer are AND functions. Unless the *entire* residual block preserves variance, the model will not. Similarly, weight init *and* static scales *and* learning rate must all work in conjunction to get hparam transfer. **The above modifications are not independent, additive tweaks–they are a minimal set to achieve the desired properties.**\n\nWhile the existing results already provide ample justification for the components of µS, to be completely sure that we address the reviewers’ comments, **we performed further experiments ablating norm placement and residual modification choices.**\n- Norm placement: Pre- vs. Post-branch norm with µS models (see https://imgur.com/a/qhq9CTP)\n    - Post-branch-norm converges better than pre-norm with µS in FP8. Supports the theoretical motivations for post-branch-norm from Section 2.1 (i.e., maintaining residual stream variance).\n    - Supplements our existing norm ablation in Fig. 4b.\n- Residual modification: Fixed, running-mean and standard residual modification (see https://imgur.com/a/eByqlYO)\n    - µS models that use standard residuals (i.e., no coefficients) do not properly converge. Supports the theoretical motivation for fixed residual modification from Section 2.2 (i.e., maintaining stream variance).\n    - Supplements our existing residuals ablation in Fig. 5.\n\n**In the final paper, we will compile all of these ablation results into a single subsection to clarify the contribution of the µS components. We hope this addresses the reviewers’ feedback about ablation studies.**\n\n## Other modalities, architectures, and tasks\n> Extending experiments to diverse architectures, tasks, and hardware platforms would provide stronger evidence of µnit Scaling’s robustness and versatility.\n\n> Limited benchmark diversity restricts insights into µS’s broader applicability.\n\n> Can µS be effective in models from other modalities or for different types of tasks?\n\nWe agree that applying our ideas to more modalities and architectures would be interesting. However, we believe improving LLM training is already enough scope to have a large, real-world impact. Further, nothing in µS is specific to H100s – µS is useful whenever properties listed in Fig. 1 are desirable.\n\n## Article layout improvements\n> The layout of the article increases the difficulty of reading.\n\nIf the reviewer could please elaborate on what we can improve, we are glad to address it.", "claims_and_evidence": "Most of the key claims in the paper are well-supported by experimental evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-suited to the core problem of FP8 training for LLMs. The design effectively addresses stability issues, improves efficiency, and reduces tuning costs. However, extending experiments to diverse architectures, tasks, and hardware platforms would provide stronger evidence of µnit Scaling’s robustness and versatility.", "theoretical_claims": "The paper introduces a modification to the conventional softmax attention mechanism by applying a square root to the softmax scores. The provided proof establishing the relationship between sequence length and variance is sound, and the advantages of using Square-root softmax attention are effectively demonstrated through both theoretical analysis and experimental results.", "experimental_designs_or_analyses": "The experimental design effectively demonstrates µnit Scaling’s strengths in FP8 training stability, throughput, and hyperparameter transfer. However, limited benchmark diversity restricts insights into µS’s broader applicability.", "relation_to_broader_scientific_literature": "Square-root Softmax Attention:\n\nConventional softmax attention mechanisms in Transformers are known to amplify large activations, which can destabilize FP8 precision due to overflows in matrix multiplication. The Square-root Softmax Attention in this paper offers a lightweight yet effective stability solution for FP8.\n\nPost-Branch-LayerNorm:\n\nConventional Transformer architectures place LayerNorm before the residual connection. µS adapts post branch layernorm ideas specifically for FP8 precision, ensuring better variance control in deep LLMs.\n\nResidual Modification Schemes:\n\nResearch on Deep Residual Networks has demonstrated that scaling residual branches can mitigate gradient explosion or vanishing gradients. This paper introduces fixed residual modification that stabilizes variance across deep Transformer layers.", "essential_references_not_discussed": "Most of the related works are cited and discussed in the paper.", "other_strengths_and_weaknesses": "Strengths: \n\n1.\tThe paper is well-written with minimal typos.\n\n2.\tFigures int the paper are well designed and easy to understand.\n\n3.\tThe introduction and demonstration of the method are sufficient.\n\nWeaknesses:\n\n1.\tThe layout of the article increases the difficulty of reading.\n\n2.\tLack of experimental evidence demonstrating the necessity of different components of the method.", "other_comments_or_suggestions": "Summarizing the symbols and variable names used in the paper in the appendix can improve the readability of the article.", "questions_for_authors": "1.\tAre there ablation experiments on the role of different components of µS?\n\n2.\tCan µS be effective in models from other modalities or for different types of tasks?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "n/a"}}, {"review_id": "e2AuSAf88c", "reviewer": "Reviewer_GMxA", "summary": "The paper introduces a new 8-bit training method for LLMs, by keeping all tensors close to unit variance.  The method has fewer hyper-parameters than earlier methods, and is more straightforward to implement.  It also enables more accurate hyperparameter transfer from small models to large ones.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback on our work and are glad that they find the method’s theoretical basis and empirical results rigorous and compelling.\n\n## Training duration\n> Only thing to make this claim stronger would be to train for more steps. Training steps are a bit short, but this is understandable and probably due to compute limitations.\n\nWe agree with the reviewer that training for a longer duration would have been nice to make our claims even stronger. As the reviewer has anticipated, we had only a limited amount of compute and had to design our 1B, 3B, 7B, and 13B training runs accordingly. We are excited to see future work applying µS at even larger scales.\n\n## Changes vs. SP and µP and their roles\n> For instance table 1 is informative to show changes wrt standard transformer, a similar table would be very useful to compare against \\mu P and SP as well.\n\nWe agree that denoting differences between µS and other training methods is useful. Table 1 in the paper shows changes that µS makes relative to the standard transformer (i.e., standard parametrization, or SP), which is the main comparison we aim to make. As a supplement to Table 1, we have prepared another table that enumerates differences between µS and other training methods (see table: https://imgur.com/a/f4dglps), which we plan to include in the appendix of the final paper.\n\n> What modifications induce the stability observed over SP? \n\nNumerical stability is a result of variance preservation in µS, making tensors better representable with low-precision formats. The subtle point we would like to emphasize is that variance preservation is an AND function; many components of µS all work in conjunction to achieve it. Unless the *entire* residual block preserves variance, the model doesn’t preserve variance. The entire model needs to be variance preserving in order to achieve better numerical stability over SP, and this is a conjunction of several modifications: linear layer scaling factors, post-branch-norm, fixed residual modification, and unit variance initialization.\n\n> Which method allows to reduce hyperparameters over \\mu P?\n\nWe need to tune fewer hyperparameters with µnit Scaling than µP because:\n- **Improved stability leads to more simplicity with µS**\n    - As discussed earlier, improved stability is the result of several components of µS together. Because training is more stable, we do not need to tune more hyperparameters to achieve reasonable performance with µS. Of course, if we added more hyperparameters and multipliers to tune it is reasonable to expect marginally better performance, but this is simply not necessary with the µS approach.\n- **Design choices eliminate extraneous hyperparameters**\n    - Enforcing near-unit variance in µS models by design eliminates hyperparameters related to initialization and individual layers. The examples below contrast hyperparameters tuned with µP (see Section F.4 of Tensor Programs V) with µS.\n        - µP tunes the weight initialization scale. µS initializes weights from $\\mathcal{N}(0,1)$ by design.\n        - µP tunes the attention temperature. µS maintains queries and keys with near-unit variance by design and does not require this.\n        - µP tunes the embedding multiplier. µS keeps activations near-unit variance by design and avoids this.\n\nAs a result, **training performant LLMs with µS only requires three hyperparameters**: learning rate, weight decay, and the residual coefficient.", "claims_and_evidence": "The main claims of the paper are well supported:\n- Claim:  Proposed method achieves comparable accuracy to bf16 training.\nAuthors compare against equivalent models trained with bf16 at several scales (up to 13b) and show training loss is within ~1%.  Only thing to make this claim stronger would be to train for more steps.  Training steps are a bit short, but this is understandable and probably due to compute limitations.\n- Claim: Proposed method is more practical than previous fp8 training methods\nThe method has only 3 hyperparameters which compare favorably to comparable methods, and these parameters also transfer better from small to large models, making the method more practical.  It also looks easier to get stable training with this method.  I'm conviced.", "methods_and_evaluation_criteria": "Metrics are training loss and standard LLM benchmarks, which is the right set of metrics here.  Training steps could be longer, if more compute is available, to make a stronger case.", "theoretical_claims": "Theoretical claims look correct.", "experimental_designs_or_analyses": "Experiment design is mostly straightforward as it should be.  The setup and presentation for the hyperparameter transfer results are rigorous and convincing.", "relation_to_broader_scientific_literature": "This is the part I'm most unsure about, as I haven't followed the recent literature very closely.  Most of the methods presented (e.g. post-layer norm, unit variance initialization and modifications) have been used in prior work, and the authors cite those works where appropriate.  It is hard for me to judge, which changes introduced here result in the improvements over prior work like \\mu P.  For instance table 1 is informative to show changes wrt standard transformer, a similar table would be very useful to compare against \\mu P and SP as well.  \n\nFor example what modifications induce the stability observed over SP?  Which method allows to reduce hyperparameters over \\mu P", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "What modifications induce the stability observed over SP?\nWhich method allows to reduce hyperparameters over \\mu P?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Reviewers generally agree that this is paper represents progress towards all-FP8 training of large models.
In comparison to existing work (mu-P, u-mu-P, unit scaling), the paper:
- maintains more weights in FP8
- uses post-branch-norm to address the sequence-dependence of output variance in masked attention
- provides hyperparameter transfer over more hyperparameters

Concerns about comparisons to existing work (some essentially concurrent) are well covered in the rebuttal.

A remaining concern is the small number of training steps - given the number of runs represented in this paper, 
it would be extremely reasonable to ask for a realistic number of steps on at least one or two large scale runs.
Without seeing any indication of behaviour on longer runs, all of the above advantages cannot simply be assumed 
to apply.
Decision: Accept (poster)