u-$\mu$P: The Unit-Scaled Maximal Update Parametrization
Authors: Charlie Blake
Constantin Eichenberg
Josef Dean
Lukas Balles
Luke Yuri Prince
Björn Deiseroth
Andres Felipe Cruz-Salinas
Carlo Luschi
Samuel Weinbach
Douglas Orr
Venue: ICLR 2025 Spotlight (2025)
Keywords: maximal update parametrization
learning dynamics
hyperparameter transfer
efficiency
training
stability
scaling
numerics
fp8
low precision
Primary area: foundation or frontier models, including LLMs
TL;DR: We improve µP by combining it with Unit Scaling, leading to a simpler scheme with better default hyperparameters, lower loss, more efficient sweeping and simple FP8 training.

Abstract: The Maximal Update Parametrization ($\mu$P) aims to make the optimal hyperparameters (HPs) of a model independent of its size, allowing them to be swept using a cheap proxy model rather than the full-size target model. We present a new scheme, u-$\mu$P, which improves upon $\mu$P by combining it with Unit Scaling, a method for designing models that makes them easy to train in low-precision. The two techniques have a natural affinity: $\mu$P ensures that the scale of activations is independent of model size, and Unit Scaling ensures that activations, weights and gradients begin training with a scale of one. This synthesis opens the door to a simpler scheme, whose default values are near-optimal. This in turn facilitates a more efficient sweeping strategy, with u-$\mu$P models reaching a lower loss than comparable $\mu$P models and working out-of-the-box in FP8.

Published as a conference paper at ICLR 2025
U-uP: THE UNIT-SCALED PARAMETRIZATION
MAXIMAL
UPDATE
Charlie Blake" Graphcore
Constantin Eichenberg* Aleph Alpha
Josef Dean Graphcore
Lukas Balles Aleph Alpha
Luke Y. Prince Graphcore
Björn Deiseroth Aleph Alpha
Andres Felipe' Cruz-Salinas Cohere
Carlo Luschi Graphcore
Samuel Weinbach
Douglas Orr
Aleph Alpha
Graphcore
ABSTRACT
The Maximal Update Parametrization (uP) aims to make the optimal hyperpa- rameters (HPs) of model independent of its size, allowing them be swept using cheap proxy model rather than the full-size target model. We present new scheme, u-uP, which improves upon jP by combining with Unit Scaling, a method for designing models that makes them easy train in low- low-precision. The two techniques have a natural affinity: μP ensures that the scale activations is independent of model size, and Unit Scaling ensures that activations, weights and gradients begin training with a scale of one. This synthesis opens the door to a simpler scheme, whose default values are near-optimal. This turn facilitates a more efficient sweeping strategy, with u-uP models reaching a loss that is equal to lower than comparable uP models and working tue-of-the-box in FP8.
(a) More Efficient
Sweeps
(b) Better HP Transfer
Simple FP8 Training
3.6
3.6
10
Best Validation Loss
3.5
J.P
3.4
3.2
pP
S00080
pil ase
3.3
3.0
trp 2006
3.2
2.8
pip 2000
10
10a
10P
10a
2
2-1
2 2 2
Run Count
nP u-pP Learning Rate
Training Step
HP Sweep Strategy
uepa
Rundam
Widths
RP
Independent Search
128
1024
Prodision
Mults
512
4006
Combinni Mults
Figure 1: (a) Two different HP sweeping processes used for uP and u-uP proxy models. Unlike μP, u-pP admits independent (1D) search due to careful HP design. The first part of independent search is an LR sweep, which alone reaches near-optimal loss for u-uP. (b) Using the best proxy HPs from (a), we train many models at different widths and LRs. The best LR for width 256 is -optimal for 4096, showing LR transfer along with lower loss. (c) We re-train with simple un-scaled to (float cast on matmul inputs. This would fail for other models, but u-uP trains with minimal degradation.
Work done while at Aleph Alpha. Supervisory role. Correspondence charlieh graphbrrrrappcoor.ii wwwstottii..ciieeeo.eoo gegtaann.ccheeneeggeeeelehhhpppp
'Equal contribution.

---

Published as a conference paper at ICLR 2025
INTRODUCTION
The challenges of large-model training extend beyond the domain of engineering; they are also algorithmic nature. Effective approaches for training smaller models are not guaranteed to work at the multi-billion- multt-billion-paaantte parameter scale used for today's large language models (LLMs). These difficulties can be framed in terms of stability, which we consider in three forms:
1. feature learning stability, which ensures that parts of the model do not learn too fast slow relative to each other.
2. hyperparameter stability, which ensures that the optimal HPs for small models remain unchanged as the model size grows.
3. numerical stability, which ensures that foaatig--point representations during training stay within the range of a given number format.
The Maximal Update Parametrization (μP) (Yang & Hu 2021 Yang et al. 2022) targets the first two sources of instability. μP defines set of scaling rules that in principle make model's optimal HP values consistent across model sizes and ensure 'maximal feature learning' in the anfinit-wwidth limit. The practical benefits of this that models continue improve they get larger, and that practitioners can re-use a set of HP values (especially the learning rate) found for a small prexy version of their model, on a larger targer model. This is vital for modern LLM training, where the cost of sweeping over candidate HP values for the target model is prohibitive. Consequently, uP has
been adopted by several open L.LM training efforts (Dey et al. 2023a b. Liu et al. 2023 Hu et al. 2024 and there are indications of its use in state-of-the-art L.LMs
However, there exists a gap between the extensive theory underpinning μP and its effective use in practice. This relates to issues surrounding efficient HP search, HP transfer, interpretability. ease-of- use and low-precision training. Some of these problems have been observed in the literature (Yang et al. 2022; Almazrouei al. 2023 Lingle 2024); others we outline here for the first time. As result, uP does not necessarily provide the kind of simple, stable scaling for which a user might hope.
To address this, propose the Unit-Scaled Maximal Update Parametrization (u-uP). u-uP combines uP with another closely-relate rellted training innovation, Unit Scaling (Blake et al. 2023). μP ideally provides consistent training dynamics across model sizes, but says little about what those dynamics should be. Unit Scaling addresses this by proposing an ideal principle for dynamics: unit variance for all activations, weights and gradients. Unit Scaling was initially designed ensure stable numerics, but in the context of μP the principle of unit-scale brings many additional benefits. We show that provides a foundation upon which the broader range of drawbacks identified for uP can be addressed.
BACKGROUND
2.1
THE MAXIMAL UPDATE PARAMETRIZATION
Tensor Programs V Yang et al. 2022) defines parametrization as 'a rule for how to change [HPs] when the widths of neural network change'.. They show that μP is the only parametrization giving "maximal feature learning' in the limit, while standard parametrization (SP) has imbalanced learning.
One consequence of this improved stability that learning dynamics under uP are ideally independent of model-size, are optimal HPs. This facilitates method known as μ TTansfer,, which describes the process of training many smaller proxy models to evaluate candidate HP values, then using the best- best-performing ones to train larger target model.
JPP, SP, and the Neural Tangent Kernel (NTK) (Jacot et al., 2018 are all instances of prrraaetttrriationss This assumes model under training where weights are defined as:
wi ~ (O, Bly) W w; Ow (VLo, ..., VL.),
W. = Aw Wts
The GPT-4 technical report OpenAl 2023) hints at the use of μP by including Yang et al. (2022 in its references, without citing directly. The multipliers present in the Grok (xAl 2024 codehase also suggest the of μP.
2

---

Published as a conference paper at ICLR 2025
with time-step and 4 Lu. VL2) is the weight update based previous loss gradients.
parametrization such as μp is then defined by specifying how scalars Aw; Bw Cw change with model width. This can be expressed in terms of width-dependent factors aw, bw ew, such that Aw X ow. Bw xc bw, Cw x cw. The values these factors take are what characterize a particular scheme. For uP these are given in T For depth, similar result has been proved using depth-uP (Yang et al. (023b), albeit in a restricted setting. When we refer to uP in the paper we assume the depth-uP scaling rules (Table 2 'Residual' column).
A key property of the Abe-paramettizzaioo is that one can shift scales between Aw, Bw, Ow in a way that preserves learning dynamics (i.e. the activations computed during training unchanged). We term this abc-symmetry. For a fixed 0 >> O, the behavior of a network trained with Adam is invariant to changes of the kind:
Aw + Aw /0, Bw < Bw 8, Cw + Cw 0
(2)
(reproduced from Tensor Programs V, Section 1.2.1). This means that parametrizations like μP can be presented in different but equivalent ways. ABC-symmetry a key component in developing u-uP.
Transferable HPs. μP focuses on the subset of HPs whose optimal values we expect to transfer across axes such as width and depth. We term these yTransferable HPs. All uTransferable HPs function as multipliers and can be split into three kinds: aw, dW, '/ where Aw ε ow, Bw & ow, Cw x 1/w The difference between these multipliers and the ones that define parametrization that they are specified by the user, rather than being function of width. ow and 1/ are rarely introduced outside of the yP literature, but can be valuable to tune for both μP and SP models. the μP literature the term 'HPs' often implicitly refers to y'Transferable HPs. We adopt this convention here, unless specified otherwise.
Base shape. Two additional non-alTransferable HPs introduced by pP are the base-width and base-depth. This refers to a mechanism where a user specifies a particular shape for the model, where its behavior under μP and SP are the same. The uP model still scales according to the abc-rules, so for all other shapes the two models will be different.
In summary, the absolute expressions for Aw, Bw, Cw under uu are given by:
aw
bw
Cw 7/ CW
Aw aw OW
Bw < OW bwinee
CWiaa
Though base shapes are necessary for uP, they are not typically swept. Rather, they are considered a preference of the user, who may wish to retain the behavior of an existing SP model at given shape.
Choosing HPs sweep. In theory, the search space of yTransferable HPs includes ow, ow, 7w for every parameter tensor W in the model. In practice, far fewer HPs are swept, with global grouping often used for aw and Ijw, and many aw are dropped or grouped across layers.
The sets of HPs chosen for sweeps in the pP literature is explored in Appendix E. Tensor Programs uses a random search to identify the best HP values, which has become the standard approach to sweeping. The number of runs in sweep is typically in the low 100s, incurring non-negligible cost (though usually less than a single training run of the target model). This high number partly owes to dependencies between HPs (shown in Section making the search space hard to explore.
Table I:
The scaling rules defining μP. The type of weight is determined by whether fan-in & fan-out both depend on width (hidden), only fan-out does (input), or only fan-in (output). Hence fan-in is always a multiple of width here.
Weight (W) Type
ABC-multiplier
Input
Hidden
Output
parameter (aw) 1
4/tanin(W)
uP
initialization (bw)
1/ fan-in(W)
Adam LR (cw)
1
1/fan-in(W)
1
3

---

Published as a conference paper at ICLR 2025
2.2 LOW-PRECISION TRAINING
All the major potential bottlenecks of model training-compllnn compute, communication and storage see roughly linear improvements the bit-width of their number format is reduced. In modern LLM training, the compute cost of large matrix multiplications (matmuls) means that substantial gains are available if these can be done in low- wow-precision (< 32 bit) formats, which makes them one of the most promising avenues towards increased efficiency in deep learning.
Recent AI hardware offers substantial acceleration for the 8-bit FP8 E4 and E5 formats. However their reduced range means that they cannot directly represent some values generated during training. Various methods have been introduced address this, such the per-tensor dynamic re-scaling in Transformer Engine (NVIDIA 2024b). However, this comes at the cost of added complexity and potential overheads. For more in-depth treatment of low- oow-precision formats, see Appendix g
2.3
UNIT SCALING
An alternative approach low- ww-precision training is Unit Scaling (Blake et al. 2023), which also uses per-tensor scaling factors to control range, but instead finds these factors via an analysis of expected tensor statistics at initialization. These are fixed factors, calculated independently of the contents of tensor, at the beginning of training. As such, the method is easy to use and only adds the overhead of applying static scaling factors (which we show be negligible Appendix K).
These factors are chosen to ensure the unit variance of activations, weights and gradients at initial- ization. This is useful criterion as it places values around the center of floating-point formats' absolute range. This applies to all tensors, meaning every operation in the network requires a scaling factor that ensures unit-scaled outputs, assuming unit-scaled inputs. Unit Scaling does not provide mechanism for re-scaling tensors dynamically during training, but due to its ideal starting scale for gradients, activations and weights this may not be required. Empirically this is shown be true across multiple architectures, though it is not guaranteed.
We provide an example deriving the Unit Scaling rule for matmul op in Appendix E.2 resulting in the scaling factor: 1/ diaaninn We accompany this example with a full recipe for applying Unit Scaling to an arbitrary model.
3 THE UNIT-SCALED MAXIMAL UPDATE PARAMETRIZATION
In this section we show how μP can be adapted to satisfy Unit Scaling, and provide a new set of HPs thich-thanks to Unit Scaling-are more interpretable and separable than those commonly used for μP, unlocking several practical benefits. Although some features of u-up presented here are transformer specific, we stress that u-uuP can in principle be applied to wide range architectures, since both μP and Unit Scaling are very general approaches.
3.1 COMBINING μP WITH UNIT SCALING
Whereas Unit Scaling provides rules for scaling all operations, uP only does so for parametrized ones. It's these operations we need to address to arrive at unified scheme, resolving differences in the scaling rules each recommends. We begin with the expressions for the Aw. Bw, Cw scaling factors Equation (3) and substitute the μP scaling rules defined in Table [ This results full implementation of μP, which is shown in the top half of Table We set out to turn this into
valid Unit Scaling scheme, which requires unit initializations (Bw 1) and matmuls with the Unit Scaling factor we identified in Section 2.3 (Aw < /faniinn.
Our first step is to drop the ow and base-fan-in HPs entirely, and associate the aw HPs with subsequent non-lineat functions instead of weights-decisions decisions we justify in the rest this section (this results in the simplified intermediate implementation in Table IT). Our input weights now have unit initializations as desired, and a unit parameter multiplier, which is also the appropriate scaling factor (as input layers here embedding lookups, matmuls).
4

---

Published as a conference paper at ICLR 2025
Table 2: Scaling rules for uP versus u-uP, including associated HPs (assuming the extended set in
Table These rules constitute the definition of u-pi,, along with the unit-scaled ops in Appendix B
Weight Type
ABC-multiplier
Residual
Input
Hidden
Output
parameter (Aw)
Clemb
1 (or Okattn)
Qout
base-fan-in fan-in
base-depth
depth
μP
initialization (Bw)
Fiiit
Jiitt
base-fan-in fan-in
Onit
Adam LR
(Cw) 17 Nemb
17
base-fan-iu fan-in
7
base-depth
depth
parameter (Aw)
v Van-in
fan-in
depth
u-uiP
initialization (Bw)
1
1
Adam LR
(Cw) 1) Vfan-out
7) Vfan-ln
7)
depth
Residual multipliers are applied to the end each branch, rather than the output linear layers. uuuuP's HPs associated with operations, not weights, so are not included here (see Section 3.3 "To maintain unit scale we apply 1/ /fannout scaling in the backward pass (sec Appendix H
Hidden weights now have the implementation: Aw < 1, Bw Vfan-in" Cw 1) lan-in which
differs from our Unit Scaling criteria. However, using abe-symmetry (Equation (2) we can shift
scales by / ffn-in, arriving at a unit-scaled scheme: Aw <- Viau-in' Bw + 1, Ow + Vlan-in'
Finally, our output layers also have unit initialization, but parameter multiplier of Aw <- 1/fan-in. This differs from the Unit Sealing rule, but in the forward pass this is permissible as there are no subsequent matmuls in transformer. In the backward pass this mis-scaling would propagate, apply the desired <- 1/V /Tanoou factor. Using different forward and backward scales in this way is usually not allowed, but is valid for output layers due to the cut-edge rule (Appendix H).
make is to the input LR scaling rule, which show in Section 3.4 is more The final change
effective if cw <- 1 is replaced with cw +- /fan-out. With these changes made, arrive at our final u- u-up scheme, given in Table 2 Note that the scaling rules this table must be combined with the standard Unit Scaling rules for other non-matmul operations. These are covered in Appendix B
3.2 OUT-OF-THE-BR -BOX
TRAINING
When training a transformer model with u-up most tensors have stable scale during training, except a small number of critical tensors that exhibit scale growth, which can cause extreme values to go out of FP8 range. We empirically identify these to be the inputs to the attention dense projection and final FFN matmul well the weight of the decoder2 (for details see Appendix A.7
Based on our analysis, we propose the following FP8 -ixeeiprecissin scheme for u-uP transformers:
For layers using non-critical tensors, we cast the input and weight to 34M3, and the gradient with respect the output E5M2. This is done in the forward computation, well the two backward computations (for the gradients w.r.t. the weight and input).
For layers using critical tensors, matmuls are performed in BF16 (along with non mon-matmul operations). We keep optimizer states in FP32, leaving the use of FP8 to future work.
In some cases are able to deal with the critical tensors by casting them to E5M2. However, observed instabilities applying this in a large-scale setting. In small-scale settings we also empirically find that applying E4M3 instead of E5M2 for the gradients is possible, but again becomes problematic in the large-scale setting where gradients require higher dynamic range. Under our mixed- precissin scheme, approximately 70% of the matmul computations in a Llama transformer block are performed in FP8. If desired, a dynamic per-tensor scaling could still be applied to the critical tensors.
The decoder becomes negligible in terms of model flops as width and depth of the model increase, so we generally keep this operation in higher precision.
5

---

Published as a conference paper at ICLR 2025
3.3 A PRINCIPLED APPROACH TO HYPERPARAMETERS
Approaches for selecting which HPs to sweep are poorly motivated in the literature (see Appendix C.2)) Our objective in u-up is find simple, well-justified and effective alternative. To this end, propose the following ideal criteria:
1. Minimal cardinality: the of as few HPs possible.
2. Maximal expressivity: the ability to still express any model defined using the per-tensor
Ow, I/ HPs outlined in Section 2.1 (in practice, we relax this slightly).
3. Minimal interdependency: the optimal value of each HP should not depend on the value of other HPs, simplifying the search space.
4. Interpretability: there should be a clear explanation for what an HP's value 'means' in the context of the model.
The u-up HPs given in Table are designed to satisfy these criteria, to the fullest extent possible. The place- ment of these HPs in the model is given in Table
Table 3: Typical transformer HPs used under different schemes. Basic HPs in bold are most impactful and are commonly swept. Ex- tended HPs in non-bold are not always swept, often set heuristically or dropped.
Cardinality & expressivity. We arrive at our set of HPss in three steps, starting with the full aw, ow. 7/w for each weight tensor W. Firstly, we can remove any one of these HPs by permuting under abe-symmetry, such that one HP = 1. As we want our weights to begin with unit scale, we choose drop dw. Sec- ondly, we observe that several of the ow HPs com- bine linearly with other aw HPs, providing an op- portunity to re- re- p--armettizze with single HP. For instance, the scale of self-attention softmax activa- tions is proportional to the product of ow multipliers: std(Iaatta) x aw, aWx. In this instance we use a sin- gle parameter (termed (sattn-softmax) and associate it with the attention operation, rather than the weights.
SP
μP
u-pi
7 a-scheme
"
Finit Glemb Memb
Offin-act
Clattn
Gattn-softmax
Cout
Olres
base-width
aree-atta-ratio
base-depth
Oloss-softmax
We apply the same principle to all operations, unless they unary and k-homogeneous for k 2 0, in which case they propagate scale and don't 't require an HP (see Appendix G.I). This results in the set of HPs shown, with their placement in the model given in Table 7 We note that this procedure also works in other architectures than transformers, and naturally will produce other kinds of multipliers.
Thirdly, we
single global 1) and group a HPs across layers. This breaks expressivity criterion, but we argue represents the best trade-off between expressivity and cardinality. We show in Appendix A.3 that having tuned global HP and our extended HPs, the further benefits of tuning per-tensor f/ HPs (which modify the global 17) is minimal.
Interdependency. The second stage above, moving HPs from weights into operations, not only reduces the number of HPs, but also minimizes the interdependence between those that remain. We find that u-uP's optimal HP values depend less on each other than under μP (see Section
Interpretability. The combination of unit scale and reduced dependencies between HPs means that each a can be interpreted as determining some fundamental property of the model initialization. For example, the Olommax HP defines the (inverse of) the softmax's 's temperature for a unit-scaled input. We also introduce a new scaling scheme (defined in Appendix G.2.2) for residual connections, designed to give HPs independence and clear interpretation: Oress defines the contribution of the residual connections to the output scale, and ares-attu-ratio defines the relative contribution of attention versus FFN branches. Finally, we choose not to include base shape HPss in u-pi. They do not add to expressivity, lack a clear interpretation (besides aligmment to a base model at a particular shape), break the interpretation of HPs given above, and complicate implementation.
3.4 A NEW EMBEDDING LR RULE
Although theoretical transfer properties have been proved for μPP not all its HPs have had TTannsfe shown empirically. We do so for the extended yP transformer HPs in Figure 15 where observe

---

Published as a conference paper at ICLR 2025
poor transfer across width for the embedding LR multiplier Nemb. femb. This suggests that the corresponding LR scaling rule for Cemb is mis-specified.
We show in Figure that changing from the yP rule Cemb = 1 to /ttooo///aaaaeea fan-out corrects this failure of HP transfer. As result, we improved loss under u-uiP for larger model-sizes relative to up. Our adoption of this change is a key factor in the improved performance of u-up over uP in Figure We offer no theoretical justification for this change, which we leave future work.
3.5 HYPERPARAMETER SEARCH
As shown in section Section 2.1 the standard approach to HP search for μTransfer via a random sweep over all HPs simultaneously. Sweeping individual HPs separately is challenging due to the dependencies between them. In contrast, u-uuP's HPs are designed to admit such a strategy due to our interdependence criterion. Because of this, propose simpler sweeping strategy for u-uP which we term independent search (outlined in detail in Appendix (A.5). Independent search involves a sweep of the LRR followed by set of one-dimensional sweeps of the other HPs (which can be run in
parallel). The best results from the individual sweeps are combined to form the final set of HP values.
4 EXPERIMENTS
Our experiments use the Llama architecture (Touvron et al., 2023a) trained on WikiText-103 (Merity al. 2017) (except large-scale runs in (.4)) We apply best-practice LLM training techniques from the literature (see Table 5) In accordance with our analysis in Appendix C.T we remove parameters from norm layers, independent A, and avoid training on too many epochs.
4.1 QUANTIFYING HYPERPARAMETER INTERDEPENDENCE
Our principled approach to HPs (see Section 3.3 contains the requirement that their optimal values should depend minimally on the value of other HPs. We now investigate this empirically, conducting
a 2D sweep over every pair of HPs for yP and u-up (see Figures IT and 12 for pairwise results).
To derive an empirical measure of HP dependency, we introduce the notion of transfer error (see Algorithm M.. We take the best value of the transfer HP for each non-optimal value of the fixed HP, and use it with the optimal value of the fixed HP. The transfer error is the difference between the losses obtained and the minimum loss. Figure 3 shows this measure for each pair of HPs, reflecting the improvement in HP dependency transfer errors leading to large degradations, and the potential to sweep HPs in more separable way.
result of scheme. This gives u-uiP reduced risk of small
Reosll: Maas wy Geab CV V.C.). Gemb Remb "eenh "Fennh//eenl Cembbbase
3.6
Procr Opticum
Width
Proposed Transfir
aP Transfer Role
128 (RaseWidth)
3.4
Vabdatiou Loss
Orets
1024
3.2
20-88
4096
Embedding LR scale
Grub 1/ 1/vTau-out
2.8
Oytimam
Impooved
Gruh +-
Tonnies
24
2
2-12 12
2
2-
2
C
Learning Rate
Figure 2: (Left) holding the embedding LR multiplier (Cemb) constant, vs. scaling with widthh both with fixed global LR. This suggests the uP embedding LR rule (Cemb) should follow the latter scaling. (Right) we test this by sweeping the global LR under the two scaling rules. The new rule leads to lower loss large models. (Dot/cross markers represent the same runs across both graphs).
7

---

Published as a conference paper at ICLR 2025
transfer HP
Datt
Gatte
destpst
Yimb
01.
Transfer Ezzor
Phit
Com-attn-nrie
lixed
01.
Gattn
Offirec
Oudpul
Oouppu
J.P
u-jiP
Figure 3: A visualization of the dependencies between pairs of HPs under each scheme. Transfer error
u-pP
Training Steps
Batch Size
Depch
Validatio Less
3.5
16384
128
32768
250
3.2
2-2
2
2.2
27
22 Learnnng Hate
24
2-2
20
Learning
Learning
Figure 4: Learning rate transfer for u-uP over training steps, batch size and depth. See Figure (b) for transfer over width and Figure for regular uP transfer results. The default shape parameter for other panels is shown in bold. The shaded area shows the 95% confidence interval for the mean.
4.2 HYPERPARAMETER SEARCH
Our new HP scheme, designed for improved separability, enables proxy model HPs to be swept more efficiently. This shown in Figure (a). We conduct standard random search for yP and u-uP, along with the independent search outlined in Section 3.5 (and Appendix A.5).
Independent search begins with a simple LR sweep. This alone is sufficient for u-uP to reach near- optimal loss (using only runs). During this phase other HPs are fixed at 1, which for u-uP means that the inputs to operations are generally unit-scaled. Consequently, we conclude that unit scale at initialization is close to ideal scaling for effective learning here. In contrast μP still requires non-LR HPs to be swept to attain reasonable loss. The final 'combined mults' phase causes the loss to spike for μP. This is due to the HP dependencies shown in Figure 3 which mean HPs cannot be swept independently and used together, necessitating random search which can require hundreds of runs.
4.3 HYPERPARAMETER TRANSFER
We train many models and plot transfer of LR across width (Figure T (b)), steps, batch size and depth (Figure44) 4), and transfer of other HPs across width (Figure 15). Note that u-uP (building on uP) is designed to give transfer over 3 : the other axes we report for practical purposes. We find that:
The optimal LR is constant across width under u-puP There is small drift for training steps and batch size, and a larger with depth. Hence we recommend proxy models which primarily differ in width, moderately steps and batch size, and least in depth.
2. Whereas uP sees diminishing returns for larger widths, u-up continues benefit from width. We attribute this primarily to our improved embedding LR rule.
3. Non-LR HPs also have constant optima across width under u-uP. This is not true for μP,
where Nemb has poor transfer (see Section 3.4). along with dmit (see Appendix C.2
4. The optimal values found for non-LR HPs are all close 1. In practice this means that dropping these HPs entirely is potentially viable for similar models and training setups.
3 As use depth-uP this could he said about depth as well, hut as Yang et al. 2023b) show that transformers don't attain depth-transfer under depth-uP do expect strong transfer across depth.
8

---

Published as a conference paper at ICLR 2025
Table 4: O-shot benchmark results 7B scale.
Scheme
Format
MMLU
HellaSwag
OpenBook QA
PIQA
TriviaQA
WinoGr
SP
BF16
29.6
52.4
27.8
76.5
22.2
63.3
u-up
BFI6
29.0
53.4
31.6
77.1
23.4
63.7
u-uP
FP8
31.2
53.4
29.6
77.6
21.3
65.7
4.4 FP8 TRAINING
In this section we justify the simple mixed- preccisinn scheme descrihed in Section 3.2 and demonstrate that it can be used to train u-uPP models out-of-the-box.
shhows the RMS all linear layer inputs for a moderately sized trans- former. RMS captures the larger of the mean and scale of a distribution, and such is good test of whether tensor is likely to suffer over/underflow in low-precision. We observe that u-uP tensors largely have RMS starting close 1 and remaining at the end of training, supporting
Proof-of-conee Proof-of-concept.
scheme.
As an initial proof-of-concept we train a u-up model using our FP8 scheme over 8k steps, using HPs from proxy model, as shown Figure I this scale critical tensors can still be cast to FP8 using E5M2, while gradients can even use E4M3.
(c). We see only a small degradation versus FP32, and
Larger scale. Next consider more realistic training scenario. Using the same architecture, and following the steps set out in our u-uP user-guide (Appendix D), we train our target models on 300B tokens of the SlimPajama dataset (Shen al. 2023) (see Appendix A. for training details).
We begin with an independent search (Section
3.5
over our u-uP proxy model's HPs. Here we make
the following observations:
When using a relatively small proxy model (8 layers and 512 width), the HP-loss landscape is rather noisy. By doubling the width
discern optimal HP values more clearly.
2. The most important HPs are 1) and
All others can be left at the default 1.
3. The optimal values of these HPs are 7) 23.5 and Orres-attn-ratio 2-2.0 and thus differ non-trivially from the observed HPs in our smaller-scale experiments.
We then train u-uP models of approximately 1B, 3B and 7B parameters, using our FP8 mixed- precision scheme (see Section 3.2). We also train two basclines at each size: the first a BF16 version of our u-uP models, and the second is set of SP models using the weight init scheme from the Pythia model family (Biderman et al. 2023) and the LR scheme from Llama 3 (Dubey et al. 2024), scaling inversely with width and using LR of 3e-4 at 7B scale. For the SP baseline, non-independent weight decay of 0.1 and parametric RMS norm as Llama (Dubey et al. 2024). which is standard practice. The loss curves are shown in Figure 6 All FP8 runs converge and show no significant loss degradation. In comparison to SP. the u-up models have a qualitatively different training curve with higher loss for most of training that catches up latter stages, hinting at fundamentally different optimization trajectory. In terms of downstream performance, both of the u-uP 7B models are competitive with SP. particular, the scores of the FP8 model are mostly on par with the BF16 6 models (see Table
nPP
D.pP
Wetlwwloo
artinallon
wnight
grod
grod
2- RMS
2-
25
RMS
Figure 5: Per-tensor RMS / 2 n across u-uP and uP models initialization (left) and after training (right). Dashed and solid red lines show each format's min. normal and subnormal values.

---

Published as a conference paper at ICLR 2025
2.4
1B
6F
2.4
171
FP8
2.3
TB
233
723
22
Triinin
Training
2.1
19
0K
30K
10K
6K
10K
0.K
Training Shep
Trainin Step
Figure 6: Large-scale training runs. (Left) u-up BF16 u-up FP8. (Right) u-up BFI6 vs SP BF16.
5 RELATED WORK
Low- now-precision training Techniques to facilitate FP8 training include those covered in Appendix J
and more (Wang al. 2018: Mellempudi et al.) [019 Perez et al. 2023). These largely concern the
quantizing of activations, weights and gradients, though Peng et al. (2023) also explore FP8 optimizer
states and cross-device comms, which we consider interesting avenues of exploration. Quantization
made harder by the emergence of large outlier values when training at scale, with several techniques
proposed to mitigate this (Bondarenko et al. 2023 Sun et al. 20244 He et al. 2024). Combining these with u-uP is a natural extension of our work, facilitating simpler or lower-bit quantization.
Parametrizations The neural tangent kernel (NTK) (Jacot et al. 2018 and mean field parametriza- tion (MFP) (Mei et al. 2018 Bordelon Pehlevan 2023) are alternatives to pP. In terms of width exponents, u-ul is identical to MFP except for the embedding learning rate scaling, as shown in Table I Everett et al. (202)). The equivalence classes given by [vereet et al. (2024) also mirror our notion of abc-symmetry. They additionally show that under MFP gradients become sufficiently small to require re-scaling the Adam epsilon term. u-up naturally solves this problem in a similar way, via the up-scaling of our gradients to unit scale.
Compatible parametrization frameworks Apart from u-uP, several other recent efforts have also introduced frameworks for training under particular parametrizations, which have the potential be compatible with u-up. Large et al. 2024) introduces the modular norm over the weight-space with the aim to ensure stable updates that provide LR transfer, like μP. [verett et al. (2024) explore the notion of alignment between parameters and data, showing that other parametrizations with per-layer LRs can outperform standard μ.. These parametrizations could admit unit scaling under abc-symmetry, though this is outside the scope of this work.
Signal propagation Unit Scaling and yP are mainly concerned about statistical properties of single tensors. Another line of work studies how the covariance between two distinct inputs propagates through the model. For MI.ss this was analyzed in Poole et al. (2016); Schoenholz et al. (2017). revealing phase transitions in deep networks related to expressivity and trainability. For transformers, more recent work (Noci et al. 2022) studies the collapse of token representations. Coming from this angle, they derive the same 1/ /epph residual multiplier scaling that is used in u-up.
6 CONCLUSIONS
We introduce u-uP, modified and improved version of μP that satisfies Unit Scaling. Through careful analysis guided by first principles we identify an interpretable set of HPs that has minimal interdependencies and facilitates an efficient independent sweeping strategy. We show that the stability properties of yP combined with Unit Scaling enable a simple and robust FP8 mixed precision scheme that works in realistic large scale training scenario.
Limitations and future work. Some choices like the modified embedding LR rule are only justified by empirical observations, and currently lack a theoretical explanation. Additionally, neither uP nor Unit Scaling give guarantees for network quantities to be well-behaved over the course of training. In particular we would like understand the issue (or feature) of scale growth in the critical layers better and look into possible mitigations. We also believe that low-precision training techniques can be pushed further, with u-up offering ideal starting point for future optimizations.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ebtesam Almazrouci, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli, Ruxandra Co- jocaru, Mérouane Debbah, Étienne Goffinet, Daniel Hesslow, Julien Launay, Quentin Malartic, Daniele Mazzotta, Badreddine Noune, Baptiste Pannier, and Guilherme Penedo. The Falcon series of open language models. CoRR, abs/2311. 2023, doi: 1044550/ARRX.2231.6660 URL 48550/5500000 arxiv 2311 16367
Anonymous. Straight to zero: Why linearly decaying the learning rate to zero works best for LLMs. In Submitted to The Thirteenth International Conference on Learning Representations, 2024. URL httpss://opeet reveeeeeeweeeeeeeeeeeeeeeeeeeeee net / forum? BighrrlBggsM under review.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O' Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, 2023. URL https: /ttp:///aaxiv../ org. abs/ 2304. 01373
Charlie Blake, Douglas Orr, and Carlo Luschi. Unit scaling: Out-of-the-box low- low-precision training. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 2548 2576. PMLR, 2023. URL httph::tt////. ptttps://proced:: i ngs. ml wwwpppsssss//.2 /ake222 ht.
ake23a
Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Quantizable transformers: Removing outliers by helping attention heads do nothing. In Alice Oh, Tristan Nau- mann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Informa- tion Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http::httpppptteeeee nips httppssscpppppp.pppeee fi les hetTe//ppppr//22///aaaaese 2al00755tt8577777aaaaaaa00000--0
htmll
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural n Journal of Statistical Mechunics: Theory and Experiment, 2023 (11):114009, nov 2023. doi: 10. 10088/17425568////e
UtL.tttss://// doi. http:///xx..i..r//1//0./ org/10
1088/1742 5468/4665000 add1b0
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, IA, USA, November 28 December 9, 2022, 2022. URL htt.:::://ppperrmmmipe. nips. httpp/rrssiiissccccppppeeeeeeeee Aott
httpe/perrrrlllleeeaprr///// htppp//ecccppppeee.llele//e////p h2ahh/hh///h////// html]
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3.int80: 8- bit matrix multiplication for transformers at scale. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, nips httppp/eeresspppsocpcppeeeeeeeie httpappperr.ieeeeeppp////0221110 httm
2022, 2022. URL
c
Nolan Dey, Gurpreet Gosal, Zhiming Chen, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestness. Cerebras-GPT: Open compute-optimal language models trained on the cerebras wafer-scale cluster. CoRR, abs/2304.032008 2023a. doi: 10445500AAXXIV2204.02208 URL
03208
Nolan Dey, Daria Soboleva, Faisal Al-Khateeb, Bowen Yang, Ribhu Pathria, Hemant Khachane, Shaheer Muhammad, Zhiming Chen, Robert Myers, Jacob Robert Steeves, Natalia Vassilieva,
Marvin Tom, and Joel Hestness. BTLM-3B-8K 8K: 7B parameter performance in a 3B parameter model. CoRR,
2023b. doi: 1044550/ARXiV.2300.11 1568. URL hntpp
//doi.org//11 48550 wtw.wrwwww........... 2309. 11568
11

---

Published as a conference paper at ICLR 2025
Nolan Dey, Shane Bergsma, and Joel Hestness. Sparse maximal update parameterization: A holistic approach to sparse training dynamics. CoRR, abs/2405.15743, 2024. URL http:///rxiv
2ttg/bb/22455. 15743
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshom, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson. Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia- Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, and Frank Zhang et al. The llama 3 herd of models, 2024. URL https: /arxiv org bbg/abs// abs 2407 21783
Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A. Alemi, Roman Novak, Peter Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, and Jeffrey Pennington. Scaling exponents across parameterizations and optimizers, 2024. URL URL/ttps:: pecttps:/arxiv...e
2bbs/4072 05872
Dirk Groeneveld, Iz Beltagy, Evan Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magmusson, Yizhong Wang, Shane Arora, David Atkinson, Russell Authur, Khyathi Raghavi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar, Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Valentina Pyatkin, Abhilasha Ravichander, Dustin Sehwenk, Saurabh Shah, Will Smith, Emma Strubell, Nishant Subramani, Mitchell Wortsman, Pradeep Dasigi, Nathan Lambert, Kyle Richardson, Luke Zettlemoyer, Jesse Dodge, Kyle Lo, Luca Soldaini, Noah A. Smith, and Hannaneh Hajishirzi. Olmo: Accelerating the science of language models. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), ACL 2024, Bangkok, Thailand, August 16, 2024, pp. 15789-15809 Association for Computational Linguistics, 2024. doi: 10. 016666//1/02244AA1100000 18653/v1/2024.co -ong. 841
URL https:: //opi //opo/////oio org/10
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro von Werra, and Mar- tin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. CoRR. hbs/2405.18392 2024. doi: 10.48550/ARXIV.24000 18392. URL /ttp::://oooooooo
48550/arXiv. 2405. 18392
Bobby He, Lorenzo Noci, Daniele Paliotta, Imanol Schlag, and Thomas Hofmann. Understanding and minimising outlier features in neural network training. CoRR, ab2/2405. 2024. doi: 10. 48550/ARXIV.24005. 19279. URLhttps. doi. oro.oogg///.. 111.1.885000 2ttp://0..550000....222222 2405 19279
Shengding Hu, Yuge Tu, Xu Han. Chaoqun He, Ganqu Cui, Xiang Long. Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zhen Leng Thai, Kai Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng. Dahai Li, Zhiyuan Liu, and Maosong Sun. MiniCPM: Unveiling the potential of small language models with scalable training strategies. CoRR, abs/2404.06399 06395, 2024. doi: 10.4550/ARXX.2400006555
httL::t //oi. org 10. 405.0/650//aa (15555//rXiiv arxiv. 2404. 06395
IEEE Computer Society. IEEE standard for floating-pon point arithmetic. pp. 1-84, July 2019. doi:
10.1 (U09/EEEETTDD2009...7..
Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural tangent kernel: Convergence and generalization in neural networks. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi. and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NewrIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 8580 8589, 2018 URL 5a4be fa34e (wc.be8666e6999600000 Abst
wtw.:oocceeiinnnnnn..neeeee neurips httpp//pcc.ccaaap///////////////
12

---

Published as a conference paper at ICLR 2025
Oleksii Kuchaiev, Boris Ginsburg, Igor Gitman, Vitaly Lavrukhin, Carl Case, and Paulius Mi- cikevicius. OpenSeg2Seq: Extensible toolkit for distributed and mixed precision training of sequemce-to-ssqquene models. CoRR, abs/1805.10387, 2018. URL attp://arxiv.oom
abs/1805 10387
Tim Large, Yang Liu, Minyoung Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. CoRR, 1.N 2024. doi: 10.48550/ARXIV24000. 14813. URL httppsss//b....rgg/o.ooo...
14813
Lucas D. Lingle. large-scale exploration of p-transfer. CoRR, abs/2444.057722 2024. doi: 10. 485500/ARXI.22444.55725
4800..8550/8800 2404
057288 05728
Zhengzhong Liu, Aurick Qiao, Willie Neiswanger, Hongyi Wang, Bowen Tan, Tianhua Tao, Junbo Li. Yuqi Wang, Suqi Sun, Omkar Pangarkar, Richard Fan, Yi Gu, Victor Miller, Yonghao Zhuang, Guowei Hc, Haonan Li, Fajri Koto, Liping Tang, Nikhil Ranjan, Zhiqiang Shen, Xuguang Ren, Roberto Iriondo, Cun Mu, Zhiting Hu, Mark Schulze, Preslav Nakov, Tim Baldwin, and Eric P. Xing. LLM360: Towards fully transparent open-source LLMs. CoRR, abs/2312.065500, 2023. doi: 10.48550/ARXIV22122.6553 URL
06550
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. wwt./etett/eeeeee OpenReview.net, 2019. URL BRLLhtttps////eeeee/// Bet Biw.dd-kkkkKKKii1
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665- E7671, 2018. doi: 10. 10.1073/pnas. 180557911.. URL https: www. pnas org/ hhtgg///o////////////// 10. 10.1073/pnaa 18C6579115
Naveen Mellempudi, Sudarshan Srinivasan, Dipankar Das, and Bharat Kaul. Mixed precision training with 8-bit floating point. CoRR, abs/1905. 12334, 2019, URL http:://arxi..or//abbb.
1905. 12334
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings. OpenReview.net, 2017. URL https //openrevien lew net htttt//ooumm1ddddem Byydddd
Paulius Micikevicius. Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David García, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver; BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenKeview.net, 2018. URL pettps:///oenreviers net wtt.://oooooommearrrrrr....
Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisen- thwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, Naveen Mellem- pudi, Stuart F. Oberman, Mohammad Shoeybi, Michael Y. Siu, and Hao Wu. FP8 formats for deep learning. CoRR, abs/2209.054333 2022. doi: AAAS550/ARXIV.2200005.. htt.:/550000222.2220............ 2209.05433
URL
Microsoft. Maximal update paramettization (uP) and hyperparameter transfer (μTraanffe). https h//.thub.ccom///// htt.:/o/m/miicrroootot/mm/ 2024.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand. Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phan- ishayee, and Matei Zaharia. Efficient large-scale language model training on GPU clusters using Megatron-LMM In Bronis R. de Supinski, Mary W. Hall, and Todd Gamblin (eds.). International Conference for High Performance Computing, Networking, Storage and Analysis. SC 2021, St. Louis, Missouri, USA, November 14-19, 2021, pp. 58, ACM, 2021. doi: 00.1145//458817.44762233
URL httpt::pp//.m......... / .org/ /et////..14/33588881 34768881777774999
13

---

Published as a conference paper at ICLR 2025
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aure- lien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave. K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 27198-27211. Curran Asso- ciates, Inc., 2022. URthtttss://prrceeed ngs neuri ps ccc CC wwtpppppppppppptpppttp fiwpsfffeeeeprrrrrreeee 0 4022:/11447601111156666623332223 Papeer---app----ee Conferenenfeeeeceeeoo Sor.raaeeeeec.aaaaaaooooeeeeeeee
Badreddine Noune, Philip Jones, Daniel Justus, Dominic Masters. and Carlo Luschi. 8-bit numerical formats for deep neural networks. CoRR, abs/2206.00915, 2022. doi: 10.48550/ARXIV.2200.
02915. Utthttps: /tt.:oio..gggo/e/ooooo
02915
NVIDIA. Using FP8 with transformer engine. bttps://doss.avid..
deepl http:///cccaattttee.eeeeeeeee/ra
html 2024a.
NVIDIA. Transformer http:://gihhh..ccooo///////////m
2024b.
OpenAL. GPT-4 technical report. CoRR, BB/.//22.////777////// 2023. doi: 10445550AARXI.22300.07777
URL
Houwen Peng. Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze Liu, Yifan Xiong, Ziyue Yang. Bolin Ni, Jingeheng Hu, Ruihang Li, Miaosen Zhang, Chen Li, Jia Ning, Ruizhe Wang, Zheng Zhang, Shuguang Liu, Joe Chau, Han Hu, and Peng Cheng. FP8-LM: training FP8 large language models. CoRR, abs/23 10211.0111110 18313, 2023. doi: 10.48550/ARXIV.23100188
URL
https://dol wrg/10.......................... 48550/ arxiv 2310 18313
Sergio P. Perez, Yan Zhang, James Briggs, Charlie Blake, Josh Levy-Kramer, Paul Balanca, Carlo Luschi, Stephen Barlow, and Andrew Fitzgibbon. Training and inference of large language models using 8-bit floating point. CoRR, abs/2309 2023, doi: 10.44550ARXXW22300.12222 URL
https: htt.//i..oggglg/ /00 48550 wrw.v. 2309. 17224
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein and Surya Ganguli. Exponential
expressivity in deep neural networks through transient chaos. In D. Lee, M. Sugiyama, U. Luxburg. Guyon, and R. Garnett (eds.). Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL htttps://procestet
www..iis.........ee.ee Paper.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models. In Christine Cuicchi, Irene Qualters, and William T. Kramer (eds.), Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2020, Virtual Event / Atlanta, Georgia, USA, November 9-19, 2020, pp. 20. IEEE/ACM, 2020. doi: 10. w0w.////444...200......o 1/doi.org/10/ // 20..
URL https
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- May 3, 2018, Workshop Track Proceedings. OpenReview.net, 2018. URL https: /openneview net hwt.:/nnee...otiaeeeeeaeem www.ommmm--..oooee.
Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https
//openrevie ew wwwp:wwwwwweeooeeeteee
W1 UN9gg]
Noam Shazeer. GLU variants improve transformer. CoRR, abs/2002.05202, 2020. URL https
7/arxiv. wwwpiiiorr//aae http://.rrrabbb/////0//// 05202
Zhiqiang Shen, Tianhua Tao, Liqun Ma, Willie Neiswanger, Zhengzhong Liu, Hongyi Wang, Bowen Tan, Joel Hestness, Natalia Vassilieva, Daria Soboleva, and Eric P. Xing. SlimPajama-DC
Understanding data combinations for LLM training. CoRR, abs/2309.108188 2023. doi: 10.48550/ ARXIV.2309.1 1011 URL Utt.ttts:: //ooi. wttp//......///........ 48550/arX 2309. 2209..101111..
14

---

Published as a conference paper at ICLR 2025
Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputiing, 568: 127063, 2024. doi: 127063
httpps://bb...rg/111.11111. newcom. 2023
Mingjie Sun, Xinlei Chen, J. Zico Kolter, and Zhuang Liu. Massive activations in large language models. CoRR, abs/2402.17762, 2024. doi: 10.48500ARXXVV2222.1782 URL https
Xiao Sun, Jungwook Choi, Chia-Yu Chen, Naigang Wang, Swagath Venkataramani, Vijayalak- shmi Srinivasan, Xiaodong Cui, Wei Zhang, and Kailash Gopalakrishnan. Hybrid 8-bit floating point (HFP8) training and inference for deep neural networks. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Bus Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference
Neural Information Processing Systems 2019, NeurlPS 2019, December 8- 8-14, 2019, Vancouver, BC, Canada, pp. 4901 9901- 2019. URL https: ptttpss//prrceeed/ ngs neurips, 1ps Cc/ paper/20p00 9/hash/ Abstract 3551FDD8477887833222212222222222 at .tmll
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix. Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. CoRR, 139971 2023a. doi: 1044850/ARXX..200..1559 ww0.44855501i11. 2302 1397I
URL https:
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Fenton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. CoRR, abs/2307.09288, 2023b. doi: 10.48550/ARXIV.2300.. 48550/arXiv 2307 09288
URL https:://dii org/100
Naigang Wang, Jungwook Choi, Daniel Brand, Chia- Yu Chen, and Kailash Gopalakrishnan, Train- ing deep neural networks with 8-bit floating point numbers. In Samy Bengio, Hanna M. Wal- lach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Ad- vances in Neural Information Processing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 7686-7695 2018. URL https::/ newrips. hwtsscccppae wttpsssccppppppcccceeaaap html
Xi Wang and Laurence Aitchison. How to set adamw's 's weight decay as you scale model and dataset size. CoRR, abs/2405. 2024. doi: 10.48550/ARRIV.24055.1 13698. URL https:
//doi. 41550//aXXv.2205.13333..
Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dickstein, Kelvin Xu, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilitiess. CoRR, abs/ 143222 2023. doi: 10.0::004880 10.48855/ARRXI22309.11
URL
14322
xAI. hotook/..tttss/ggghbbboo.o/////
2024.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng. Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jianxin Yang. Jin
15

---

Published as a conference paper at ICLR 2025
Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue. Na Ni, Pei Zhang, Peng Wang. Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Xuejing Liu, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yun- fei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, Zhifang Guo, and Zhihao Fan. Qwen2 technical report. CoRR, abs/2407.106711 2024. doi: (048550/ARXXV.240..00677 URL https: www.oo........o.ooooo..ooo...... w0w.:///.....21.......
Greg Yang. Tensor programs I: Wide feedforward or recurrent neural networks of any architecture are Gaussian processes. CoRR, abs/ Abss/191..1247, 2019. URL
12478
Greg Yang. Tensor programs II: Neural tangent kemel for any architecture. CoRR, ab4/2006.14548, 2020a. URL BtRLhhtss://arattt org/ abs/ 2006. 14548
Greg Yang. Tensor programs III: Neural matrix laws. CoRR, abs/2009.10685, 2020b. URL
pettps://arxiae wwg//ooggg/oooooooooooooo.oooooo abs/a//s////// 10685
Greg Yang and Edward Hu. Tensor programs IV: Feature learning in infinite-widte neural networks.
In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference
Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event. volume 139 Proceedings of Machine Learning Research, pp. 11727-11737 PMLR, 2021. URL
mlr.press/ http////lll/agglc..... /
html
Greg Yang and Etai Littwin. Tensor programs Il:: Architectural universality of neural tangent kernel training dynamics. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 11762-11772 PMLR, 2021. URL http::/ http:://rooeeiiiic ngs mlr, press / www.//119gg9.... 1ttmn
Greg Yang and Etai Littwin. Tensor programs IVb: Adaptive optimization in the infinite-width limit. CoRR, arwiv. 230..00818...114..
2023. doi: 10445550ARXXIV220800555 URL https:://oo// /rg/ 01814
10. 405448550/aa
Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: Tuning large neural networks via zero-shot hyperparameter transfer. CoRR, abs/2203.0044222222 2022. doi: 10.48550/RXXV..200.. 481.1..48550100
03466. URL httpp:://do wwtp//g.......ooooo.. 485544455////0 att.//////.iiv../ 2203.03466
Greg Yang, James B. Simon, and Jeremy Bernstein. A spectral condition for feature learning. CoRR, abs/2310. 17113, 2023a. doi: 10.45500ARXXV.231017785 URL httpps://ooi. orgg/0 48550 11550/arXivv 2310 17811111111
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: Feature learning in infinite-depth neural networks. CoRR, abs/2310.02244, 2023b. doi: 1048550/RRXIV2210.022222 (RRLhhtts:: /ttp:p:::/dd...o
2310 02244
Chao Yu and Zhiguo Su. Symmetrical Gaussian error linear units (SGELUs). CoRR, abs/1911.03922 2019.
Biao Zhang and Rico Semnrich. Root mean square layer normalization. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference
Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver. BC, Canada, pp. 12360 12371, 2019. URL httt:tt:///oocceeelllittoo .ng.uriipppsnnnn..pp.......... wwwrips...... wttp:pppp..pppppp/...... 9/ 9/hash/ 1e8a 8 73
html
16

---

Published as a conference paper at ICLR 2025
A
ADDITIONAL EXPERIMENTAL DETAILS
A.1 EXPERIMENTAL SETUP
Our experimental analysis of u-uiP was conducted by adapting the codebase used for Tensor Programs V, allowing us to compare yP and u-pi in the same setting. We change various experimental settings from the uP paper to make our experiments better reflect standard training procedures, particularly the dataset which we switch from WikiText-2 to the larger WikiText-103 (Merity et al. (017). Where not specified otherwise, the default setting used our experiments are given in Table 5 These also
represent the settings of our proxy model.
Dataset
WikiText-103 (Merity et al. 2017
Sequence length
256
Vocab size
32000
Training tokens
138M
Architecture
Llama (Touvron et al. 2023a) (Transformer, PreNorm, RMSNorm, SwiGLU, RoPE, "untied" embeddings), non-trainable RMSNorm pa- rameters.
Width
256
(scaled up 4096)
Depth
Number of heads
(scaled up to 64)
Head dimension
64
Total parameters
19.5M (scaled up to 1.07B)
Batch size
64
Training steps
8192 (0.97 epochs)
LR schedule
Cosine 10%, 2000 steps warm-up
Optimizer
Adam W (31, (2, c) = (0.9, 0.999, 10-8)
Weight decay
2 independent (Loshchilov & Hutter 2019
Dropout
0.0
uP HP search range
22 10,10,22 10,2 2-
Nemb E [
Tinit, aemb, Oatin: Couupuu [2-2,27
u-ui HP search range
17 E [2-1 (2-1,23) 23]
Qatin 2,
Oresidual, Oresidunl-attn-ratin Offin-act Coutput [2 1.-3.0 22]
μP HP defaults
Gemb = Qattn = Qoutput Nemb =
u-up HP defaults
Oattn Table 5: Default hyperparameters and training settings.
Oresidual raat O'ffn-act Onutpuu
17

---

Published as a conference paper at ICLR 2025
A.2 VALIDATING OUR EXPERIMENTAL SETUP
In this section series of ablations to validate decisions relating to experimental setup given above. In particular, we examine the effect of using repeated data, the effect using a shorter warmup duration, and the effect of different final learning rates the end of decay.
A.2.1 REPEATED DATA
As outlined in Table 5 our standard training setup uses 0.97 epochs of the WikiText- 103 dataset (50x larger than the Wiki'Text-2 dataset used in Tensor Programs V). However our batch size and training steps scaling experiments in Figure 4 we train on up to 4x the amount of data than in our standard setup, and hence use up to epochs.
Though this still small level of repeated data, this moves our training slightly into the over-fitting regime. Based this change, here investigate the hypothesis that this regime has better worse transfer of the optimal LR than the non-overfiting regime, and hence our results could be mislcading. To do so, repeated these experiments with the
number of tokens, but using the much larger SlimPajama dataset (Shen et al., 2023) where we use < epoch.
The results for this experiment datasets, for both batch size and training steps (albeit with a higher loss, due to the more varied nature of SlimPajama). From this conclude that the effect of repeated data from of WikiText- not significant.
in Figure 7 The shape of curves is very similar across the two
wikitext-103
Slim SlimPajama
4.50
5.00
4.25
Validation Loss
4.75
Batch Size
3.75
4.50
3.50
4.25
128
256
3.25
4.00
3.00
2-1
2
23
2
2-1
2
2
25
Learning Rate
Learning Rate
wikitext-108
SlimPajama
4.25
5.00
Validatton Loss
4.75
4.00
Training Steps
3.75
4.50
4096
8192
3.50
4.25
16384
32768
3.25
4.00
3.00
3.75
21
25
2
2
Learning Rate
Learning Rate
Figure 7: repeat of the batch size and training steps experiments in Figure 4 but using the larger SlimPajama dataset where no data is repeated. In both settings our validation loss basins take the same shape, indicating that our analysis using the WikiText- 103 dataset holds.
18

---

Published as a conference paper at ICLR 2025
A.2.2 WARMUP DURATION
For our experimental setup (Table 5) we use longer duration of warmup than in our large-scale setup (Table 6. We do so out of caution, as we use fewer tokens-per-batch for the smaller-scale experiments and may require longer warmup. However, doing so also creates the risk of spending too large proportion of training doing warmup, which could affect transfer.
To investigate this effect, we run two experiments. Firstly, we re-run the experiment for LR transfer over training steps, shown in Figure on a quarter of the warmup steps. This is shown in Figure (left). The main effect appears to be higher loss for larger learning rates, but the optima are largely unchanged. The only exception is the 4096-step run, where the optimum shifts left and the loss improves slightly. This appears to now align the optimum better with the other training durations, but leads to narrower basins as result, suggesting trade-off for this particular experiment.
However, our other experimental runs use the 8192-step configuration, which has consistent optimum regardless of warmup duration here. To investigate the effect of reduced warmup width transfer at this particular step-count, we re-run our experiment in Figure I (b) under the shorter
warmup duration, shown in Figure (right). The only significant impact of this change is narrow
the basins, inducing no significant change in the optimal LR. As such, we conclude that using 2000 steps of warmup in our experimental setup is a reasonable choice, and both give the same width transfer.
84
26
Wain
Thantng Saya
2.1
125
Valeane Lap
416
4290
240
15384
1051
33700
Winzap
.22
24
Lousing
Lounting
Figure 8: (Left) Learning rate transfer across training steps under different numbers of warmup steps. (Right) Learning rate transfer across width under different numbers of warmup steps. In this setting (training steps = 8192) the optimal LR is consistent, meaning either warmup regime can be used, though the longer gives wider basins.
A.2.3 LEARNING RATE DECAY TARGET
In all experiments
cosine decay of our learning rate down to 10% of the maximum.
This follows the standard approach taken by most L.LM training projects (Touvron et al. 20236 Biderman et al. 2023 Groeneveld et 2024 Almazrouei et al. 2023 Yang et al. 2024). However, recent research has indicated that this may not be the optimal decay target, with implications for LR transfer. Hägele et al. (2024) show that the choice of target percentage can alter the shape of transfer curves and potentially shift the optimum value (Figure 21, right). They also suggest that using fixed target value may work better than percentage (Figure 22, right), which could be swept separately. Anonymous (2024) separately suggest that linear decay to zero is the most effective scheme.
Though using the optimal decay scheme is not necessarily essential the validity of our method, any implications of different schemes transfer properties should be investigated. To do so, we run two experiments. The first sweeps the learning rate for our standard model at various percentages and fixed values of cosine decay target, including zero, in Figure (left). Lower decay targets perform better here, including zero, suggesting that this simple rule may be ideal.
We then re-run width transfer experiment from Figure (b) but with our LR decaying to O, and plot the result in Figure 9 (right). This leads to slightly better results for large learning rates, though for large models this difference diminishes. Fortunately the effect this decay target has the shape of curves (and hence optimal LR transfer) is minimal, indicating that our conclusions are not effected significantly by the choice of decay target.
19

---

Published as a conference paper at ICLR 2025
Min LR
Watn
325
255
Valieebe
532
0.17
1001
2045
1035
1093
10096
Uhvil
122
Loarinn Rati
Figure 9: (Left) A learning rate sweep over LR targets of different types (percentage, fixed and zero) on standard model. (Right) Using the zero and 10% learning rate targets, LR transfer over width.
20

---

Published as a conference paper at ICLR 2025
A.3 PER-TENSOR LEARNING RATES
In Section 3.3 we relax the requirement for each weight tensor in the u-uP model to have an associated tuneable learning-rate multiplier top of the global learning rate. Whilst this does reduce the theoretical expressivity of the u-uP scheme, Figure Figure|TO]shows shows that using a single globally optimized learning rate is already at or close the optimal choice for all weight tensors, and therefore is reasonable to drop these multipliers in favor of reducing the number of HPs. However, a practitioner attempting to absolutely maximize the task performance of their model could experiment with tuning a few key per-tensor LRs, in particular the embedding table.
3.30
Self Atteution Q
Self Attention K
Sclf Attention
3.25
3.20
3.30
Validation Loss
Attention
Linear
FFN Linear
3.25
3.20
3.30
FFN SwiGLU
Embedding
Unembedding
3.25
3.20
2
2
20
22
24
2
2-
20
22
24
2-
2
20
22
24
Per-Tensor LR Multipliers T/
Figure 10: Independently varying per-tensor lcarning rate multipliers 7w, using the u-uP model of width 256 from Figure with optimized global learning rate 21.5 the starting point. Where applicable, the same multiplier is used for tensors of the same name across transformer layers. Each subplot fixes all but one multiplier therefore the midpoint of each subplot is precisely the u-uP256 256 model from Figure
A.4 HYPERPARAMETER INDEPENDENCE
In Section 4.1
we explore the question of HP independence under μP and u-uP. The following plots
in Figures and 12 show the result of 2D sweep over every pair of HPs under each scheme. All other HPs are held at 1 when not swept, except the 7) which is held at 2 7.5 for uP and 21.5 for u-uP, and Nemb which is held 2 for μP.
These results show visual dependence between μP hyperparameters as diagonal structure in the grids, such as (femb; omit) and (17, a.tt)). We quantify this the plot in Figure where we use measure of HP dependence termed transfer error. This is explained verbally in Section 4.1 and we provide an algorithmic description in Algorithm T
We note that differences transfer error between the two methods may also be influenced by the flatness of the optimum. The HP and loss values used for our transfer error calculations are those Figures [I and 22
21

---

Published as a conference paper at ICLR 2025
Algorithm 1 Transfer Error
Require: A 'fixed' HP with candidate values F (/1, ................................ fn ). a 'transfer' HP with candidate values T = {h,,.. (....,,m)) bm), a function that gives the final validation loss for the pair of HPs L: F x T (assuming all other HPs are fixed at default values).
crr +
f", tt" +- argmin(L)
for in F do
if ʃ + l" then 1 <- Argamin(((())) (L(J)
err 1= t))
end end for return err((n- - 1)
Besh
Puah
Outt
Guoppe
335
Valkdotion Loss
-3.3
1
Dracs
Ousgpe
Figure 11: Hyperparameter coupling sweep for pP. Note strong coupling between optima, e.g. the cases of (femb, Tnit) and (17, Cattn). See also: u-uP, Figure I2
22

---

Published as a conference paper at ICLR 2025
Omecattavcatic
Patta
Mooppe
Dreck
3.5
Validation Loos
Ouart
-3.
14
1
Dasigt
1
Oostpa
Datts
Ousger
Figure 12: Hyperparameter coupling sweep for u-uP. Note less coupling than with μP, see Figure 11
23

---

Published as a conference paper at ICLR 2025
A.5 HYPERPARAMETER SEARCH
Here we outline the particular search processes used for our μP and u-ui HP sweeps in Figure (a). The random search samples uniformly from grid defined over all extended HPs (extended HP sets are defined in Table E with grid values defined in T). We perform the random search over 339 runs, each of which of shorter searches at various run-counts by taking a random sample of the results, resulting in the smooth curve over run-count shown.
full training of the w5dth-256 proxy model. We then simulate the effect
The independent search consists of the following phases:
1. Perform a ID line search for an optimal learning rate, with other hyperparameters set to their default values (9 runs).
2. For each hyperparameter in parallel, perform ID line search (330 runs).
3. Combine the best settings from step 2, and re-evaluate (6 runs).
The number of runs the ID line search an order of magnitude higher than required in practice. We do so to form fair comparison with the random search, which benefits from this large number of runs. The number of runs for the ID line search could be reduced further by using binary search, though this would require sequential runs and limit the extent of parallelism.
A.6 HYPERPARAMETER TRANSFER EXPERIMENTS
Baseline μP transfer Figure 13 is a companion plot to Figure 4 in the body of the paper, showing the LR transfer of the baseline μP model over the same axes. u-uP shows marginally more stable HP transfer here relative to the baseline, and at a consistently lower loss.
LR transfer over width The transfer experiments shown in Figure (b) use the non-LR HPs found in Figure [ (a) (indicated by the circled points), rather than using default HP values. For the u-uu sweep we take the HPs the end of the LR portion of the independent search, as these are already close-to- ooptimal,,,,,, and means only 9 runs were required in the sweep. In contrast, for uP it is necessary to use the results of the random search over a large number of runs.
LR transfer over other axes For the training steps, batch size and depth transfer experiments in Figure 4 all HP values are fixed to except LR which is swept. As with width transfer, u-uP outperforms μP here using these default HP values. Reducing training steps is done by fixing the number of warm-up steps (at 2000) and still cosine-decaying the learning rate to 10%; all that changes the number of post-warm-up steps. We found this be more effective than cutting-short the decay schedule. For both Figure (b) and Figure 4 we sweep the LR over a logarithnically-spacee grid of step 21/2x, with runs for each point.
Additionally, in Figure 14 we show learning rate transfer over sequence length for both μP and u-uP fixing either tokens per batch or sequences per batch. In both scenarios u-puP shows not only better absolute training performance, but also better transfer behaviour sequence length increases. Since our default proxy sequence length is 256, using μP to transfer to sequence length 2048 would result in minimal improvements even degradation in validation loss, whereas the u-uP shows much greater and more consistent improvements.
Other HP transfer over width For our non-LR HP transfer results in Figure 15 we note that good transfer under yP has not been demonstrated for all HPs used in the literature. This is particularly true for the Nemb HP, which has poor transfer under μP. Our investigation here led to our identification of the need to adjust the embedding LRR scaling rule outlined in Section 3.4 In many cases users have not swept this HP, but instead swept the corresponding parameter multiplier Oeebh- How this HP interacts with the embedding LR scaling problem identified (and our proposed fix) remains to be explored, though we note in Figure 15 that it also appears to have poor transfer.
Combined HP transfer Whilst Figure 15 demonstrates the transfer of individual hyperparameters over width, Figure 16 instead demonstrates the simultaneous transfer of all hyperparameters when co-optimized on the small-scale proxy model, as is done for y'Transfer. The uP and u-uP points
24

---

Published as a conference paper at ICLR 2025
pP
Training Steps
Batch Size
Depth
4096
Valtcattoon Loss
3.6
16084
128
32709
200
3.4
3.2
2-20
2-
2-
2
2-
2-
2-M
2-1
2-
Learning Rate
Learnngg Rate
Learning Rate
Figure 13: Learning rate transfer for uP over training steps, batch size and depth. This a companion to the equivalent u-uP plot in Figure 4
P
u-pi
3.8
Validation
3.4
Sequence Length
64
3:2
128
258
3.0
512
2048
Fix toloms/tetchh 16k
Valiation Loss
Fix sequences/ /batth
3.4
Min
3.2
2-10
2-
2-
2F Learning Rate
Learning Tate
Figure 14: Transfer of learning rate over sequence length for uP (left) and u-pP (right). As sequence length varies, we can fix the number of tokens per batch by inversely varying the number sequences per batch (top). Alternatively we can fix the sequences per batch and allow the number of tokens per batch to vary with sequence length (bottom). In the latter case, larger sequence lengths mean the model sees more tokens during training, though as per Table this translates to >1 epoch WikiText- 103 when sequence length goes above 256.
taken from Figure I with hyperparameters swept on model of width 256 using full random HP search and simple learning rate sweep for uP and u-uP respectively. The Standard Parametrization scheme, as shown in Table requires choosing a learning rate and a weight-initialization scheme. We follow the initialization scheme of Pythia (Biderman et al.] 2023), and transfer learning rate using heuristic scaling factor of base-width/ width, as is done in Dubey et al. (2024).
25

---

Published as a conference paper at ICLR 2025
Thelnin Looe
3.6
8.2
2.8
Width
128
Lourning Rate
Durta
Qharger
1084
Thining Lave
3.6
4000
3.4
3.2
3.0
2²
11 u.p
2.5
Trilanng Lose
3.4
3.0
2.8
Width
27
21
2-4
2-3
Learning
Datt
Ohargea
256
1024
3.6
4005
Thilnnng
3.3
28
Dloak de
Figure 15: Transfer of model hyperparameters over width for yP (top) and u-upP (bottom). When one hyperparameter is being swept, all others are fixed at 1, with the exception of Learning Rate 7 = 2, 2-75) for (u-uP, uP).
2
3.2
SP
nP
Validltion Loss
3.1
u pp
3.0
2.8
2
29
210
211
2
Width
Figure 16: Transferring hyperparameters from width 256 up to 4096 using three different hyper- parametrization schemes. μP and u-uP results seen in Figure T whilst Standard Parametrization follows the initialization approach of Pythia (Biderman et al. 2023
26

---

Published as a conference paper at ICLR 2025
A.7 NUMERICAL PROPERTIES
Our analysis of the numerical properties of u-uP focuses on the RMS of tensors that we wish to cast FP8: linear module input activations, weights and output gradients. From the RMS training statistics plots in Figure and Figure 17 we note that
1. uP has gradients and weights with low RMS, at risk of FP8 underflow, whereas u-uP starts with RMS 1.
2. Many input activations do not grow RMS during training (due preceding non-trainable RMSNorm), however the attention out projection and FFN down projection have uncon- strained input activations that grow considerably during training.
3. The decoder weight grows during training. Since it is preceded by RMSNorm, the model may require scale growth in order to increase the scale of softmax inputs. Other weights grow slightly during training.
4. Gradients grow quickly but stabilize, except for attention out projection and FFN down projection, whose gradients shrink as the inputs grow.
We also evaluate how RMS growth is affected by model and training hyperparameters the tensors that showed the highest end-training RMS, shown in Figure 18 This shows that the main parameter affecting scale growth is learning rate, with end-training RMS increasing to the right of the optimal LR basin, training becomes unstable. End-training RMS is remarkably stable as width, depth, training steps and batch size are independently increased.
Altn
Attn K
Atta V
Altn Out
2-1
RM3
Down
27
input weight
RMMS
gnd
2-31
25500
5000
7500
0
25000
5000
7500)
254/
5000
7540
25/00
50000
7500
Step
Step
Stop
Stop
Atta Q
Asta K
Atta V
Attn Ous
BMS
2-3
Loyer
2.1
FEN Up
FFN Gate
FTN Down
Decoder
Inpot
wnight
12MS
god
2.3
5000 2.4
2000
75001
2000
7500
2200
20000
7500
2000
5000
7000
Step
Step
Step
Step
Figure 17: RMS during training, for all parametrized matmul inputs, for jP (top) and u-uP (bottom). Model width 256, default hyperparameters, 7) = (2' 2 for (u-pP, μP).
27

---

Published as a conference paper at ICLR 2025
3.4
Trin loss
[a] RMS
(b) RMS
100
10
(e) RMS
256
512 Width
1024
8192
16384
32768
128
256
Learning Rate
Depth
Training Steps
Batch Size
FP32
FP8
Figure 18: The effect of hyperparameters on FP8 training loss and on the end-training RMS of critical tensors: (a) decoder weight, (b) last-layer FFN down-projection input and (c) last-layer FFN down-projection output gradient. Only learning rate has a substantial effect on the end-training RMS. Vertical lines show the default setting of that hyperparameter, used for all other plots.
28

---

Published as a conference paper at ICLR 2025
A.8 SCALE TRAINING DETAILS
Our large-scale training settings are given in Table 6 These are largely the same as our standard experiments (Table but with many more tokens used for training and scaling up larger model-size.
Dataset
SlimPajama (hen et al. 2023)
Sequence length
4096
Vocab size
65536
Training set tokens
600B
Architecture
Llama (Touvron et al. 2023a) (Transformer, PreNorm, RMSNorm, SwiGLU, RoPE, "untied" embeddings), non-trainable RMSNorm pa- rameters.
Width
[2048,3075 4096]
(1024 for proxy model)
Depth
[16, 24, 24,
(8 for proxy model)
Number of heads
[16, 24, 32]
(8 for proxy model)
Head dimension
128
Total parameters
1.077,3.12B. 6.988)
Batch size
1024
Training steps
72000
(~ 300B tokens; 20000 for proxy model)
LR schedule
Cosine to 10%, 500 steps warm-up
Optimizer
AdamW (2,)) = (0.9, 0.95, 10-$)
Weight decay
2- 13 independent (Loshchilov & Hutter 2019
Dropout
0.0
Table 6: Large-scale training settings.
We mixed- mixed-precision during training with optimizer states in FP32 that sharded via ZeRO stage (Rajbhandari et al., 2020). We retain the model weights in BF16 and apply our FP8 scheme as described in Section 3.22 to the tensors participating in matmul operations throughout the transformer block. All other tensors remain either in BF16 (embedding, readout layer, norm, activation function) or FP32 (Flash Attention (Dao al. 2022)).
Each model was trained on several Nvidia A A100 (80GB) or H100 GPUs, with all FP8 experiments conducted the H100 chips utilizing their native FP8 support. For the FP8 operations we use PyTorch's torch... corch__scaled mm function as backbone.
29

---

Published as a conference paper at ICLR 2025
B UNIT-SCALED OP DEFINITIONS
Table 7: Implementations of unit-scaled ops, building on Table A.2. from the Unit Scaling paper (Blake et al. 2023). These are considered part of u-uP and should be used in the place of standard operations.
Op
Unit Scaling factors
matmul(1, (2, w) ww
Ber
Vian-in
Vian-out
V Vbstch-size
attention(q. k) =
0 3g 3g Bg
softmax (aatta Thead (gk") O Cmask)
1/ log interpolate
Sdeend 1,
log(s)
oattn
Tgaat)
92m -Baxaaie 1 pog_interpolate
3 Fgaat + sigmoid(Cren-actt Fgate)
(2:
Haract
residual (osidual_add(zeeee (Tresid, Tskip) = a O.Fresid. + b DTskip
Jo241
V+
(see G.2.2 for full details, inc. values for т, which depends ares and Ores-attn-ratio)
sog_sotoaaaxllsssttttaa rrg_softmax( (oosss x)
a == B=s/v--
RoPE(r)
a 3
(i.c. no scaling)
RMSNorm(z) see (ingle 2024
a 3 =1
(i.e. no scaling)
The Unit Scaling paper provides scaling factors for various ops, in order make them unit-scaled. However, these ops do not cover every case required for the Llama architecture used our experi- ments, nor do they cover we outline series of new unit-scaled ops for each of our required architectural features, as well as existing unit-scaled ops, given in Table
updated residual layer implementation. To address this, in this section
The presentation here is derived from that the Unit Scaling Compendium given in (Blake et al. 2023, Table A.2). This makes reference to the factors o., -1, Sr. the forward pass, and 33 are the scaling factors for the gradient of the op's inputs in the backward pass. For each op, value or rule is provided for determining the required mult ensure unit-scale. The correct value for these multipliers is derived by analyzing the scaling behavior of each op, given some reasonable distributional assumptions about the input and incoming gradient tensors (see Appendix F.2 for an example). Below we provide an in-depth overview of each new modified unit-scaled op we introduce here.
is the output scaling factor in
Unit-scaled dot-product attention The Unit Scaling paper considers the attention layer scaling in terms of its separate components: the various matmul operations and the internal softmax. Linear operations are scaled using the standard rule, and the softmax scaling is given a B = factor.
From an implementation perspective, the self-attention layer is more typically broken down into weight-matmuls and a fused scaled-dot-produu attention operation. This is the case we handle here, accounting for three complicating factors not considered in the Unit Scaling paper:
As we use decoder-style transformer our experiments, our softmax operation has a causal mask applied to its input.
30

---

Published as a conference paper at ICLR 2025
2. We follow the up guidance of using 1/dhead scaling in self- layer, rather than the usual 1/ 1/Vdhead.
3. We place a Cattn multiplier immediately before the softmax, which is an HP that users may tune.
As result our dot-product attention takes the form:
attention(g, k v) softmax (Nattn-softtmax deead (g k.) ④ Cmask) u
The addition of an HP before the softmax introduces an additional challenge for Unit Scaling, as our scaling multipliers will need to account for this value when preserving unit scale.
This operation is sufficiently complex that we found an empirical model of its scale to be more accurate than any mathematicalll-derr derived rule (future work may consider justifying our model mathematically). We find that the scale of dot-product attention is approximately
log(s)
(attention(g. v)) Mlg_interppoltte
1
8
where
log "g_inteerpolac(a, Bupper: blower) 6 log(pupper )+ (x) (fg(b)wwer)
The corresponding scaling rule therefore to divide by this factor both the forward and backward pass, outlined in Table
SwiGLU FFN Llama uses a SwiGLU (Shazeer 2020) layer for its FFN, which introduces two new operations for unit-scale: a SiLU (Yu & Su 2019) (a.k.a. swish (Ramachandran al.. 2018)) operation and an element-wis wise multiplication. We take similar approach to our dot-product attention, and consider unit-scaling the following fused operation:
gated_silut (Zing Tyate) Tin + Tgaate (athh Tgatt)
For the surrounding weight-matmus, follow the standard Unit Scaling rules.
Again, we use an empirical model of the scale of this op, which is surprisingly similar to the not-product attention model:
a
Figate) log_lonteetlnoooo interpotreee
2
Insoet
dividing through by this factor to get our scaling rule.
Residual layers Our implementation of residual layers for u-uP is more complex than other operations, as adjustments are required to:
Make pre-norm residual networks support Unit Scaling (see Appendix F
2. Introduce our new, principled residual HPs (see Appendix
Our residual layer scheme is presented in full in G.2.2 For readers interested in our justification for this, see the sections noted above.
We also follow the example of Unit Scaling and delay the application of residual multiplier in the backward pass to the base of the branch (see (Blake et al. 2023), Figure 3c). This does not change the model, and enables unit-scale be maintained on the residual branch regardless of the value of the multiplier.
RoPE embeddings We also require unit-scaled implementation of Rotary Position Embeddings (RoPE (Su et al.) (2024), which are applied just before the scaled dot- dot-product attention operation. As RoPE essentially consists of pair-wis pair-wise rotations of elements by different degrees, we observe no meaningful scale-change result of it's application, and hence leave it unchanged.
31

---

Published as a conference paper at ICLR 2025
2024) we opt to non-trainable version of RMSNorm (Zhang & Sennrich 2019), in order to facilitate better transfer. As a result, we also leave this operation unchanged. Were trainable RMSNorm to be used, the recipe would follow closely that of the LayerNorm presented in the Unit Scaling paper's compendium.
RMSNorm
Following (Lingle
Scale constraints One final, minor deviation from the scheme outlined in the Unit Scaling paper is the way in which we apply scale constraints (see their Section 5.2). The essence of scale constraints that for perfect unit scaling, sometimes the ideal scale for the forward pass differs from those the backward pass. In some special cases (e.g. at the ends of the network) the use of different scales can be valid, but in the general case a single scale must be agreed upon. The solution the Unit Scaling paper is to use the geometric mean of the forward and backward scales.
We propose instead to simply use the forward scale over the backward scale(s) in these cases. We do so for the following reasons:
1. For these architectures we find empirically that where there is a disparity in ideal forward and backward scales, it is not large.
2. By taking the forward scale, we can ensure strict unit-scale in the forward pass.
The value of the latter point is in terms of what it means for the interpretation of our u-puP multiplier HPs. Consider the Offn-act multiplier; with strict unit scale we can say that the standard deviation of activations immediately before this multiplier is 1. Therefore the standard deviation immediately after is Offi-act- As this multiplier is (by design) the last operation before the ffn activation function, we can say that the interpretation of Offin-ct is simply to set the input standard deviation to the FFN's activation function. Similar arguments can be made for other u-up multiplier HPs. This interpretation only holds because
the forward-scale in our constraints.
C THE CHALLENGES WITH μP IN PRACTICE
C.1 NoT ALL TRAINING SETUPS GIVE TTRANSFER
Lingle (2024) shows that directly applying JP to decoder LM fails to provide LR transfer across width. Given that the primary use of μP in the literature has been LM training of this kind, this result suggests a significant limitation. How do we reconcile this with the strong LR transfer across width V?
shown for language models in Tensor Programs
We answer this in Figure 19 The first training setup (a) is aligned with that used in Tensor Programs (their Figure 4). There are several atypical aspects to their training setup, primarily the use of constant LR schedule and high number of epochs; we outline the precise differences between setup (a) and (b) Table g This overfitting regime makes validation loss unusable, and transfer misleadingly good. When remove these and shift to a standard Llama training setup (b), optimal HPs begin to drift with width (see Figure 21 for an ablation of individual changes). This confirms Lingle's findings that standard μP is in fact poor fit for modern LM training. We fix this (c) by the removal of parameters from LayerNorms/RMMS of independent weight decay for Ad as suggested by Wortsman al. (2023) (see) Wang & Aitchison (2024) for further analysis). With these changes adopted, we recover the strong transfer shown in Tensor Programs V's experiments. Each change is evaluated independently in Figure 20 which shows that the dominant effect is a narrowing of the learning basin due to non-parametric
as suggested by Lingle, and the introduction
RMSNorms, leading to better learning rate transfer.
C.2 IT's NOT CLEAR WHICH HYPERPARAMETER TO SWEEP
The problem of selecting HPs to sweep can be framed as choosing subset of the per-tensor aw, 7/w HPs outlined in Section 2.1 and grouping across/within layers. As shown in Table uTransfer experiments in the literature have done this in a variety ways. Practitioners have not
3 As in other work, we use μP as shorthand for the method outlined in Tensor Programs V. including μ/Transfer. Strictly speaking, μP ought only to refer to the parametrization outlined in Tensor Programs IV. Lingle suggests independent weight decay is unstable, but we find it to be more so than Adam or standard AdamW.
32

---

Published as a conference paper at ICLR 2025
TTnaaee Programs arttings (on: GPT, Wikitest-2)
(b7 Ldemnn aettings (on Llamn, Wikitect-103)
(6) Standaad Elann settings stability fiors
(cn Llama, Wilatext-100)
Valvlatin Lose
10
128
256
512
4.0
1021
20445
0106
Thining Loss
Min
3.4
3.4
0
2.8
2.8
2
Learning Rase
Loarzing Rate
Learning Rate
Figure 19: Effective yTrranffer does not hold across all training setups. (a) We show strong transfer for the unrealistic setup used in Tensor Programs V (too many epochs; constant LR). (b) Moving to more standard Llama training setup, transfer breaks down. (c) This is restored by the introduction of two improvements transfer stability: con-parametric norms and independent weight decay.
(x) Indepencent weight decay
Non-gurametrie rorms
[e]
Independent weightt deeng Not-garamettte norms
Width
3.8
128
3.6
512
Valiiaion
3.4
2045
22
4096
Standard
3.0
settings
28
Min
2-11
2.
2-
2.2
2-11 2. 2 2-3
2-11
2.
2-
Learning Rate
Leaming
Learning
Figure 20: The effect of the individual transfer stability fixes from Figure 19 (a) In this setting switching from non-independent to independent weight decay has only a minor effect, though Wortsman et al.] (2023) Figure 6 suggests it may be highly valuable in other settings. (b) Non- parametric norms give narrower learning rate basin, leading to better transfer. (c) The combination of these, for comparison, matching Figure 19 (c).
33

---

Published as a conference paper at ICLR 2025
Feature
Tensor Programs
Standard Llama
Dataset
wikitext-2
wikitext-103
Vocab Size
33278
32000
Nsteps
10000
8192
Batch Size
20
64
Optimizer
adam constant
adamw
LR Schedule
cosine
Weight Decay
0
0.00012
Positional Encoding
absolute layer norm 0.2
rotary
Norm
ms__norm norm
Dropout
NLayers
2
4
Use Gated FFN
False
True
Activation FN
relu
swish
FFN Ratio
4
2.75
Final Norm
False
True
Base Depth
1
Zero Init Readout
True
False
Table 8: Comparison of Tensor Programs V's standard settings (as best we can tell) and our Standard Llama setup, corresponding to (a) and (b) in Figure T9
Trisce Progouss V Vattiins
Fizal RMSSNom
Lisss undeel Chaston dheayy
5.5
Wilttes-103
Thining Low
3.5
125
366
515
Firal RVVSSNom Wiloicat-Ji J0G
5.5
I Lluta model Choiue LR doeay calakcat-100
Liana medeel Chostze LR docar RhiSSom Wilited---S
Laams
3x Sanch ebm ..]
Triinng Loo
4.0
0.5
000
388
3.5
3-12
1-112 3-14
Lacmnnng
Figure 21: An ablation of the more standard Llama training settings against the Tensor Programs settings from Figure 19 This shows that the flat basins with poor transfer are not due a single change, but the combination of larger dataset (training < epoch) and the stronger Llama model are largely responsible. Note that 'Llama model' here indicates a group of changes: mms norm, rotary embeddings & swiglu FFN.
justified these choices, appearing to rely on mixture of precedent and intuition. We outline two major downsides to the lack of principled approach.
Firstly, not all groupings of HPs are suitable. Consider the commonly-used global Finit. HP. At initialization the activations going into the FFN swish function have std std (Zswish) X OW whereas the self-attentin softmax activations have std(Aamt)) Ow, ØWK A global HP thus has linear effect on the FFN and a quadratic effect on attention, suggesting that this grouping may be unideal.
Secondly, not all HPs are independent of one another. The key example of this is the interaction between ow and nw. The relative size of a weight update is determined by the ratio 7j /ow, not by either HP individually. Because of this, the optimal values for and 77 depend we demonstrate empirically in Section This can make the problem of HP search much harder, and may be why hundreds of random-search runs have been required for sweeps in the literature.
each other, which
34

---

Published as a conference paper at ICLR 2025
C.3 BASE SHAPE COMPLICATES USAGE
Most practitioners are unlikely to require alignment with an SP model, in which case it is unclear what base-width (and base-depth) should be used. The literature has aligned standard base-width of 256 (see Table 9p. but this appears to lacking a principled motivation hherau though the fact that they are not dropped entirely suggests they may be beneficial under u-uuP.
Implementing base-shape HPs (see Equation (3) can also add complications from an engineering perspective. The proposed implementation in the mup library (Microsoft (024) reflects this, requiring an extra 'base' model to be created and the original model to be re-initialized. This can interact awkwardly with other model-transforms for features like quantization, compilation, etc:
import mup
proxy_model MupModel (6 [6_model] -- ...)
proxy width
base_model MupModel (d_model -2006 ...)
base width
mup.set_ base _hapees (proxy_model, base_model)
re-init proxy_model
C.4 μP APPEARS TO STRUGGLE WITH LOW-PRECISION
Finally, we note an interesting contradiction observed in the relationship between μP and low- precision. One of the stated aims for μP that its activations have (1size coordinates in the limit (Yang et al., 2022 Desiderata J.1). This desideratum is specifically given in order that values can be represented using finite-range floating-point numbers (Yang Hu 202T] Section 3). Yet despite numerical stability being central to the theory underlying pP, this not leveraged to ensure that uP models can actually be trained in low- owwprecision. Indeed, for the LLM runs in Tensor Programs V the SP model trains successfully in FP16, while the uP model diverges (attributed underflow of gradients). We remedy this with u-uP.
D A GUIDE TO USING U-up
We bring together our u-up scheme presented in Section 3o form simple recipe for applying it to a model. The u-uP scheme is designed and validated on Llama-style architecture, so it may not be applicable or effective on other models, particularly those with substantially different architectures. Exploring this question is an important avenue for future work.
Before applying our scheme, users are encouraged to apply the following pre-requisites to their training setup, based on our analysis of effective u'Transfer in Appendix C.I
Remove trainable parameters from normalization layers
Use the independent form of AdamW
Ensure training in the under-fitting regime (i.e. avoid excessive data repetition)
Having done this, our recipe for using u-upP is as follows:
Replace operations & optimizers with u-uP versions: Each operation should be replaced unit-scaled version (these wrap the existing operations, with added static scales in the forward and backward passes). We have pre-calculated scales for common operations in Appendix B Parameters should be initialized with unit variance, and Adam(W) adjusted to use the scaling rules defined in Section 3.4 (we refer to the optimizer as Adam in this section, but AdamW should be used if weight decay required. Other optimizer scaling rules can be determined by the same process we outline).
2. Choose a set of HPs to sweep: From the set of HPs outlined in Table select those to be swept. We recommend the extended set, though basic LR sweep can be effective.
3. Decide proxy model config: The cost of proxy model training should be such that the sweeping process is much less than target model training, while still being as representative possible. We base our recommendations the results in Figure 4 In general, width is
35

---

Published as a conference paper at ICLR 2025
the most reliable feature to transfer. Training steps and batch size also give good transfer, so moderate changes here are permissible. Depth is the least reliable feature for transfer, so we only recommend modest changes in depth. We keep the number of warmup steps constant, but always decay to the same final LR when varying the number of steps.
4. Perform independent HP search: Following the process outlined in Section 4.1 and Appendix A.5
5. Train the target model: (though for our large-scare scale experiments we found the scales of two operations drifted enough over time that some lightweight dynamic re-scaling was required).
This can done in FP8 simply by placing casts on matmul inputs
E ADDITIONAL BACKGROUND MATERIAL
E THE MAXIMAL. UPDATE PARAMETRIZATION
Theoretical background We do not cover the theory underpinning μP in this paper, presenting only its resulting scaling rules (Table T.. For readers interested this theory, the extensive Tensor Programs series Yang 2019 2020a Yang & Littwin 2021 Yang, 2020b Yang & Littwin 2023 builds up a framework from which up is derived ( (Yang& & Hul (021). For those requiring a more accessible introduction, Yang al. (2023a) show that μP be derived in a simpler and more general way by placing spectral scaling condition on the norm of weights and their updates.
Approaches to HP sweeping in the literature Table outlines the ways in which users of μP in the literature have approached HP sweeping. These all follow the approach used in Tensor Programs of random sweep, sampling combinations from the joint space of all HPs. The authors of Tensor Programs V note that other more complex methods may be more efficient, but these are considered beyond the scope of their work and have not been used widely. A Bayesian search method was used for the development of MiniCPM (Hu et al. 2024), but the authors give no further details as they use 400 runs in their sweep it is not clear that this approach makes HP search easier.
Table 9: Sweeping configurations used for selection of μP models from the literature. The sweeping process is similar across models, the only differences being the choice of discrete or continuous distributions and their ranges. Model references: T.P.V WMT14 (Yang et al. (022). T.P.V BERT large (Yang et al. 2022. T.P.V GPT-3 (Yang et al. 2022). MiniCPM (Hu et al. Cerebras-Gri GPT (Dey et al. 2023a). SuPar (Dey et al. 2024).
2024
Model
proxy/target tokens used
proxy/target model size
sweep size
base width
HPs swept
T.P.V WMT14
100%
7.1%
64
Ctout Crattn
T.P.V BERT large
10%
3.7%
256
?
7), Nembs Qouts O'attn O',, Obiia
T.P.V GPT-3
1.3%
0.6%
350
O, Oremb; Crout Orainn O
MiniCPM
0.008%
0.45%
400
256
7), O, Olemb; Orresidual
Cerebras-GPT 1.1%
1.5%
200
256
Otemb
SuPar 6.6%
6.4%
350
256
7). O, Olemb
E.2 UNIT SCALING
An example: the unit-scaled matmul op Here outline the procedure for calculating the scaling factor of matmul op, which practitioners can use as guide for scaling new ops that we not cover in this paper (see Appendix B).
There are two potential approaches here. The first is to derive scaling factors from an analysis of an op's dynamics. Specifically, given the assumption of unit-scaled inputs, the appropriate scaling factor the reciprocal of the expected output scale. For basic matrix- matrix matmul we have,
matmul(X, W) XW,
X R dharch diamin
W
B dianiin Xdianout
36

---

Published as a conference paper at ICLR 2025
where weights and activations are sampled i.idd from centered Gaussian:
X1 (0, Wis WN From this derive the expected output scale (i.. o (matmul))::
dianci
matmmn(X, W)sk =
a (matmull W)a) = v ddaa aw ox-
Under Unit Scaling we have aw = ox = 1, and hence the scaling factor required to ensure a unit-scaled output s/Vdian-in This gives our final unit-scaled matmul:
u-matmul(X, W) matmul(X, W)/ diau-in
The distributional assumptions made here hold at initialization, but do not over training. A more precise model for the asymptotic behavior of neural networks under training is given by the Tensor Programs framework, but for the purposes of numeries this precise treatment of scale at initialization appears to be sufficient.
The second, less ideal approach calculating scaling factors is experimentation to infer this relationship empirically. output scale over a range of dran-in values (or whatever HPs one expects the output scale to depend on); fitting a curve to the observed data.
this case, one would sample random initializations and compute the
Applying unit scaling To apply Unit Scaling model and train in low- cow-precision, the following steps are required:
1. Scale parameter initializations to have zero-mean and unit variance.
2. Replace operations with their unit-scaled equivalents (including and especially the loss, matmuls and residual-adds).
3. Constrain the scales of operations which are required have the same forward and backward factors.
4. Place simple to (fp8) cast on the inputs to matmuls.
Step 3 relates the problem of conflicting scales in the forward and backward passes. A single linear layer in differentiated model requires 3 matmul ops in the forward and backward passes, each requiring a different scaling factor Vdiaani 1 Vidaoout Viduatch-oiac However, using these directly would give invalid gradients. The compromise here is that the activations and activation gradients have their scaling factors constrained such that they equal (the Unit Scaling paper recommends taking the geometric mean; we modify this for u-up in Appendix Bio to simply use the forward scale everywhere). Weight gradients still be given their own scaling factor due the cut-edge rule (as explained in Appendix H).
Step reflects the key benefit of Unit Scaling. Unlike other methods it changes the learning dynamics of a model, but the advantage is that unit-scaled models then 'naturally' generate well-scaled tensors. This means that low-precision arithmetic ideally becomes as simple placing operation before matmuls as outlined.
F UNIT-SCALED PRE-NORM RESIDUAL LAYERS
The popular pre-norm residual network architecture is simple to implement, but problematic to combine with Unit Scaling. It exhibits scale-growth in the skip-stream at initialization, due to the repeated addition of residual connections without subsequent normalization. Here we present a surprising and useful finding: that for any pre-norm model there exists mathematically--quivall model where this scale-gee scale-growth is eliminated, through the careful re-scaling of residual connections.
Note that this section focuses applying Unit Scaling to standard pre-norm models. Only once have addressed this problem are we able to do the same for u-upP models, as shown in Appendix G.2 Readers only interested in final u-uP residual implementation may skip ahead to Appendix G.2.2
37

---

Published as a conference paper at ICLR 2025
F.1 SCALE GROWTH PRE-NORM RESIDUAL NETWORKS
Let's consider a pre-norm residual network of depth L:
(2) ro-,,
(4)
(x) r1.(((((1(z)) + Re-1(x), 1= 1,,19
(5)
R (x) (1+1(R((x))
(6)
with embedding multiplier and residual branch multipliers for = 1, L. To satisfy pre-norm, all fi are coro-homogeneuus functions, i.c. f((Az) f((x).
The scale of the skip-stream at initialization as a result of Equation ( is
o (Ri) + ((R-1), [ 1=
assuming ria(f1)² > 0. This shows that scale inevitably grows with the addition of each residual
layer.
This scale-growth clearly incompatible with unit scaling, which aims for a(R) = 1 for all == O, L + 1. In the following we present an elegant solution this problem making use of symmetry transformation available in pre-norm residual architectures.
F.2 RESIDUAL SYMMETRY IN PRE-NORM ARCHITECTURES
To resolve the problem of scale shift residual networks demonstrated by Equation slightly more general ansatz:
we try
Ro(x) I, Ri(x)) = at fi (Rg + by (x))
(8)
(9)
( = f (Rr(()))
(10)
with coefficients at bj. We want choose these coefficients so that the outputs of R are unit-scaled the outputs ft, R are. similar calculation in Equation leads to the sufficient condition
2 + bj 1,
(11)
which can be easily satisfied. Having restored Unit Scale, we are faced with another issue. It seems that Equations (8) (TO) describe different network than Equations (4) to (6). whereas idcally the relation from input to final output should unchanged when converting the network to Unit Scaling.
Note that the coefficients a,, b are not uniquely defined yet, we should find an additional constraint to get unique solution. To find this constraint, let us consider original residual network in Equations (4) to (6) and analyze how the variance propagates through the network we assume all the fi satisfy Unit Scaling and (x) = 1. Let of_, denote the variance of R Then simple inductive calculation shows that
mathematical intuition tells that
oi_1 = -1==1t ?
By Equation (5) the output of Rj adds quantity of scale from the residual connection and quantity of scale d from the skip connection. Intuitively, the ratio of these scales should be more important for the overall network dynamics than their absolute values. Thus our constraint becomes preserving the ratio of scales from the original model, through our choice of at, by:
ai
(rifi) 61-1
T
by
which, recalling Equation [
(up sign) uniquely defines our multipliers at by
aj
b
(12)
38

---

Published as a conference paper at ICLR 2025
In summary, we propose the modified residual network
Ro(x)= I,
(13)
R((x)= 71
V /
fi(R(-1(z))+
1 V + Rx-1
(14)
(x) (RL)(R
(15)
77
(16)
Our main result of this section that this network is indeed mathematically equivalent to the network
defined in Equations 4 to (6). under simple additional structural assumption:
Lemma F.I. Consider RG, Rs defined as in Equations 5 and (14 respectively. Then Ra
Rt/V E for all 0,.,,, O..., L.
Remarkably, this result does not assume the individual network operations fi actually satisfy Unit Scaling. is purely a consequence of the pre-norm residual structure. However, only under Unit Scaling can the factors be interpreted as the ratio of scales between skip and residual branch.
As a consequence of the lemma, the final residual output Rz (x) is the same as in our original network up to a fixed multiplier. Due to the zero eero-homogeneity of the final output function fL+1 this
gives Ro+1 = fL+1 R2(2)/ V = fL+1(R2(z)) = RL+11 proving the mathematical
equivalence of residual scheme. Modern LLM architectures like Llama (Touvron al. 2023a pre-norm residual networks of this kind. Hence they admit faithful unit-scaled reparametrization.
F.3
PROOF OF LEMMA F.
Proof. This is proved by induction. For the base-case 1 1, we have T rr/ro. giving
1
(1)x) 7 +1 fi(z) +
7 +1
(rifi(x) + f/02//
v18 +
= RRI/V 18 =RI/v18+1 +7.
Then if the statement holds for 1- we have
fx(2))
11
(f(e)1 +
(2)
V +
V 7
V E R 1
f((R(-1(2)) +
Ei
El ?
((f((((1(()))
r (()1-(())
/
Σ
i=1
Rx-
1
= (()((R(-(()))+++
Σ p?
Σ ?
Σ m
+ R R
Σ r2
= -R((x)/
Σ 2
1=0
39

---

Published as a conference paper at ICLR 2025
F.4 UNIT SCALING FOR TRANSFORMER RESIDUALS
The above scheme describes Unit Scaling for arbitrary pre-norm residual networks. We now apply it to the case of pre-norm transformer residual layers.
We describe transformer in terms of the residual network given in Equations 4 Our fi functions alternate between self-attention layers and seed-forward layers. Implementations differ in the handling of how residual multipliers 72 correspond HPs. In many cases practitioners simply ignore these r,, but for the sake of expressivity we assume the two types of residual layer each have their own HP, as well as the embedding. In other words,
Cemb
1= 0
71
Cattn-residual
| is odd is even, and
Afn-reiiuual
To convert this to a Unit Scaled network we apply Equations (3 to (6)), from which can derive the following closed-fome expression for T/.
attn rosidual
is odd
emb
la
la rassriiuull
T/
resiiual
is even.
Lemb (e 1) 1) attn-residual
la
where L 11-1
1-1
This gives us a unit-scaled pre-now pre-norm residual implementation for standard transformer, which is mathematically equivalent non-unit-scaled version. In the next section augment this by adding two HPs, in a carefully-designee manner that satisfies our criteria for u-uP HPs, giving us our full residual implementation.
G JUSTIFYING THE U-uP HYPERPARAMETER SCHEME
Here we justify our particular choice of u-up HP, given in Table (with their placement defined in [7. We discuss this topic briefly in Section 3.3 stating that all our HPs (excepting the LR) are HPs, and under u-uP they are now associated with operations instead of weights. All operations have an a HPs, unless they are unary and k-homogeneous for k 2 0.
We begin this section by explaining why we apply this rule to the model and how it results three of our u-pp HPs. We then consider how best to hyperparametrize our residual layers, building our criteria for HPs given in Section 3.3 and the unit-scaled pre-norm residual scheme in Appendix
G.1 MULTIPLIERS FOR NON-HOMOGHNEOOUS OPS: Oatin-softuma: Offan-cct Closss-offtaax
In this section derive the rest of our u-up multipliers. We want identify the minimal set that can still express all different choices of pre-op scales in the model. The crucial observation that every pre-scale multiplier of unary operation h f(ah) can be propagated through the network if f is k-homogeneous for some k > 0, i.e. f(az) a ((x), leaving the model and its optimization unchanged. We can iterate this along the computational path until either the next operation is non-homogeneous. non-unary (we are at the end of a residual path), the next operation 0-homogeneous (e.g. a norm).
In the first case the accumulated scales are absorbed in the pre-op scale of the non-homogeneous operation (where we introduce a multiplier), in the second case they are absorbed in the residual addition for that branch (where again introduce multiplier), and in the final case the scale disappears (so we start over). We now go through the Llama forward computation and follow this paradigm identify multipliers in Table TO
40

---

Published as a conference paper at ICLR 2025
Table 10: A walkthrough of the Llama architecture, showing how our desttn-softmax Offn-act and oosssssf multipliers are derived via an analysis of scale- poopappropagtion..
Op
Scale propagation behavior
Embedding
show in Appendix G.2.1 that the embedding multiplier can be absorbed in the residual multipliers, meaning one is not required here.
Attention RMSNorm
This operation is 0-homogeneous and thus start over. Both are linear, meaning their scale propagated. Multipliers are therefore not required.
Query & key projection
Query-key matmul
Again linear. As query & key are both generated from the same input, this operation is 2-homogeneous wrt. that input. Hence it also propagates scale.
Softmax
The softmax operation is non-homogeneous, Thus the pre-op scale of the softmax becomes our first multiplier: daatta-sfftmax-
Value
The value layer is linear and hence propagates scale.
Softmax-vale value matmul
Again linear and hence propagates scale.
Attention projection
This operation is linear and lies at the end of the attention residual path. Hence there are no more multipliers in the attention block.
Residual add
This operation is non-unary and hence receives our second (and third) multipliers: Gress Cresattn-ratio- The manner and motivation for using two multipliers here is justified in the next section.
FFN RMSNorm
This operation is O-homogeneous and thus we start over.
FFN input scale
The input layer is linear, hence propagates scale.
Sigmoid input
This function is non-homogeneous and thus we have our fourth multi- plier: Offin-act
SiLU weight
This layer is also linear and propagates scale.
Product
The entry-wise multiplication of the outputs of sigmoid, input layer and SiLU weight is homogeneous and thus propagates scale.
FFN output
This layer is linear and the end of the residual path. Hence there no more multipliers in the FFN residual block.
Residual add
See above.
Output RMSNorm
This operation is 0-homogeneous and thus start over.
Output head
This layer is linear, hence it propagates scale.
Loss
The cross-entropy loss is non-homogeneous and leads to final multiplier: Clossssofttxx
G.2 RESIDUAL BRANCH MULTIPLIERS: Cress Ogus-attn-rratio
In this section we derive our two u-uP residual HPs. We start with the basic, non-unit scaled model we began with the previous section, outlined in Equations to (6). We described a set of aemb, Qattn-residual: Offnnresidual HPs associated with this model in Appendix F.4 However these
HPs poorly satisfy our cardinality, independence and interpretability criteria from Section (333 so in the Appendix G.2. we present parrametrrizaion of these HPs designed to better satisfy these points. Appendix G.2.2 we then combine these HPs with the final unit-scaled pre-norm residual scheme derived in Appendix resulting in complete u-uP residual scheme.
41

---

Published as a conference paper at ICLR 2025
G.2.1 IMPROVED HYPERPARAMETER FOR TRANSFORMER RESIDUALS
To avoid cluttered notation, in this section we rename
O'res
O'p
Ores-stta-ratio = Upp
Oeemb Otey
Gattn-residual
Qa
Offin-residall O'f.
To make the presentation more clear, derive our new HPs using the standard residual scheme from
Equations (4) to (6). For the actual unit scaled implementation one needs to transform the multipliers
following Equations
(T6), which do in Section 22222
To facilitate our analysis, we can view the transformer residual output as the sum of three terms:
Rr 12 + +R(,) + -11,2
R/( = Cc
L/2
R
aa
f (
L
1=1
L/2
RL6
a
as
f2i (2(())))))
VL/2
and define the average residual scale.
o a(R(al j + (1)
2
Note that we have added in the depth-pP multipliers here, though similar analysis can be performed for non-depth-uP models. As above, functions alternate between self-attention layers and feed- forward layers.
With respect to our interpretability criterion, we propose two new multipliers that correspond to dynamics in the network which we suggest are important to control at initialization. The first the ratio of the average scale of the residuals' contributions those of the embedding, dr, //a( (R(()) The second is the ratio of the scale of the attention-residuaass contributions to or (R( those of the feed resuuur- a, = o ((() (1(( Not only do these two ratios control key dynamics of our model, but we can use them replace our existing (ae, (aa: at) multipliers.
Let us first examine these two quantities for standard (non-unit-scaled model). Residual functions of the same kind have the same expected output scale at initialization in pre-norm networks, meaning we can denote the output scale o (R))) of all self-attention functions as oa, and of all feed-forward functions oy. We thus have the following scales the output:
o(R(.) (x)()),
L/2
o(R(
aa
Σ fzi 1 (R21-1)
adoa.
LL2
1,/2
a(R(f))
af
Σ
Sei (Rzi)
=ajof
VL/22
(ap0a)²
+ (a
2
42

---

Published as a conference paper at ICLR 2025
Recalling our definitions of a,, Op above, this gives us:
Ga
Pp a os
ay (aaDa)² (ajo))²
(z)o
0 +1 of as
2 ((x) ae
The original au ay multipliers can then be written in terms of ar: Op:
= apaf
0a
0f = 0, (xx)
2
03 +
We have replaced two of the three original multipliers, but still have dependence on a, here in our expressions for o and R2(e) which we now remove by dividing it out of our residual branches and embedding. We use the hat ()) symbol to denote terms that have been divided-through by dee This
new system of equations is equivalent to our old one thanks to the zero-homogeneity of the final post-residual layer:
RL+1(x) = fL+1( + R + R(
= S + R + RL())
f + R + R(
This gives R(() o.x/O. x, removing our first occurrence of dee Following the division through R and f(f) we obtain:
L/2
A(a) =
Σ
òa
f21-1
(0--1(R11-20
1R21-1);
LL2
L/2
600 Σ
fail (22).
LL2
6a apå da
2 a²
â ar of
This system of equations is the same as the original, but with the two a, terms dropped, meaning our model's multipliers can be expressed in terms of only o, and app Using the above equations, any pair of values for (ay, ap) can be translated back into an equivalent set of values for (ae, Qa. ay) such that the output R,+1 (xr) is the same, meaning that our multipliers are no less expressive than the original set. This satisfies our desired criteria of minimizing the number of multipliers while maintaining expressivity.
We can simplify further in the case of unit-scaled models, which are designed such that (x)) a of are all 1 at initialization. In this case our re- re-parametrization becomes:
áa =apuf
(17)
2
6 ar
(18)
az +
â. 1.
(19)
43

---

Published as a conference paper at ICLR 2025
This is the basis of our claim that Unit Scaling is what enables more intuitive set of multipliers. Not only do the multipliers a, and ap represent important dynamics in the network at initialization (the ratio of residual-to caaiuual-to-ebeeee ee-eeddiednngg scales, and the ratio of via unit scaling that these equations become simple enough to implement in practice. Using equations Equations (17 (19) for non-unit scaled network may still be effective, but the interpretation we've ve given to and ap no longer hold.
scales), but it's only
Our final desired property an empirical one: that the most effective choice of one multiplier depends as little as possible on the choice of the other multiplier(s). We demonstrate that our multipliers satisfy this property better than the standard set residual multipliers Section 4.1
G.2.2 THE FULL U-uP RESIDUAL SCHEME
Here we give the full definition of our u-pp residual scheme, summarizing the results of previous sections. general pre-norm transformer implemented as:
(o = C,e,
(20)
(x)x) a (
1=
(21)
R = fL+1(Ru
(22)
where a,, by and are scalar multipliers, and the alternate between self-attention and feed-forward layers. We consider our baseline set of jP residual HPs here to be (Cemb, Cattn-residual: Offnnresiduul). which we implement (assuming depth-uP branch scaling) as:
Cestun-r
I is odd (self-attention)
VL/2
01
Offin-esidual
is even (feed-forward)
LL/2
bj
e Gremb-
The corresponding u-up set of residual HPs is (ares; Ores-atta-ratio). which implement
af 2
(23)
7
1
6
(24)
7
+1
c=
1,
(25)
(26)
/ odd
L + lo2 + tā?
1-
71
2
l
(27)
2
1 is even
1
+
1)62 + lā?
2
0
2
ates-attn- à²
2
(28)
2
à Pessattn-ratia ratio
Ores
(29)
This is the u-uiP residual scheme. satisfies the three properties that we initially set out to achieve: the variance at initialization of our R((x) is always 1, our HPs have clear and useful interpretation, and our scheme is as expressive as the baseline (which is neither unit-scaled or has interpretable HPs).
H THE CUT-EDGE RULE
In the section review the notion of constraints used for scaling operations in computational graph. For a more thorough, generalized treatment, please refer to Section 5.1 and Appendix E.4 of the Unit Scaling paper (Blake al. 2023).
44

---

Published as a conference paper at ICLR 2025
For simplicity, will only discuss the cut-edge rule in the context of typical neural network. For each operation f. parametrized by taking imput a and emitting output y, user must choose how to scale y, Va and Vo (gradient of loss w.r.t. and 0 respectively). In the simplest case, where there no further data dependencies, we can simply choose factors that preserve unit scale. In more complex scenarios, we must balance the need for each tensor to be unit-scaled and for gradients be correct up to constant factor.
In particular, a problem emerges in the presence of residual blocks in which y = f(a; (). In these cireumstances, V is computed as the sum of residual gradient V of and skip gradient Vy. If 8i we choose not to insert scaling factors into our graph, V @f and Vy y will have some ratio of scale r. at
However, if we have chosen rescale the gradient of operations in f. then V, Øf will have been rescaled by some 8. This means the new ratio of V 0f these together, Vz is longer correct gradient up constant factor.
ae and V, will be r. Therefore, when adding
How do you remedy this? we can ensure that the scaling factors are the same for both the input gradients and outputs of op, we will have 1. This ensures that gradients for inputs to residual blocks are correct up to a constant factor.
How do you decide when you are free to preserve unit scale, and when constrain scaling factors to be the same? We previously define the cul-edge rule (Blake et al. 2023) for computational graphs where nodes represent forward pass operations and edges represent operation outputs. If an input edge is cut-edge, i.e., the number of connected components in the graph would increase upon deletion (examples in typical transformer model: output of embedding gather, output of a residual add, output of final norm, output token logits, weights), there is no need to constrain the scales of the operation's ' output edge and the input edge gradient. For all other input edges (e.g., inputs to residual add, intermediates computed along a residual branch), the scales of gradients and outputs should be constrained.
I FROM μP TO U-HP
Here we outline additional details to help readers follow the process of deriving u-ui from the combination of Unit Scaling and μP. Our first step of dropping ow and base-fan-in, and moving aws to functions, results in Table [T] This intermediate scheme does not yet satisfy Unit Scaling, but simplifies the HP rules in preparation for further changes. Note that we have also removed femb as we don't 't include this HP in our u-uP extended HP set. We have included residual scaling rules here, in accordance with depth-uP, which we intend u-up to satisfy, though our standard jP implementation doesn't It use it.
Table 11: An intermediate scheme resulting from dropping those HPs from jP which are not needed under u-uP.
Weight Type
ABC-multiplier
Residual
Input
Hidden
Output
parameter (Aw)
lan-in
v
depth
initialization (Bw)
1
1
Vfan-in
Adam LR (Cw)
fan-in
1]
Vdeett
J LOW-PRECISION AND ITS TRADE-OFFS
Number formats for deep learning The standard numerical representations used in deep learning the set of formats defined by the IEEE 754 floating point standard (IEEE Computer Society 2019). IEEE floats comprise three elements: a sign bit, exponent bits, and mantissa bits. The number of exponent bits determines the range of format, while the mantissa determines the precision'.
45

---

Published as a conference paper at ICLR 2025
We refer readers to Blake al. (2023), Section 3.1 for comprehensive overview of floating poattngg-pinnt representations.
The default format used for training is the single-precision floating-poine point format, commonly known as FP32, with some hardware providers automatically casting it to the smaller TF32 compute mode for accelerated arithmetic. The 16-bit FP16 and BF16 formats were later introduced, and more recently the FP8 E5 E4 formats (Sun et al. 2019 Noune et al. 2022 Micikevicius et al. 2022). The higher range of E5 has typically been used for gradients, while the higher precision E4 has been seen as necessary for weights and activations. Other aspects of training such as the optimizer state and cross-device communication have also been put into FP8 (Peng al. (003). though not all tensors are amenable to being run in the lowest precision (Dettmers et al.] (022) without degradation. The use of multiple formats is known as mixed precision (Micikevicius et al. 2018). comparison of these formats is given in Table 12
Table 12: A comparison of deep learning formats. E indicates exponent bits, and M mantissa bits. The smaller formats typically give more FLOPS, at the expense of reduced range and/or precision.
Format
E
M
max
min normal
I min subnormall
FLOPS (vs TF32)
FP32
8
23
3.4 X 10S8
1.2 x 10 38
1.4 X 10-45
< 1 x
TF32
8
10
3.4 10
1.2 X 10-38
1.1 x 10-41
1x
BF16
8
7
3.4 x 1038
1.2 x 10 10-38
9.2 x 10
2x
FP16
5
10
65504
6.1 x 10-5
6.0 x 10
2x
FP8 E5
5
57344
6.1 X 10-5
1.5 x 10
4x
FP8 E4
4
3
448
1.6 x 10-2
2.0 X 10-3
4x
The benefits of low-precision Using numerical representations with fewer bits facilitates the design of more efficient arithmetic in hardware, typically leading to a linear increase in peak FLOPS (as shown in Table As lcrge-scale training efforts are typically compute-boure comuute-bound due to the size of matmuls (Narayanan et al., 2021), putting the inputs to these operations in low- oowpprcission formats has substantial impact on training efficiency. Low-precision formats also reduce the other two common performance constraints: for memory memory-handwidth hand--athhhhhhhun models they require fewer bits to be transmitted, and for memory-size-bont models they require fewer bits to be stored.
The challenges of low- oow-precision Unfortunately, moving to low-precision formats also increases guantization error. For values within the representable range this takes the form rounding error, and for values outside it, clipping error (both overflow and underflow). Rounding eπror tends to be an intrinsic problem: the number of mantissa bits dictates the expected accuracy of representations and this cannot easily be changed. In contrast, clipping error is often eliminated by scaling tensor so that its values lie within the range of a format. Note that multiplicative change in values of this kind doesn't 't affect the (relative) rounding error, due to the exponential spacing of values. Most research into making low-precision work has focused on the problem of scaling tensors in this way.
Simply casting all tensors FP16 or FP8 tends impair training, largely due to clipping error. For FP16, this primarily affects gradients. Micikevicius et al. (2018) address this by introducing fixed global loss-scale HP, which multiplies the loss value in the backward pass, artificially up- up-scaling gradients to lie within FP16 range. Automatic loss scaling (Kuchaiev et al., 2018) builds upon this idea, making the loss-scale dynamic value that tuned during training.
The later BF16 format has the same range as FP32, making loss scaling unnecessary. For FP8 no such range-equivalent format can exist, so the problem of clipping error must be addressed. Most FP8 implementations have done so by moving from global loss-scale to a local scale for each FP8 tensor. In pseudo-code, this takes the form:
scale (A)
scale (B) to_ fp8 (A / a)
Confusingly, the term low-precision tends to indicate using <32 bit-width formats, so in this context precision also reflects the number of exponent hits well as the usual mantissa bits.
46

---

Published as a conference paper at ICLR 2025
fp__p8 (B b) matmul (A, B)
where we assume that matmul takes inputs in FP8 and directly produces the output in higher precision.
The result of the scale () operation can either be fixed scale determined before training (Nuuce et at. 2022), or in the case of Transformer Engine (NVIDIA) 2024a), computed dynamically as function of the 'absmax' of the input tensor (though they introduce delay across time-steps, to facilitate an efficient fused kernel). Increasing granularity and computing scales dynamically using this kind of method inevitably adds complexity (from both a logical and implementation perspective), as well the potential for computational overhead. Unit Scaling generally avoids need for matmul input scaling.
K BENCHMARKING SCALED MATRIX MULTIPLICATION IMPLEMENTATION IN PYTORCH
Given that the end-goal of leveraging u-mup's 's low-precision properties is to speed up training and reduce memory usage, it's reasonable ask why don't 't investigate this experimentally. The answer relates to the relative immaturity the FP8 training software stack a lack of open, efficient FP8 kernels for compute and communication mean significant additional engineering effort is required to attain expected speedups over the full model.
Here show that u-uP'ss static scaling factors add overhead to matmuls in FP8, and hence ought to be able to reach close to the maximal FP8 throughput attainable for the full model.
Boat32
(1P163550
Matuel tya
300
wxid
110
2500
1030
6000
380)
3000
9000
5000
530
1000
4100
2003
CBEMM
Figure 22: Square matrix multiplication throughput in TFLOPs with and without scaling factors applied to the output across 32-, 16-, and 8-bit float dtypes on NVIDIA H100 PCle. Naive implemen- tation in PyTorch.
Standard strategies for FP8 training require expensive statistics gathering (e.g., amax) per tensor. A key benefit of u-up for FP8 training is that it instead provides us with static scaling factors rescale operation outputs. Even naive implementation in pytorch can achieve minimal drop in hardware utilization.
Figure 22 demonstrates hardware utilization for FP8, FP16, and FP32 matrix multiplications on single NVIDIA H100 PCIe card. For FP16 and FP32, torch. miich.matme matmul used, whereas forch._scall
is used for FP8. Comparing "scaled" to "unscaled" matrix multiplication demonstrates a 30%, 20%, and 10% drop in hardware utilization for each data type respectively. In the case of FP8, where the drop in utilization is most pronounced, utilization be recovered by passing the scaling factor as a scale associated with one of the two input tensors.
It should be noted that of PyTorc version 2.3, torch. scaled, always computes amax as well as the matrix multiplication. The performance of FP8 matrix multiplications could higher without this overhead.
47

---

Published as a conference paper at ICLR 2025
The above analysis focuses on throughput; significant memory savings also possible through the use of FP8, though how this affects the total memory footprint depends on various additional variables and the overall distributed training setup. The following factors play significant role: typically the main memory bottlenecks are the optimizer states, which are kept in full precision. This footprint be reduced by applying ZeRO sharding (Rajbhandari et al. 2020), though for significant gains the number of data parallel processes needs to be sufficiently large and ZeRO stage 2 or 3 are required. In these settings the memory footprint of activations and gradients becomes significant, and quantizing these to lower precision promises further memory savings, though may be non-trivial (Peng et al. 2023).
L ATTENTION OUTPUT RMS GROWS WITH MODEL DEPTH
Residual Activations
Residual Gradients
mip__dwn
layer idx
layer idx
mip_ nat__gaee
Tensor Name
mip_up
mip_nee norm
attn_out
attn_gky
attn_norm
Skip Activations
Skip Gradients
4
3
layer idx
2
2
2-1
2
2
2-1
20
2
RMS
RMS
Figure 23: Scale of intermediate tensors grows with depth at initialization. Top left: Intermediate activation tensor RMS along the residual branch. Only the attention outputs after the first layer are not unit-scaled. Bottom left: Skip activation tensor RMS. Scale growth attention outputs drives growth in skip activation scales. Note that layer_i dx= 0 corresponds to the embedding output, and 1 layer_ide <idx corresponds the final layer outputs. Top right: Intermediate gradient tensor RMS along the residual branch. Growth in the attention output scale drives growth in attention qkv gradient seales. Bottom Right: Skip gradient tensor RMS. The scale output activations induces global rescaling of the gradients.
48

---

Published as a conference paper at ICLR 2025
A core assumption in deriving per-op scaling factors is that each input operation has zero mean, vari-var, and uncorrelated elements at initialization. This is trivially true for weights and by extension the token embeddings taken as input to the transformer trunk. However, this is not guaranteed for intermediate results and gradients if an operation the computational graph induces corrllation in the elements. In such scenario scaling factors will not return unit- mnitvvaiancce outputs as we will not have corrected for these correlations in the inputs. As we then increase the depth of the network, where the same operation left to amplify correlations, can end up with variance in intermediate results and gradients scaling with depth
Figure 23 illustrates this phenomenon in a unit-scaled four-layer Llama model with width=256. All activation tensors in the residual branches are unit-scaled, except for the output of the attention layers. We also see that the variance of attention outputs grows with depth. Since Ilama models use pre-no norm on the residual- -branch,, residual- -branch inputs will revert to unit-scale again until they reach another instance of the correlation-indums induuiing operation. As we add under-scaled attention layer results back to the skip-branch, our skip tensor variances grow with depth as our residual-add assumes unit-variance inputs. This has a knock-on effect on the global scaling of the gradients since the Jacobian of the final norm will scale the gradient by the inverse of the final skip tensor variance.
So which operation induces correlation in the attention output at initialization? For the default case where all multipliers are set to I, our 1/d scaling of attention logits results in sufficiently high temperature that attention probabilities are effectively uniform. With causal masking, we effectively across the value tensor along the sequence dimension. As result, each
take running subsequent token representation is correlated with the last. Since we derive appropriate scaling factors for the first layer, we do not see scale growth emerging until the second layer, where correlations accumulate during the next effective running mean.
We leave it to future work to offer a solution to scale growth created by correlation in intermediate tensors. We note that this is scale growth emergent at initialization, but we also see scale growth in other intermediate tensors during training. Whether scale growth during training is related to the phenomenon outlined here remains to be seen.
49


Reviews:
[{"review_id": "59Fx32fsbN", "reviewer": "Reviewer_gFHM", "summary": "The authors present a combination of the maximal update parametrization (muP) and unit scaling, coined u-muP. It brings together the two main ideas of 1) hyperparameter transfer and 2) the ideal principle of unit variance of activations, weights, and gradients. The authors implement this idea with decoder language models, showing both the transfer of parameters, the performance, as well as scaling to 7B models and FP8.", "questions": "I have raised questions/weaknesses in the section above, which I would be happy to discuss and eventually raise my score. Beyond, I am curious what the authors think about HP transfer to larger scale and longer training. This connects to the first point above. For instance, there seems to be a slight shift of the optimal LR to the right when growing batch sizes (Fig. 4 middle). This could be problematic for planning large scale runs (where batch sizes have become enormous). Similarly, do you have concerns about transferring to much longer training lengths (e.g. more than 1M steps for Llama 3, unfeasible for LR sweeps). To be clear: I do not expect the authors to have a solution for this, I am just curious about their thoughts.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "I am very positive about this paper -- I think it is both a valuable contribution and an important direction for future work, as it is both practically and theoretically motivated and I agree with the concept of unit scale. I also appreciate the demonstration of failed HP transfer of muP for typical Llama-like models, which I have experienced myself. The experiments are very broad and consider not only HP transfer, but dependence between parameters, numerical properties during training, FP8 and 7B scale with downstream evaluations. Certainly, the 7B experiments are most convincing. \n\nSome side note: As the authors note themselves (so I do not see it as a weakness, but future work), unit scaling does not give guarantees for the behavior during training, so I would be particularly interested to see this method combined with models like the outlier protected block (He et al., NeurIPS 2024 https://arxiv.org/pdf/2405.19279), investigating the outliers during training, for which this model might enable even better/easier FP8 training.", "weaknesses": "While I am an advocate for the paper, I want to raise some points/irregularities that came to my mind while reading and I think would need to be addressed, both to improve the work, its insights or my score. They concern, in particular, the experimental setups:\n\n- Why not use a larger dataset for the HP transfer experiments? I understand it is to compare to the setup of Yang et al., but I am asking because it is really small compared to modern training settings. For instance, in Fig. 4, 34k steps imply ~4 epochs over the dataset? Similarly, a warmup of 2000 steps is relatively long compared to the overall steps? In comparison, the large scale training only used 500 warmup steps.\n- When sweeping the LR, is the final LR always changed to 10% of the chosen rate? The final LR should either be swept independently or kept fixed to a low enough value for a proper LR cooldown, otherwise this can skew results.\n- Comparison to SP, in particular 7B: If understand correctly, you use the exact same model for SP and u-muP. Does this mean you also use the tweaks (non-trainable RMSNorm, independent WD) for the SP model? Since the LR setup for Llama was chosen without those changes enabled, I think a fair comparison would be to use the original model, in particular with coupled weight decay.\n\nI particularly think the point on the comparison to SP is very important — the main focus of the paper is on a comparison to muP, but there is the simple reason of being more convincing to adopt the method because of its performance and not just its elegance (e.g. for practitioners, many of which haven’t adopted muP yet either).\n\nRelation to prior work: I think it would be important to add references and discussions to the field of signal propagation in neural networks, e.g. Noci et al. https://arxiv.org/pdf/2206.03126 and the many references within their related work.\n\n## Update\nI am very happy with the rebuttal and the authors' response, hence I have increased my score to 8.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "kliAlIjQrN", "reviewer": "Reviewer_CvCm", "summary": "This paper combines two approaches—(1) mup and (2) unit-scale—to enable hyperparameter transfer for low-precision training (FP8). The paper is overall well-written and contains many details and ablations.\n\nI do like this direction of research, and low-precision training is important for speeding up training, lowering training costs, and enabling more research.\n\nHowever, I have several concerns regarding the paper listed below. I am happy to raise the score if these are addressed properly.\n\n\n\n---- update ----\n\nI am delighted with the promised change. in particular, down play the batch, depth scaling results; detailed discussion of embedding scaling and potential limitation; \n\ni raised my score to 6 ( i would give 7 if that option is available).", "questions": "Independent weight decay vs coupled weight decay? need more clarification. I am not sure which one to use. Lingle proposes to use wd=0.1 * LR, while og mup, wortsman and here propose scale-independent weight decay.  \n\nFig 2. Are the legends regarding C_embed correct? It says C_embed = 1 is better than C_embed  = 1/ root(fan_out). \n\nThe uses of non-trainable rmsnorm \"scales\" seem non-standard to me, have you done ablations by yourself?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "I think this is important research direction that has many practical values. The paper contains a lot of useful details (many of them are in  the appendix). I really appreciate that. \n\nThe embedding scaling rule seems interesting and novel but also controversal.", "weaknesses": "- Transfer across batch, depth, training steps is not convincing (Fig. 4).\n - (important) The learning rate in the embedding seems unnatural and contradicts the original mup paper. In the infinite-width setting, the update will go to zero, and the input layer is frozen; this doesn't seem right to me.\n - (Important) Everett also studies hyperparameter transfer thoroughly and is highly related to this paper. Please make a more comprehensive comparison. In particular, the mean-field parameterization is very close to the unit-scale proposed here. Please clarify what's new and what the differences are.\n- Citations to previous mean-field papers are needed and should have been discussed.\n- The results of the paper seem to contradict some results of the original mup paper. I would love to see a table/section summarizing them + an explanation of why the original mup setup is not right.\n- I  also want to see a wall clock time comparison between fp-8 vs. bf-16 runs at different scales, which is why we want to use fp-8. In addition, it may be good to know how much memory can be saved.\n\nAgain, I think this can be a good paper and I would like it to get improved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "nwfP6QcteF", "reviewer": "Reviewer_uAPB", "summary": "This work proposed u-$\\mu P$, which combines $\\mu P$ with unit scaling to facilitate hyper-parameters (HPs) search. u-$\\mu P$ does not require a base shape and the HPs have less interdependency. Moreover,  u-$\\mu P$ empirically maintains unit scaling for the activations, weights and gradients, which enables FP8 low-precision training.", "questions": "* Maybe the authors already mentioned this, but I wonder how is the performance of u-$\\mu P$ with embedding LR =1 compared with $\\mu P$?\n\n* I am curious about how whether u-$\\mu P$ can be applied to non-transformer models and how to select the HPs (as in line 304-316) in this case? I am not asking for additional experiments, any useful insights or discussion would be appreciated.\n\n* Typo in Figure2 right panel: the solid and dotted lines are swapped?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and the idea is clearly delivered. The authors have conducted extensive experiments to compare the performance of the proposed methods with the original $\\mu P$.", "weaknesses": "* As the authors discussed, this work lacks a comparison with other proposed methods (e.g., Large et al., 2024).\n* There is no theoretical justification of why choosing a different scaling for the embedding learning rate.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Spotlight)