Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers
Authors: Gavia Gray
Aman Tiwari
Shane Bergsma
Joel Hestness
Venue: NeurIPS 2024 poster (2024)
Keywords: Efficient deep learning
gradient noise scale
critical batch size
language models
Primary area: optimization_for_deep_networks
TL;DR: While using a trick to compute per-example gradients efficiently we discover that normalization layers statistics predict GNS accurately.

Abstract: Per-example gradient norms are a vital ingredient for estimating gradient noise scale (GNS) with minimal variance. Observing the tensor contractions required to compute them, we propose a method with minimal FLOPs in 3D or greater tensor regimes by simultaneously computing the norms while computing the parameter gradients. Using this method we are able to observe the GNS of different layers at higher accuracy than previously possible. We find that the total GNS of contemporary transformer models is predicted well by the GNS of only the normalization layers. As a result, focusing only on the normalization layer, we develop a custom kernel to compute the per-example gradient norms while performing the LayerNorm backward pass with zero throughput overhead. Tracking GNS on only those layers, we are able to guide a practical batch size schedule that reduces training time by 18% on a Chinchilla-optimal language model.

Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers
Gavia Gray Cerebras Systems Toronto, Canada gavia. wawiaargra@ceeebbrss.o net net
Aman Tiwari Subjective London, UK
Shane Bergsma Cerebras Systems Toronto, Canada
Joel Hestness Cerebras Systems Sunnyvale, CA
Abstract
Per-example gradient norms are a vital ingredient for estimating gradient noise scale (GNS) with minimal variance. Observing the tensor contractions required to compute them, we propose a method with minimal FLOPs in 3D or greater tensor regimes by simultaneously computing the norms while computing the parameter gradients. Using this method we are able observe the GNS of different layers higher accuracy than previously possible. We find that the total GNS of contempo- rary transformer models is predicted well by the GNS of only the normalization layers. As result, focusing only the normalization layer, we develop a custom kernel to compute the per-example gradient norms while performing the Layer- Norm backward pass with zero throughput overhead. Tracking GNS only those layers, we are able to guide a practical batch size schedule that reduces training time by 18% on Chinchilla-optimal language model.
1 Introduction
The gradients gathered during the backward pass while training neural network are typically inspected via their Frobenius norm, the magnitude of the vector. This gradient vector may be viewed as the sum of gradients computed over each individual example in the minibatch. Each of these has its own norm. In this work, we develop method to access these norms that works at any scale, for three common layer types in deep learning models: linear, normalization and embedding layers.
One primary application of per-example gradient norm is in estimating the Gradient Noise Scale (GNS) [39], metric that has been shown be useful in training large scale models [9]. The uncertainty of the GNS estimator depends directly on the size of the batch used to compute the small batch gradient norm as shown in Section 2.1. So, the most precise estimate of the GNS is obtained by computing the gradient norms for each example in the minibatch: the per-example gradient norm.
To demonstrate GNS measurement in practice we perform experiments contemporary language model architectures, providing a detailed visualisation the movement of the GNS components throughout training, presented in Section 4. By inspecting these components it was found that the GNS of the model is highly correlated between layer types, which we give an intuition for in Figure 1.
However, the practical utility of measuring GNS with per-example gradient norms is only present if it can be gathered without affecting training time. Focusing on LayerNorm [4] layers, we note the main speed bottleneck is the memory I/O when not implemented as fused kernel. To demonstrate this, we develop a custom kernel to compute both the backward pass and the per-example gradient norms the same time. Using this kernel the throughput overhead of gathering the per-example gradient is zero, even outperforming Py PyTorch's LayerNorm at larger dimensions. We apply this to practical batch size schedule case study in Section 5.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Minibatch
Minibatch
Minibatch
Aggregated
1
4
4
4
4
T
lu
l
l
l3
x la
la
L3
Figure I: Gradient noise scale (GNS) is typically computed by comparing per-minibatch (aggregated- across-layers) gradients to gradients "Aggregated" across minibatches. We estimate GNS with lower variance by making each minibatch single example, and maintain per-layer GNS estimates. We find the magnitude gradients (visualized by the length red arrows) to be consistent across layers, enabling overall GNS to be computed very cheaply using only gradient stats from LayerNorm layers.
To reiterate, the contributions of this work are:
A minimal FLOP algorithm and implementation for computing gradients and per-example gradient norms of linear layers simultaneously.
Observations that the measured GNS for LayerNorm layers is highly correlated with the GNS of the remaining layers.
Development of an example kernel to implement tracking the GNS of LayerNorm layers that does affect network throughput (tokens/sec).
Demonstration of real application of GNS tracking in batch size schedule experiment that obtains an 18% wall-time speedup in training a Chinchilla-optimal [29] LLM.
2 Background
2.1 Gradient Noise Scale
GNS is metric derived from observing second order Taylor expansion of the change in loss function under the following assumption on the noise in the gradient estimate [39],
Gest(o) N ((()), 1 =
(1)
where Gea is the observed gradient, B is the batch size, and 9 the parameters of the model. Here, G' is the unobserved "true" gradient and E is the covariance of the gradient estimate. The Taylor expansion mentioned is,
1 _ ( (r(HE) E[ cGest)) = [(() -Gp2
B
Where is the learning rate and H is the Hessian of the loss. On the right hand side is factor that depends on B. It may be shown [39] that the optimal step size and optimal change in the loss is achieved when B Buoise == tr tr(HE) (HE) /G" HG. Averaging this optimal step over an entire run, and measuring this value by grid search, yields Ben which describes a batch size that meets an optimal tradeoff between cost and training speed. It is shown by analysis and experiment that Broise Beni
As this depends on the Hessian, which is typically unavailable, McCandlish et al. [39] suggest making the assumption that the Hessian is diagonal, which yields
tr(2)
B simple
GTG
Similar algorithms for other layer types described in Appendix B

---

1-4, 5-1
0.025
1-128, s-1
1-8, s-1
|-128, 5=2
0.025
sm1
1-32. s-1
1-128, 5
1n sm1
1-228, 5-16
0.020
GNS stderr
1-128,
GNS stderr
0.015
1-128, 5-64
1-128,
0.015
0.010
0.010
0.000
0.000
10
Figure 2: The variance of the GNS estimator for different Baig (left) and Bsmall (right) sizes. Bbig and Bemall = in legends. Stderr is estimated using a jackknife resampling method for ratio estimators [12]. For the same number of samples processed, smaller small always has lower standard error, while the size of the large batch, Bugg does not affect the standard enror.
To compute simple McCandlish et al. [39] define the unbiased estimators S and ||9|| as:
| =
1
Bsmall |G Bamall?) l12) = GT (Bug ||G Bug 1
(4)
Boig Bamall
(IG
G l12) s tr(
= 1/ 1/ Boig
(5)
where Baig and Bsmall are the batch sizes used to compute the gradients G GBug and G G.an respectively (potentially corresponding Aggregated and Minibatch gradients depicted in Figure 1).
Baulll 11g is trivially computed using the gradients accumulated for the optimizer but ||G I.mnlllle is not. One option is to use the gradients communicated between Distributed Data Parallel (DDP) nodes, but this has two downsides: (1) the variance of the estimate tied to the DDP configuration and (2) the estimate is not available in all training configurations. For example, experiments on single GPU cannot use this method. One can also access the gradients during gradient accumulation, but this similarly depends the training configuration. A full taxonomy of the options for computing ||G ll is provided in Appendix A.
For each observation of ||G Il, we may observe multiple ||G 12- typically Buig/ Bsmall of them.
On each step the estimate of || B.n. 1|2 is therefore mean over Bugg/ Banall samples, whose variance is reduced according to the law of large numbers. However, the GNS is a ratio of the unbiased estimators in Equations and
it may be clear how this affects uncertainty in the GNS estimate. Figure explores this relationship by simulation a setting where the GNS is set to while varying Bug and Bamall- We find it is always better (less uncertainty) use the smallest possible Bsmall to estimate the GNS, while the choice of Boig is irrelevant.
2.2 Efficient Per-example Gradient Norms
Goodfellow [26] proposes trick to compute gradient norms for individual examples in minibatch, which would provide the minimum variance estimate the GNS as described Section 2.1. Neglecting the original derivation, by writing the desired squared norm tensor contraction the trick may be reproduced automatically via einsum path optimization [49, 15]. The tensor contraction for per-example gradient norms, ni. of a linear layer in the 2D setting is,
2
n Σ (w 2.6
Dhsk
=>
TerTriiha Vha. i,k
where v are the activations prior to a linear layer, 1/ are the gradients of the loss with respect to the outputs of the linear layer and w' are the gradients of the loss with respect to the weights of the linear layer.

---

Li et al. [36] extend this trick to the three dimensional case. For inputs X E R BxTx/ and outputs Ye E R K the per-example gradient norm n1 is,
n bik Σ Toistlak) Toustha (XX", Y'
which has O memory complexity in the sequence length T.2. Index sets are b E [1, Bl, [1 & E [ f,
E [1. At some point, the I/O cost of computing the per-example gradient norms by computing the full w/ explicitly will be cheaper. Noting this fact motivated the work in Section 3 and the practical relationship between these resource costs is explored in Section 3.1.
2.3 Related Work
Gradient norms One common motivation for computing per-example gradient norms for differ- ential privacy. By bounding the gradient for any single example, we can ensure each example has a limited impact on the final parameters [45, 36]. Per-example gradient clipping has been performed with convolutional networks [45] and sequential models, e.g., LLMs [36]. These methods allow control over per example gradient norms even when training with large batch sizes. Approaches like these are implemented in the differential-priveer PyTorch, but are less efficient than the methods proposed in this paper. An alternative mechanism to manifest per-example gradient norms is to simply use a batch size of one. While not efficient enough for training large-scale networks, such sequential training may arise in situations such as reinforcement learning, where per-example gradient clipping has also been performed (to improve stability [52]).
library Opacus [56], and have support natively in
Gradient noise scale The Gradient Noise Scale [39] has been widely used for training large-scale neural networks. For example, Brown et al. [9] note the GNS was measured during training and used to guide batch sizing when training GPT-3. Dey et al. [19] mention that operating near the critical batch size, as dictated by the GNS, is important for hyperparameter transfer under the maximal update parameterization [54]. Even when not explicitly mentioned in publications, open source code often implements the GNS (e.g., see codebases [21, 13] for GPT-NeoX [7] and Hourglass Diffusion Transformer [14)).
Measurements similar to the GNS have also been used in a range of prior work to guide batch sizing for minibatch SGD [10, 17, 5, 55]. Chen et al. [11] show experimentally that wider networks can be trained using larger batches; they also establish theoretical connection between wider networks and gradient variance, albeit for simple two-layer networks. In contrast, Shallue et al. [47] found empirically that narrower Transformers scale better to larger batch sizes. Smith and Le [50] propose noise scale based not gradient variance. but on the learning rate, dataset size, and batch size (similar to the notion of temperature in Section 4. Zhang et al. [60] find the critical batch size depends the choice of optimizer. Faghri et al. [22] introduce gradient clustering and stratified sampling approach to minimize minibatch gradient variance, and use this approach as a tool to help understand optimization.
Gradient variance Beyond computing the GNS, our method can support other applications where measuring the distribution of per-example gradients is useful informative. Gradient variance has been used to classify the difficulty of examples [1], which can be used, for example, to surface problematic examples for human auditing. The question of whether gradient distributions tend toward Gaussian in the (central) limit of theoretical significance [50], with implications toward the ability of SGD escape sharp minima and land in wide basins [63, 41, 48). Bounded gradient variance is also assumed in some convergence analysis [8, 62), as noted in [22].
Perhaps the most familiar of gradient variance is of course in adaptive optimizers like Adagrad, Adam, and others that reduce step sizes in Bigh--raadentt-ooe noise directions [20, 57, 46, 33, 44]. Hilton et al. [28, App. C] directly relate Adam second moment statistics to a component-we
version of the GNS. Optimizers typically estimate gradients jointly across training steps and minibatches, however vSGD [46] leverages separate components for gradient momentum and for gradient variation across samples. Zhang et al. [61] find the variance of gradient norms across examples predictive of whether
"This specific Einstein contraction is not used by Li et al. [36] but appears in the Backpack library [15] We provide the vector algebra contraction path chosen by Li et al. [36] on the right.

---

vanilla SGD outperforms adaptive optimizers, however recent work has shown Adam to outperform SGD even in the (noise-free) full gradient descent setting [34, 35].
3 Simultaneous Per-example Gradient Norms
As described in Section 2, computing GNS requires small batch gradient norms. Typically, these may be gathered during gradient accumulation or DDP communication. However, these methods are not universally applicable and may not be available in all training configurations. In this section we describe method for baking the computation of the per-example gradient norms into the computation graph, making it universally applicable. The typical tensor contraction used to compute the backward gradient in a linear layer using the input activations, x, and gradients, g, is,
w//1- = Σ A........I
in other words, sum over vector outer products for every vector the trailing dimension. In principle, it is possible to access the intermediate tensor containing the batch dimension Waat = Σ &bb...9b.... This allows to compute the per-example gradient norms with FLOPs scaling the same rate as the normal, non-per-example backward pass (Figure 3), albeit at increased I/O cost due to having to materialize the intermediate tensor.
generic algorithm to compute the per-example gradient norms simultancously with the weight gradient in standard linear layer is provided in Algorithm using einsum for readability and portability. 4 The reason for the correction in step 4 can be seen by considering the gradient of loss function L with respect to the weights a single example b, wb.
V wa L(as)
- V L(xs),
1
B
computing the squared norm of this will therefore contain a factor of 1/B², which must be corrected for.
Algorithm Linear Layer Simultaneous Per-Exampll Gradient Norm Computation
Require: gradient tensor g of shape (B,.., L), input activation tensor of shape (B, K)
Ensure: weight gradient tensor w' of shape (K, L),, mean of per-example squared norms l/ b...l -> bkl',x,, g)
I: w + einsum
2: Sw < einsum(*bAe bkl + b² w/
3: w' <- einsum( (bbbl kdl', w()
4: [|w/112 < 1/B x einsum(sw:: $b + x B² # reduce by mean then apply correction
5: return w', l//l
3.1 FLOPs and I/O Costs
The computational cost of computing per-example gradient norms can be broken down into FLOPs, in Figure 3, and I/O, in Figure 4, with matrix multiplication current devices being potentially bottlenecked by both. We estimate ideal FLOP and DRAM I/O costs, assuming optimal reuse of data loaded from DRAM into SRAM with recomputation. In practice, duplicate computation may be used to improve wall-clock time and to fit within hardware limitations of the amount of shared memory available. We compare here against the efficient per-example gradient norm method described by Li et al. [36), which the authors note is only efficient (in terms of I/O cost) when 2T < PD, where T is the sequence length, P is input and D is output dimension of the linear layer. This bound is discussed further in Appendix E.
In terms of FLOPS, Figure 3 shows the simultaneous per-example gradient norms almost always preferable, only being more expensive for very short sequence lengths in small models. The reason for this is shown on the right hand side; the number of FLOPs required to compute the simultaneous per-example gradient norms independent of the sequence length.
'A complete taxonomy for small batch gradient computation given Appendix A. 'Additional algorithms for Embedding and LayerNorm layers described in Appendix B.
5

---

2.5
Semullammetu
GNS FLOPs /Model FLDPs
Siquurooo Langth
256
100
2.0-
- 1024
2045
+ 4085
134
3.0
0.-
100
Paramoters
Porametters
Figure 3: FLOP cost of computing per-example gradient norms. (Left) Total FLOP cost. (Right) Proportional cost versus model forward and backward pass. The FLOP cost of Simultancous per-example gradient norms strictly dominant to alternative methods (left) and the ratio this additional cost to the FLOP cost of processing the entire model does not depend context length (right).
Ui only
10 (proportion model forward
200
Sequence Longth
236
1024
2048
4096
0192
10a
1636
1771
65536
10ch
10a
10
Parameters
Figure Total I/O cost of computing per-example gradient norms, assuming gradients and parameters stored with 4 bytes of precision. The relative IO cost of Simultancous per-example gradient norms is less than et al. [36] for very long contexts for all model scales, approximately equivalent for models of 10B parameters and 4096 context length, and higher for shorter contexts with larger models. The 1O cost of LN (LayerNorm) per-example gradient norms alone much lower than either method.
The I/O cost shown in 4 illustrates tradeoff in computing the per-example gradient norm. The simultaneous method is more expensive at large model sizes with short sequence length because must large intermediate tensor.
To estimate model flops, we use PoTohh''s FLOPCounterMode, which only measures the FLOPs in matrix multiplications and attention computation, however these make up the vast majority of the FLOPs Transformer model.
Gradient Noise Scale in Transformer Language Models
Using the methods described in previous sections to measure per-example gradient norms and estimate the GNS, we perform experiments 111M parameter Chinchilla-optimal language model [19, 29] using the Open WebText dataset [24]. As the prior work was performed on Pile [23]. Appendix C. describes experiment check the optimality of the Chinchilla model this dataset. We also found Flash attention led to numerical instability, which we were able to mitigate with an architectural modification described in Appendix C.2.
'The code to replicate these experiments may be at https:://// http//gtttttttttteeeeeeeeeeeee com http:/ubbbcommmrrbbbasseeeerrrr nanoGNS tree http//teeeeeeeeeeeeeeeeeeeeeetet exact.

---

MLP and Attention
600
10a
500
Pott 20000
400
Tr(23
300
200
20000
100
100
100
Combined
10
10
10a
10
10
2000
4000
6000
8000
||G]|
Iteration
HSNS 10000
LayerNorm
350
10-2
300
250
200
150
10 Z00
100
10-4
Combined
0
10-4
10
10
10
2000
8800
||G||
Iteration
Figure 5: GNS phase plot: Linear/Embedding layers are separated from LayerNorm layers by row. Component estimators of Equations 4 and 5 shown (left) with the GNS over the course of training on the (right).
All experiments computed per-example gradient norms for all layers in the model with the exception of the performance results of Sections 5.1 and 5.2, which only computed per-example gradient norms for the normalization layers. Each experiment was run on Nvidia A10 GPUs, in either 12 or 24 hours depending on the precision used, 6 Float32 respectively. We used the nanoGPT6 codebase with the layers described Section 3 added.
Having an accurate estimate of the GNS statistics || and allows us to visualize the movement of both in phase space during training as shown in Figure 5. LayerNorm layers are separate from the rest of the network because their statistics are much smaller and to illustrate how the resulting GNS estimates on the right track each other. To observe these trends in another training regime, see Figure 14 in Appendix D.I.
4.1 The Temperature of Training
McCandlish al. [39, App. C] observed that the GNS measurement depends on the batch size and learning rate used in training. In fact, from the derivation outlined in Section 2. the gradient noise scale is only well-defined the optimal learning rate. Using toy model of quadratic loss function, they observed that the GNS should be inversely proportional to the temperature, T. ratio of batch size B to learning rate E:
1
B
Broise O Bsimple
T
This enables a testable prediction that the GNS will increase with increasing batch size or with descending learning rate. This prediction was found to accurately describe experiments on a small
'nttps://github
httttp///gtthhhb.oo//aaaaaaaaaa com/karpatthy/nnooor
7

---

10a
original
e/16, 8 2/4,
GNS
c/4,
10
Tokens processed
1e8
Figure 6: During the middle of training a 11 parameter language model on Open OpenWetText, the learning rate, batch size, B varied, restarting the run from the same point. This Figure replicates an experiment from McCandlish et al. [39] showing how varying the ratio causes changes in the measured GNS, but here only due changes in the learning rate. Changes in the batch size do not have the predicted effect.
convolutional model on the SVHN dataset. We repeat it here in the setting described above in Figure To match the results of McCandlish et al. [39], all interventions tested should yield the same result. We find the GNS does indeed react predictably to changes in the learning rate, but the reactions to changes in the batch size not predicted by the theory.
4.2 GNS Correlates Between Layer Types
Inspection of Figure suggests the LayerNorm layers produce similar GNS, when combined, as the total GNS of the model. Before describing how quantify this relationship we must first note that the unbiased estimators ||9 2 and are noisy. All GNS figures presented in this paper and other work smooth both of these estimators, typically with an Exponential Moving Average (EMA) filter, before computing the GNS ratio.?
So, when quantifying the relationship between the GNS of different layers, it must be compared for different smoothing factors. Here, we show the regression coefficients with respect to the alpha of the EMA filter in Figure 7. The results show that the GNS of the LayerNorm and Attention layers are highly predictive of the total GNS of the model. In both cases, the slope is approximately 1.4, meaning the total GNS is approximately 1.4 times the GNS of the LayerNorm or Attention layers.
Comparing the quality of this fit versus the quality of prior work's overall fit of the GNS to the critical batch size (measured empirically) [39], the quality seems acceptable and we do not need to apply this 1.4x correction factor, rather just note that the true Bein may be greater than the measured Biimle--
Batch Size Scheduling
We focus on two concerns that affect the practicality of batch size scheduling, First, measuring the appropriate batch size without incurring any additional training time. We find this is possible with the method described in Section 5.1. Second, whether batch size scheduling is effective in practice. We find it can offer significant savings in the required number of tokens processed in Section 5.2.
5.1 Universal GNS with Zero Overhead
Capturing GNS estimate for linear layer is powerful, but efficiently doing presents challenge. Such estimate requires accumulating per- per-example gradients of hidden size? across the sequence dimension, compared to just hidden size with LayerNorm. This increased size requires using more complex reductions in the kernel, rather than simple warp reduction followed by shared-memol atomic reduction with a final atomic global reduction (as we can implement for LayerNorm per- example gradients within shared memory). In addition, linear layer kernels are already highly optimized and require using advanced techniques to keep GPU tensor cores fed with data, so
The results described in Figure 5 are explored 1.3B parameter scale Appendix D.3.

---

1,00
0.99
GNS
slppe
0.90
Attention LayerNorm
0.37
Embeddling
0.96
Tot
0.96
0.98
1.00
0.66
0.98
1.00
alpha
alpha
Figure 7: Regression of total GNS using the GNS of each layer type. (Left) GNS of each layer type and the total GNS are plotted against the number of tokens processed for varying EMA alpha settings. (Center & Right) The slope and Pearson's 's correlation coefficient of the regression of the total GNS against the GNS of each layer type, respectively, as a function of the same EMA alpha values. The total GNS (black) on the left is predicted well by individual layer types indicated by the correlation coefficients (right), however the type with slope closest to is LayerNorm (center), only overestimating the GNS by less than 40% across EMA alpha values.
Implemientation
LayerNorm (PyTorch 2.4.0)
Karnal
Average Fwid Bwd Time
150
50
Dimensionalitt
Figure 8: Comparison of average time taken for LayerNorm forward and backward pass with gradi- ent accumulation when using PyTorch's native implementation versus custom kernel computing per-example gradient norms in tandem. Measured on an Nvidia H100 GPU.
combining such kernel with per-example gradient computation with its own memory overheads and corresponding available bandwidth reduction would be a difficult undertaking.
We thus implemented a LayerNorm-speccfic CUDA kernel that also captures GNS. In experiments with language models at different scales, illustrated in Figure 8, find this kernel has practically zero overhead compared to PyTorch 's LayerNorm implementation. The complete source code for this kernel is provided with the accompanying code for this paper".
5.2 Case Study: Batch Size Schedule
As case study we continue with the 111M parameter language model on Web described above. Over three seeds, we run both fixed batch size and a batch size schedule that increases linearly with the number of tokens processed to the original batch size. We vary the batch size during training by varying the number of gradient accumulation steps.

---

1c8
4
10
Tokens saved
3
Train loss
6 10
2
4 10
1
Original
With batch size schedule
x 10
0
10
10
10
10
10a
10a
x 10
4 x 4x
x xx000
Tokens processed
Train loss
Figure 9: (Left) Linear batch size schedule tracking the GNS over 2.2 billion tokens processed. Loss is plotted over smoothed range from 3 runs using different Seeds. (Right) The number of tokens saved over the fixed batch size run achieve the same loss.
The results of this experiment are shown Figure 9. The left plot shows the progression of the loss for both models, with the range of values captured over different seeds. The mean loss for the linear batch size schedule leads the fixed batch size throughout training. On the right, this lead is quantified by interpolating the number of tokens saved to achieve the same loss. The precise schedule used is shown in Figure 15 in Appendix D.2.
6 Limitations
In this paper, we only studied Transformers, which include Normalization sub-layers natively. While Transformers are ubiquitous in machine learning, there are many models, including variations of RNNs, CNNs, and state-space models, that do not use such layers conventionally. However, we note LayerNorm could be added these networks with very little overhead (in fact, the desire to normalize activations in RNNs was one of the original motivations for developing LayerNorm; application of batch normalization [30] RNNs was "not obvious" [4]). Nevertheless, investigating LayerNorm-based GNS in these other models requires further work.
Our work is also part of efforts improve efficiency and address the increasing costs of training and tuning large neural networks [6]. We provide both a more-efficient technique for computing the GNS, and also, by enabling use of GNS statistics, we support compute-efficien training recipes, such as use of dynamic batch sizes. While some have argued that hyperscalers may re-invest any efficiency savings into ever-larger models [42], for academic researchers, such savings could allow pushing the state-of-the-an while still getting results in a reasonable timeframe. Recent efforts enable frontier-model perfor-nerf within academic budgets are encouraging, both to reduce memory [38, 18] and save compute [37, 2]. Of course, even for such economical approaches, "extensive hyperparameter search" may still be required [31]. There is growing awareness that hyperparameter tuning has negative impact on equity in Al research, tuning success depends directly on researcher finances [51]. A correlated trend is to use better training measurements (such as gradient noise in batch and step size optimizers (Section 2.3)) to reduce dependence hyperparameters, and in this way we hope our work can also ultimately improve research equity.
7 Conclusion
This work set out to provide practical method for computing the per-example gradient norms necessary compute the GNS independent of the training configuration. In the process we discovered that not all the layers are necessary for practical estimate of the GNS and that the per-example gradient norms can computed for the normalization layers with zero overhead. This enabled practical experiments, such as batch size schedule and replicating prior GNS observations. We are hopeful that democratising access GNS statisties, on any device, will enable subsequent discoverics.
10

---

References
[1] Chirag Agarwal, Daniel D'souza, and Sara Hooker. 2022. Estimating example difficulty using variance of gradients. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 10368-10378.
[2] Sotiris Anagnostidis, Gregor Bachmann, and Thomas Hofmann. 2023. Navigating Scaling Laws: Accelerating Vision Transformer's Training via Adaptive Strategies. arXiv preprint arXiv:2311.03233 (2023).
[3] Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. 2022. High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation. arXiv:2205.0144 [stat.ML] https:/ paops://arxxe org/ abs//2/5..22// 01445
[4] Jimmy Lei Ba. Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. Layer Normalization. MtXiv:1607.064500 [stat.ML]
[5] Lukas Balles, Javier Romero, and Philipp Hennig. 2016. Coupling adaptive batch sizes with learning rates. arXiv preprint arXiv:1612.05080 (2016).
[6] Emily M Bender, Timnit Gebru, Angelina McMillan Majonr and Shmargaret Shmitchell. 2021. On the dangers of stochastic parrots: Can language models be too big?. In Proceedings of the 2021 ACM conference on fairness, accountability, and transparency. 610-623.
[7] Sidney Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, Usvsn Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. 2022. GPT-NeoX-20B: An Open-Source Autoregressive Language Model. In Proceedings of the ACI. Workshop on Challenges & Perspectives in Creating Large Language Models.
[8] Léon Bottou, Frank E Curtis, and Jorge Nocedal. 2018. Optimization methods for large-scale machine learning. SIAM review 60, (2018), 223-311.
[9] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Herbert-Vosss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language Models Few-Shot Learners. arXiv:2005. 14165 [cs.CL] https: /wwwiv.vvvvvvv////////..///aav.a org/ abs 2005. 14165
[10] Richard H Byrd, Gillian M Chin, Jorge Nocedal, and Yuchen Wu. 2012. Sample size selection in optimization methods for machine learning. Mathematical programming 134, (2012), 127-155.
[11] Lingjiao Chen, Hongyi Wang, Jinman Zhao, Dimitris Papailiopoulos, and Paraschos Koutris. 2018. The effect of network width on the performance of large-batch training. Advances in neural information processing systems 31 (2018).
[12] Denis Choquet, Pierre L'Ecuyer, and Christian Léger. 1999. Bootstrap Confidence Intervals for Ratios of Expectations. ACM Trans. Model. Comput. Simul. 9, 4 (oct 1999), 326-348. https://doi. 10g/10.1o
[13] Katherine Crowson. 2024. k-diffusion. GitHub repository.
[14] Katherine Crowson, Stefan Andreas Baumann, Alex Birch, Tanishq Mathew Abraham, Daniel Kaplan, and Enrico Shippole. 2024. Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers. arXiv preprint arXiv:2401.I 1605 (2024).
[15] Felix Dangel, Frederik Kunstner, and Philipp Hennig. 2020. BackPACK: Packing more into Backprop. In International Conference on Learning Representations. attps://openreevee.

---

[16] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. FlashAttention: Fast and Memoor-Efficient Exact Attention with 1O-Awereness. arXiv:2205.14135 [cs.LG] https:///arae org/ abs.//55.22222... 14135
[17] Soham De, Abhay Yadav, David Jacobs, and Tom Goldstein. 2016. Big batch SGD: Automated inference using adaptive batch sizes. arXiv preprint arXiv:1610.05792 (2016).
[18] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. 2023. QLoRA: Efficient finetuning of quantized LLMs. Advances in Neural Information Processing Systems 36 (2023),
[19] Nolan Dey, Gurpreet Gosal, Zhiming, Chen, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestness. 2023. Cerebras-GPT: Open Compute-Optimal Lan- guage Models Trained on the Cerebras Wafer-Scale Cluster. arXiv:2304.03208 [cs.L.]
attps://arxis org/ aataaba////// 03208
[20] John Duchi, Elad Hazan, and Yoram Singer. 2011. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research 12, 7 (2011).
[21] EleutherAI. 2024, GPT-NeoX. repository.
GitHub
[22] Fartash Faghri, David Duvenaud, David J Fleet, and Jimmy Ba. 2020. A study of gradient variance in deep learning. arXiv preprint arXiv:2007.04533 (2020).
[23] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2020. The Pile: An 800GB Dataset of Diverse Text for Language Modeling. arXiv:2101.00022 [cs.CL]
[24] Aaron Gokaslan and Vanya Cohen. 2019. OpenWebText Corpus. http://Skylion007.
gi github.io//ppp wwt/ppennibbeettorpppt
[25] Alicia Golden, Samuel Hsia, Fei Sun, Bilge Acun, Basil Hosmer, Yejin Lee, Zachary DeVito, Jeff Johnson. Gu- Yeon Wei, David Brooks, and Carole-Jean Wu. 2024. Is Flash Attention Stable? arXiv:2405.02803 [cs.LG]
[26] Ian
Goodfellow. 2015.
Efficient Per-Example Gradient Computations.
arXiv:1510.01799 [stat.ML] https://arxiv org/ abs/ wwspaa1o......... 1510.01799
[27] Gavia Gray, Anshul Samar, and Joel Hestness. 2023. Efficient and Approximate Per-Example Gradient Norms for Gradient Noise Scale. In Workshop on Advancing Neural Network Training: Computationnal Efficiency, Scalability, and Resource Optimization (WANT@ NeurIPS 2023). net/ Forforormuiii "rrrmm?idxxINNMMAPPP
[28] Jacob Hilton, Karl Cobbe, and John Schulman. 2022. Batch size-invariance for policy optimiza- tion. 1:Xiv:2110.006: [cs.LG]
[29] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. 2022. Training Compute-Optimal Large Language Models. arXiv:2203. [cs.CL]
[30] Sergey loffe and Christian Szegedy. 2015. Batch normalization: Accelerating deep network training by reducing internal covariate shift, In International conference on machine learning. pmlr, 448-456.
[31] Peter Izsak, Moshe Berchansky, and Omer Levy. 2021. How to train BERT with an academic budget. arXiv preprint arXiv:2104.07705 (2021).
[32] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. 2024. Analyzing and Improving the Training Dynamics of Diffusion Models. arXiv:2312.0266 [cs.CV] attps://art
org/ abaaabb2/1122 02696
[33] Diederik Kingma and Jimmy Ba. 2014. Adam: method for stochastie optimization. arXiv preprint arXiv: 1412.6440 (2014).
12

---

[34] Frederik Kunstner. Jacques Chen, Jonathan Wilder Lavington. and Mark Schmidt. 2023. Noise not the main factor behind the gap between SGD and Adam on transformers, but sign descent might be. arXiv preprint arXiv:2304.1333 (2023).
[35] Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. 2024. Heavy-tailed class imbalance and why Adam outperforms gradient descent on language models. arXiv preprini arXiv:2402. (0000.9944 19449 (2024).
[36] Xuechen Li, Florian Tramèr, Percy Liang, and Tatsunori Hashimoto. 2022. Large Language Models Can Be Strong Differentially Private Learners. arXiv:2110.05679 [cs.LG] https: //arxiv. wwwp:///g///aaaeeee
abs/ /b1// 05679
[37] Xianhang Li, Zeyu Wang, and Cihang Xie. 2023. CLIPA-v2: Scaling CLIP Training with 81.1% Zero-shot ImageNet Accuracy within $10,000 Budget; An Extra $4,000 Unlocks 81.8% Accuracy. arXiv preprint arXiv:2306. 15658 (2023).
[38] Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. 2023. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems 36 (2023), 53038-53075.
[39] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. 2018. An Empirical Model of Large-Batch Training. arXiv:1812.0611 [cs.LG] fttps://arxiv..e org/ h/tgg//b////// abs//822.
06162
[40] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. 2018. Spectral Normalization for Generative Adversarial Networks. arXiv:1802.05957 [cs.LG] https:
// //arxiv. www.oorggg//baaoe.ee. ab0//022..11.//////////////// 05957
[41] Thanh Huy Nguyen, Umut Simsekli, Mert Gurbuzbalaban, and Gaël Richard. 2019. First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. Advances in neural information processing systems 32 (2019).
[42] David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. 2021. Carbon emissions and large neural network training, arXiv preprint arXiv:2104 10350 (2021).
[43] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAl blog 1, 8 (2019), 9.
[44] Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. 2019. On the convergence of Adam and beyond. arXiv preprint arXiv: 1904:09900099333 (2019).
[45] Gaspar Rochette, Andre Manoel, and Eric W. Tramel. 2019. Efficient Per-Example Gradient Computations in Convolutional Neural Networks. 11Xiv:1912.06015 [cs.LG] attps://arxiv. org/abs/19tt
[46] Tom Schaul, Sixin Zhang, and Yann LeCun. 2013. No more pesky learning rates. In International conference on machine learning. PMLR, 343-351.
[47] Christopher Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl Dickksein, Roy Frostig, and George E Dahl. 2019. Measuring the effects of data parallelism neural network training. Journal of Machine Learning Research 20, 112 (2019), 1-49.
[48] Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. 2019. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning. PMLR, 5827-5837.
[49] Daniel G. A. Smith and Johnnie Gray. 2018. opt_einsum A Python package for optimizing contraction order for einsum-like expressions. Journal of Open Source Software 3, 26 (2018), 753. petpss///oo... org/ 10. 21 ooo///////.//.o/
[50] Samuel Smith and Quoe V Le. 2017. A Bayesian perspective on generalization and stochastic gradient descent. arXiv preprint 17X10:1710.06455) (2017).
13

---

[51] Emma Strubell, Ananya Ganesh, and Andrew McCallum. 2019. Energy and policy considera- tions for deep learning NLP. arXiv preprint (1Xiv:996600224 (2019).
[52] Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Mare Lanctot, and Nando Freitas. 2016. Dueling network architectures for deep reinforcement learning. In International conference machine learning. PMLR, 1995-2003.
[53] 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, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. 2023. Small-scale proxies for large-scale Transformer training instabilities. 2XIv:2309.14322 [cs.LG] https: / //arxiv.org htte//bb/233099 23092330022 14322
[54] Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. 2021. Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. In Advances in Neural Information Processing Systems.
[55] Dong Yin, Ashwin Pananjady, Max Lam, Dimitris Papailiopoulos, Kannan Ramchandran, and Peter Bartlett. 2018. Gradient diversity: a key ingredient for scalable distributed learning. In International Conference on Artificial Intelligence and Statistics. PMLR, 1998-2007.
[56] Ashkan Yousefpour, Igor Shilov, Alexandre Sablayrolles, Davide Testuggine, Karthik Prasad, Mani Malek, John Nguyen, Sayan Ghosh, Akash Bharadwaj. Jessica Zhao, Graham Cormode, and Iya Mironov. 2022. Opacus: User-Friendly Differential Privacy Library in PyTorch. arXiv:2109. 122998 [cs.LG]
[57] Matthew D Zeiler. 2012. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701 (2012).
[58] Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Josh Susskind, 2023. Stabilizing Transformer Training by Preventing Attention Entropy Collapse. arXiv:2303.06296 [cs.LG] httpp::/arxiv.org/// 06296
[59] Biao Zhang and Rico Sennrich. 2019. arXiv:1910.07467 [cs.LG]
Root Mean
Square
Layer Normalization.
[60] Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George E. Dahl, Christopher J. Shallue, and Roger Grosse. 2019. Which Algorithmic Choices Matter at Which Batch Sizes? Insights From Noisy Quadratic Model. arXiv:1907.0411 [cs.LG]
[61] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. 2020. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems 33 (2020), 15383-15393.
[62] Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi- Zhi-Quan Luo. 2022. Adam can converge without any modification on update rules. Advances in neural information processing systems 35 (2022), 28386-28399.
[63] Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. 2018. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint 1803 18030000 00195 (2018).
14

---

A Taxonomy
Gray et al. [27] included an prior version of this taxonomy in their work.
The following taxonomy describes the different methods available compute the I| ll necessary to compute the GNS as described in Section 2.1. "Gradient norm cost" below refers to the cost of computing the norm of the gradient for all parameters in the model, which is typically orders of magnitude smaller than the cost of forward or backward passes.
Microbatch: multiple Pammll 1|2 are computed over set of microbatches
DDP: Each ||G' 11 is computed before gradients are communicated between DDP
nodes [39].
Pros: Only gradient norm cost.
Cons: Variance tied to number of DDP nodes (see Figure 2), can't be used node.
-Sequential: Each || 2 are computed sequentially during gradient accumulation.
-
Pros: Only gradient norm cost.
Cons: Variance tied to the number of gradient accumulation steps.
Subbatch: During gradient accumulation, select ||G 1| partway through.
Pros: Only gradient norm cost, easy to implement.
Cons: Higher variance than Microbatch as ||G 12 is not averaged.
Per-example:
Pros: Independent of gradient accumulation or DDP configuration, minimal variance.
Exact:
* |G 11 is computed directly by the per-example gradient trick [26, 36].
Pros: Minimal cost in 2D regime.
Cons: Redundant computation required in 3D regime.
|G Bamall is computed in tandem with the parameter gradients using the method described in Section 3.
Pros: No redundant computation.
Cons: Expansion in memory causes slowdowns as described in Section 3.1.
Approximation: [| 2 is approximated by assuming input activations are normally distributed with mean zero [27].
Pros: Fewer FLOPs than Exact methods.
Cons: Not exact.
All the methods described above can be measured either online or offline. The description above focuses the online case; i.c. measuring the gradient norms during training. To use these methods offline: run the models without performing weight updates and measure gradient norms the same way. The estimators of Equation 4 and 5 can then be aggregated using by using a method to estimate measurement uncertainty such as the jackknife mentioned Figure (described in the context of GNS by Gray al. [27, App.B)). This can be useful to estimate how long to run the offline estimate.
rather than EMA
B Additional Simultaneous Per-Example Gradient Norm Computations
Algorithms 3 and 2 describe the process for computing the per-example gradient norms for the embedding and LayerNorm layers, which are typically the remaining layers in Transformer models. RMSNorm [59] is practically identical to LayerNorm in this case because the parameters the gradient is computed wrt are in the affine transform, which is the same in both layer types.
C Language Model Experiment Details
As mentioned in the text, the code to run the experiments described in this paper can be found exxcct.
https://gi http///gtttttt/////// com
15

---

Algorithm 2 Layernorm Simultaneous Per-Example Gradient Norm Computation
Require:
gradient tensor g of shape (B..., K), input activation tensor x of shape (B, K) Ensure: gamma gradient tensor > of shape (K; mean of per-example squared norms 1/761 gradient tensor 3' of shape (K,) mean of per-example squared norms ||,%/11
1: 7 + einsum( (bb... b...k bk', x,
2: 8, + einsum (bbk bb + b', 76
3: 7 cinsum ('bk + k', 7o
4: 1 < 1/B x einsum 0b + 1) x B # reduce by mean then apply correction
5: Bb < einsum(( h + bk'.g)
6: Sy + einsum( bb + b', ())
7: 3 < einsum( (bbl: + k'
8: 118.12 < 1/B einsum(s,, bb+ ') B reduce by mean then apply correction
9: return y', 11 B', ||
Algorithm 3 Embedding Layer Simultancous Per-Example Gradient Norm Computation
D), Require: gradient tensor g of shape D), input tensor x of shape (B, T), vocabulary size V
Ensure: weight gradient tensor w' of shape (V, D), mean of per-example squared norms ||w6
I: o onehot(x. V/)
2: w + einsum( blv, bld - bod' o. o.
3: Si é einsum( crbsum(-bud-a + b¹, w(2)
4: w' < einsum bbod -> ud', w¹)
5: [!wállž 1/B einsum(su.. (S.: 8b ') B² # reduce by mean then apply correction
6: return w', //w/11
C.1 Optimality on OpenWebText
We chose to use the Cerebras-GPT [19] recipes for experiments as they are designed to be Chinchilla optimal. This means that each model size should achieve the lowest possible loss for a given FLOP budget [29]. However, these recipes were tuned on the Pile dataset [23] and we used the OpenWebText dataset [24] so that results could be replicated (Pile is no longer publicly available).
To verify that the training protocol is optimal on OpenWebTers illustrate how the performance would vary as we vary the size and total tokens trained on. Model size was varied by changing the hidden size: the 70M model has a hidden size of 576, the 111M model has hidden size of 768 and the 161M model has hidden size of 960. The token budget for each model size was chosen to keep the total FLOPs constant.
we performed small study to
The learning rate was varied to observe a minima in the loss each model scale. The results shown in Figure 10. While we found that the learning rate may be increased overall, the 11 IM model was found to have the lowest loss of the three models. From these results we conclude that the training protocol is optimal within this range of model sizes and we assume 1 IIM is good enough. In other words, better model might exist between 70M and 161M parameters for this FLOP budget but it isn't outside of this range.
C.2 Flash Attention Numerical Instability
The experiments described in Sections 4 and 5.2 involve Chinchilla optimal language models at 111M scale [19]. These experiments were replicated according to the published information. We encountered diverging runs when executing in 6 Automatic Mixed Precision (AMP) consistent with the default settings in nanoGPT. These experiments were executed on NVIDIA A10 GPUs for accessible replication at small scale. By ablation it was found that these runs would diverge:
Regardless of batch size schedule
16

---

3.400
111M params
70M params 161M
3.350
Validation
3.325
X
3.300
3.275
3.250
1.1.1:::
min: 1216
3.25
10
10-3
Learning Rate
Figure 10: The loss of models trained on Open Text with 70M, 111M and 161M parameters. The learning rate varied to find the minima in the loss each model scale. The optimal learning rate for each model size is annotated.
Regardless of hyperparameters: learning rate, weight decay, LayerNorm epsilon or Adam epsilon [53]
When using PyTorch's AMP in 66 precision
When using Flash attention [16, 25]
This was surprising because prior work had trained these models successfully [19]. In that work the model was also trained using 6 AMP precision, but was trained on Cerebras CS-2 system. Due to this difference, suspected the issue was due to difference between the efficient attention kernel and the Flash attention kernel in PyTorch.
By inspecting the histograms of weights and biases in the Query, Key, Value (QKV) projection during training, we found that range grew the fastest in block (the second block in the model). In addition, we observed that the histogram of the query and key projection weights became bimodal as the gradient norm diverged. This illustrated in Figure 11. Further analysis of checkpoint taken at this point in training focused on the difference between gradients computed using the flash attention kernel and the nanoGPT pure PyTorch attention implementation using float32 precision. At initialization the gradients were not significantly different but at the point of divergence there was significant difference coinciding with increased parameter norms in that layer.
To replicate the issue from scratch, we came with simulation from a generic initialization. Inspired by the teacher-student experiment protocol proposed by Ba al. [3] (although otherwise unrelated) we set up learning task with a "teacher" and "student" model with the same architecture. Both networks begin with the same weights but add small amount of noise to the teacher's weights. The student is trained to match the teacher's output. After experimenting with hyperparameters we were able to replicated the divergence seen during training". as illustrated in Figure 12.
Using this isolated simulation we were able to test different methods to mitigate the divergence. Karras al. [32] suggested that cosine attention could address similar divergences attributed to self-attention. In Figure 13 we replicated the experiment described in Figure 12 using cosine attention and found that the divergence no longer occurred.
Separately, experimenting with precision ablation found that if float32 precision was used only in block 1 (2nd) then the divergence would also not occur. Based on this and the above, found the following two architectural mitigations for the divergence, only block (2nd):
Use cosine attention, i.e. normalize the query and key head vectors before self-attention. OR
9The code for this experiment is available at //gis://gist... g1 thub. com haam/gaviag--.. cerebras/ #77aeB9cc22985ssss999962225522.2
17

---

Query Weights
Query Biases
30000
Stet: 0.1207
Mean: 0000
Std: 0.2116
Meen: -0 0.0062
Frequency
Frequency
20000
15000
5000
0.2
0.0
0.4
Weight Valuue
Key Weights
Key Biases
M 0000001
Mean: 0.0178
40000
Sket 0.1019
Frequency
30000
Requency
20000
20
10000
0.4
0.2
0.0
0.2
0.4
1.0
0.5
0.0
0.5
1.0
We
Value Weights
Value Blases
Mean: 0,0033.
100000
Stdl
Froquoncy
80000
Frequency
60
20000
0
0
-022
-0.1
0.0
0.1
0.000
8.005 0.000
0,005 0.010 0.015 0.020 0.025
Woight Value
Blas Value
Output Weights
Output Biases
12:0000
Mean: 0000000
Mean:
Std:: 0.0191
Std: 0.0126
100000
40
Frequency
Frequen
20
40000
20000
Weight Value
Blas Valuee
Figure 11: Histograms of weights and biases for the 11 experiment described in Sections 4 and 5.2 from the attention block containing the QKV projection, self-attention and output projection layers. The histograms for the query and key projection weights and biases bimodal while the value projection weights and biases are not.
18

---

Bias Norms
Distances to Tea chtrr
30
Teacher
20
Flash Teacher
Flash
Flash Tea cher
Non-Flash
Bias No Norm
Distance
15
26
10
22
1000
2000
2000
Iterations
Iterations
Flash to Non-Flash Distance
leacher Distance Difference
20
Distance Difference
15
15
Distance
10
10
500
1000
1500
2000
500
1000
1500
2000
Iterations
Iterations
Two "student" networks, identical the "teacher" network except for the addition of small amount of noise to the teacher's QKV projection bias. As training progresses, the student using Flash attention diverges for the same inputs. Plots are, clockwise from top left: "Bias Norms" shows the norms of the bias layer in each of the networks, "Distances to Teacher" shows the L2 distance from each student to the teacher. "Flash to Non-Flash Distance" shows the L2 distance between student using Flash attention and not, "Teacher Distance Difference" is the difference between the distances to the teacher for both cases.
Figure 12:
5.2
20.8
Blas Norm
Distance
5.0
Teacher
20.4
Flash
4.9
Non-laash
20.2
4.8
Flash Teacher
Nan-Filsh to Teacher
19.8
4.6
300
400
500
Iterations
Iterations
Flash to Nan
Nan-Flash Flash Distance
Teacher Distance Difference
0.0000
0.020
Distance Difference
Distance
0.015
0.010
-0.0004
0.005
-0.0006
100
200
300
400
500
100
200
300
400
500
Iterations
Iterations
Figure 13: Replication of the experiment described in Figure 12 using cosine attention instead of Flash attention. The divergence observed no longer occurs.
19

---

Use spectral normalization [40] on the QKV projection.
Critically, only modifying a single layer does not affect the throughput of the model, the observed Model FLOPs Utilization (MFU) did not decrease by more than 1% in either case. Both of these bound the norm of the query and key head vectors prior to attention. Spectral normalization achieves this because the QKV projection is preceded by a LayerNorm layer. Using this mitigation on the 11 model allowed the experiment be replicated on an NVIDIA A10 GPU and we observed the same behaviour as running the model more slowly in float32 precision.
Similar divergences unable to verify that it is the same problem. Wortsman et al. [53] discuss how to build similar experiments to those described above but do not investigate flash attention specifically. Golden ct al. [25] investigate the numerical stability of Flash attention but neglect to demonstrate a failure mode that affects real training runs. Zhai et al. [58] focus the numerical stability of attention in general and propose a similar mitigation (their method, o Reparam, is a scaled version of spectral normalization) but do not investigate flash attention specifically.
discussed in prior literature (and in nanoGPT's issue tracker) but
is likely that the mitigation proposed will not work in all cases, such as for larger models. However, we only needed to replicate at the scale we were working at. The experiments in Figure 12 and Figure 13 are included to illustrate how bounding the norm of the query and key head vectors seems to be important for numerical stability. However, this may change in future versions of the flash attention kernel, these results were obtained with Py Torch 2.4.0.
D Additional GNS Results
D.1 Additional GNS Phase Plot
Figure 14 shows the GNS phase plot for the same model as described in Section 4 but with the linear batch size schedule described in Section 5.2.
D.2 Batch Size Schedule
The batch size schedule used in the experiment described in Section 5.2 shown in Figure 15.
D.3 Larger Scale Training
To demonstrate that the method scales to larger models, we trained a 1.3B parameter GPT model on Open WebText using 8 H100 GPUs. The results of this experiment are shown in Figure 16. The left plot shows the per-example gradient norms for all layers, while the right plot shows the per-example gradient norms for only the LayerNorm layers. The GNS computed using the traditional DDP method is also shown for comparison. In Figure 16a we observe that the LayerNorm remains predictive of the total GNS, as in the I1MM model results of Figure 7. When all non-fused simultancous per-example gradient norms were collected we observed an MFU of 40% and when only the fused LayerNorm layers were collected observed an MFU of 57%.
After completing this experiment bug was discovered in the code that decreased per-example gradient norms by constant factor. This caused an underestimation of the GNS. In Figure 16b this can seen when we compare the GNS estimated via DDP method. Initially, we assumed that this constant factor was due a failure of the LayerNorm GNS approximation to larger models. Unfortunately, we did not have the budget in time or resources to rerun the experiment so we corrected the results by multiplying by the constant factor observed in the comparison to the DDP method.
This may be representative of real world scenarios where a large model is pretrained over many DDP nodes. As the user has access to two methods to estimate the GNS, they may account for any bias or slope between the estimates. Then, if it is necessary to continue training on single node, they can use the per-example gradient norms to estimate the GNS. Similar techniques can involve enabling per-example GNS estimation for all layers for short time, or estimating the GNS offline as described in Appendix A.
"ggain following the GPT2-like|43] prescription from Dey et al. [19].
20

---

MLP and Attention
800
10a
600
ONN 10000
2
Tr(E)
10a
GNS
400
5
6
10
7
200
8
100
9
Combined
10
10
2
10
102
5000 1000115002000025000000000
||G|12
Iteration
LayerNorm
GNN -GN5-224900
10000
10a
350
300
1
10
2
250
3
Tr(2)
200
5
10 3
6
150
7
8
100
9
10
50
In_f
Combined
10-5
10
10-3
10
5000 100001500020000250000000000
||G|13
Iteration
Figure 14: GNS phase plot in Figure but focusing on the batch size schedule described in Sec- tion 5.2. Linear/Embedding layers are separated from LayerNorm layers by row and the component estimators of Equations and are plotted (left), with the GNS over the course of training (right).
E FLOP & I/O Formulae
We use the following formulae our FLOP and I/O cost estimations, where Batch Size, T = Sequence Length, K = Input Dimension, L = Output Dimension:
Table 1: FLOPs
Algorithm
Weight Gradient
Gradient Norms
Simultancous
BKL (2T
1) (L 1)
BKL + B (KL 1) BT (2K 2T. + BT
[36]
KL (2BT 1)
Table 2: I/O Weight Gradient BKL + BKT + BLT
Algorithm Simultancous [36]
Gradient Norms
BKL + B 2BT B
BKT + BLT + KL
21

---

GNS and Batch Size over Tokens Processed (alpha=0.9g)
ONS
Batch
100
80
Batch
60
0.0
1.0
1.5
2.0
Tokens Processed
169
Figure 15: The batch size schedule used and GNS observed the 11 batch size schedule experiment illustrated in Figure 15. An aliasing issue is noticeable in the interpolated linear batch size schedule that was used. This has since been fixed in the published code.
GNS vs Train Lass
1000
14
exampee (calibrated
500
2.0
10a
1.0
Train Loss
0.8
GNS (alpha=0.99)
1..
1000
W/
ex-exampl (calibrated
GNS
500
1.0
0.0
0.2
0.4
0.5
0.8
1.0
3.2
1.4
GNS Ratio (DDP) Pereeammlle)
Tokens Processed
1w10
Ratio of GNS Traces (alpha=0.99)
0.5
0.5
Mwwww
Altention LayerNoom
Emmesdiing
0.0
0.2
0.4
0.6
0.8
1.0
3.2
1.4
0.98
1.00
0.96
1.00
Takens Pracessed
1010
alpha
alpha
(b) Comparison of GNS computed using traditional DDP methods and per-example gradient norms.
(a) Regression analysis repeating Figure 6.
Figure 16: 1.3B GPT model train on Web using 8 H100s, trained twice. (Left) Per-example gradient norms for all layers gathered to replicate the analysis in Figure 7. (Right) Per-example gradient norms were gathered for only LayerNorm layers, then compared to the GNS computed using traditional DDP methods.
22

---

Solving the I/O equations above reproduces [36]'s analysis with T 222KLL at the cross-over point above which simultaneous calculation is more I/O efficient. Solving for FLOPs gives:
2KL -1
T V 2K 2L
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The claimed contributions are a method for computing per-example gradient norms in the forward pass of neural network and GNS discovery of correlation between layers. The method is efficient, see Section 5.1, and the GNS is shown to predictable, see Section 4.2.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations are discussed in Section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
24

---

Answer: [Yes]
The main theoretical results concern the FLOP costs of computing per-example gradient norms, which be found in Section 3.1. All other theoretical results are summa- rized from prior work.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results, All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: The algorithms described in Sections are sufficient to reproduce the main experimental results. To obtain the performance reported Section 5.1 the cus- tom LayerNorm kernel is provided with the shared code. The code is available at
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
25

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
The code to implement the cuda kernel described in Section 5.1 is provided with the submission. The code to implement the per-example gradient norm computations described in Section 3 and Appendix B is provided in pseudocode in the text and results may be replicated by using the shared code at https:// //gi thub. Comm/errebasRRsearchh nanoGNS/ treee/reeeeeeeeeee trree/main/ maainexxcce
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
public/gun selliicggudes//// httg///dd//codsliii /ooeeeumiiisionnn olyicy for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wwtppccppcccc/ccccc.ccccccccc..t httic//ggiideeooeooubbiiiiiiiiii CodeSubmisse ierniissoonooo (iionPoiccy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
The experiments conducted in Section 4 use the same settings as Dey al. [19], whose code and hyperparametens are available. The only changes from these details, which were made for numerical stability on GPU, are described in Appendix C.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The experiments in Section 4.2 and 5.2 are the only experiments requiring statistical testing. Section 4.2 reports the Pearson correlation coefficient and Section 5.2 reports the results over three seeds.
26

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi-
dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative етгог rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: The runtime of each experiment performed in Section 4 is provided in the text. The code to implement the custom LayerNorm kernel described in Section 5.1 is available at atttps:/ggttee
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: / /neurisu...pp// cc/ cc/public//eii htwpe/titiiinieeeinnnneeeaee
Answer: [Yes]
Justification: This research has no human participants, uses public open datasets and has generic applications in training deep neural networks.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
27

---

Question:
Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The consequences this work are difficult to anticipate, we frame here as useful for practitioners training deep models but may also influence differential privacy as described in Section 2.3.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploymeent of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper does not release any data or models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided relgeneric easing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Section 4 links to the nanoGPT repo, which our experiments are based on. The code includes sections implementing GNS for DDP that may be found in the code of
Crowson al. [14] which are properly attributed.
28

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code shared with the submission at https: //gps://gitti thub. com/ Crrebraskesearcc/naneeee CAOGGNN//teee/eeeeee ttree/maineese includes instructions for replicating experiments and notebooks for replicating figures.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cooe/mmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
29

---

Answer:
[NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
30


Reviews:
[{"review_id": "V6txRRO3WF", "reviewer": "Reviewer_5nZ5", "summary": "This work proposes a method to compute per example gradient norms as a means to compute GNS. It shows that not all layers are necessary to estimate the GNS and that the per-example gradient norms can be computed for normalized layers without any overhead.", "questions": "Perhaps the work can benefit from experiments with larger model sizes, given that even small open source models have 7B parameters.", "limitations": "As acknowledged by the authors, this paper is limited to transformers, which inherently use normalization sublayer (other architectures do not conventionally use such layers).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This work provides an efficient technique for computing the GNS. The technique is supported by various experiments, and elucidates that the GNS of the model is highly correlated between layer types. Furthermore, for LayerNorm layers, this paper develops a custom kernel to compute the backward pass and the per-example gradient, and find that the the throughput overhead of gathering the per-example gradient is 0 (which outperforms PyTorch’s Layernorm). \n- The paper also replicates prior GNS observations, which helps support the method.\n- I found the experimental results interesting, especially the case study with dynamic batch sizes.", "weaknesses": "- As the authors mention, estimating gradient noise scale is useful in training large scale models. The largest model included in the case study only has 111M parameters. I think that this work could benefit from experiments on larger architectures, given the size of language models used in practice today (e.g., at least 7B). \n- This study is limited to transformers. However, many other architectures do not use normalization sub-layers. This point is also acknowledged by the authors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thoughtful feedback, and for your support of the paper's core idea and experimental approach.\n\n### Regarding larger-scale models\n\n> The largest model included in the case study only has 111M parameters... this work could benefit from experiments on larger architectures\n\nYes, we agree that demonstrating our findings on larger models would be\nbeneficial.  Following your suggestion, we have conducted experiments on a 1.3B\nparameter model. The results are in our attached pdf document in the response\nto all reviewers.\n\n### Regarding architectures beyond transformers\n\n> This study is limited to transformers. However, many other architectures do not use normalization sub-layers. This point is also acknowledged by the authors.\n\nWe agree that this is a limitation.  However, we should have also mentioned in\nthe paper that our per-example gradient estimation methods for other layers\nstill apply (Algorithms 1 and 3). Although perhaps not as performant as\ngathering statistics from the LayerNorms -- depending on the kernel used -- the\nincrease in runtime may be acceptable. In experiments we have run it was at\nworst 30% slower to gather per-example gradient norms for all layers. Gathering\nonly LayerNorm per-example gradients did not slow down training at all. We will\nrevise the paper to make this observation. Thanks!", "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": "I8lqoglpZL", "reviewer": "Reviewer_6ptT", "summary": "This paper proposes a more efficient method for computing per-example gradient norms without significant computational overhead in terms of FLOPs and I/O cost. This method accurately estimates the gradient noise scale (GNS), useful for neural network batch size selection and scheduling. Additionally, it observes that the normalization layer effectively predicts GNS in transformer models. These findings, along with the proposed algorithm, are applied in a case study on a large language model, where effective batch size scheduling reduces training time.", "questions": "- Could the author add more experiments on different models and datasets to validate the generalizability of the method? To ensure robustness, it is also important to include ablation studies. Additionally, the results should be compared with other established batch-size scheduling methods to highlight the advantages or disadvantages of the proposed approach.\n- In section 2.3, it is claimed that the variance of gradient norms determines whether SGD outperforms adaptive optimizers. However, subsequent studies [Kunstner et al, Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be, ICLR 2023] have shown that the performance gap between Adam and SGD persists even in full-batch scenarios without stochasticity. This suggests that there might not be applications of GNS in this context. It would be beneficial for the author to address this discrepancy.", "limitations": "Limitations are well-addressed", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The proposed method efficiently computes per-example gradient norms and rapidly estimates the transformer gradient noise scale.\n- This work also investigates a downstream application of batch size scheduling and demonstrates its time-saving benefits during training.", "weaknesses": "- If the proposed method is only supported by heuristic and empirical observations, and if the primary practical use of GNS is to estimate the critical batch size, then the experiment should be much more comprehensive. Only one set of experiments was conducted on a fixed model/dataset, making it unclear if the results are robust without ablation studies. Additionally, the results were not compared to other batch-size scheduling methods.\n- The finding that the normalization layer effectively predicts total GNS behaviour in transformers has not been thoroughly tested on other transformers, and it's unclear if this generalizes beyond transformers. The authors also do not provide any intuition as to why this might be the case. \n- The paper needs to be more clearly written and easier to read. Some concepts could be better explained or motivated. For example, it's not clear how per-example gradient norms help estimate GNS, what other ingredients are needed, and what other sources of noise there are that should be taken into consideration. It's also challenging to extract key information from the figure captions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments and suggestions, and for positively\nnoting the efficiency of our GNS estimator.\n\n### Regarding the scope of the evaluations\n\n> the experiment should be much more comprehensive … only one set of experiments was conducted on a fixed model/dataset, making it unclear if the results are robust without ablation studies.\n\nWe agree that sufficient experimental evidence is an important consideration. Of course, in this paper we do not establish the relationship between GNS and critical batch size, as McCandlish et al [32] provided exhaustive experiments on this topic on\nmany data and model types.  The focus of our work is validating our novel unbiased estimator of the GNS statistics. To this end, following your suggestion, we will include additional studies for other model types, sizes, and datasets in the revised paper. We propose to expand on Figures 4-6 as follows:\n\n- Figure 4 illustrates the separation of GNS measurements by layer type and Index; this demonstrates the granularity of the estimators in practice.  We have performed the same experiment on many different model types and we will include those results in the camera version of the paper.  We will also add additional plots for other model sizes to the Appendices.\n- Figure 5 is a replication of prior work that opens questions about an experiment from McCandlish's Appendices. We will perform a replication of the original experiment on SVHN in order to check for methodological differences.\n- Figure 6 relates the measurements of GNS across layer types. The results are presented for one model size (111M parameters). We observe this result for language models specifically because that's one large-scale setting where it may be useful. Additional studies for other model types would be valuable here and we will include other model types and datasets in the revised paper.\n\n>  Additionally, the results were not compared to other batch-size scheduling methods.\n\nNote that Figure 8 is presented purely as an example use case for GNS in large-scale language model training. The actual batch size schedule is not novel, and we do not claim to establish a SOTA batch size scheduler.\n\n### Regarding the generalization of the findings\n\n> The finding that the normalization layer effectively predicts total GNS behaviour in transformers has not been thoroughly tested on other transformers … the authors also do not provide any intuition as to why [the predictability] might be the case … could the author add more experiments on different models and datasets to validate the generalizability?\n\nWe agree that it would be interesting to know if this result generalizes to other model types, such as the original post-LayerNorm Transformer, or even to image models, and we will include such studies in the final paper.\n\nRegarding why LayerNorm gradients predict total GNS behaviour, we are currently\nexploring whether the per-example layerwise gradients are correlated for the\nreason that they ultimately reflect variation in the per-example loss. While a full\ntheoretical understanding of this phenomenon may prove elusive, we agree that\nsome additional experiments to gather greater intuition would be valuable. Thank\nyou for this suggestion.\n\n### Regarding the paper clarity\n\n> Some concepts could be better explained or motivated. For example, it's not clear how per-example gradient norms help estimate GNS... It's also challenging to extract key information from the figure captions.\n\nWe regret that Section 2 failed to explain the link between per-example gradient norms and GNS measurements. We suspect the issue may be line 80, where we need to clarify that $B_{big}$ is the full batch and $B_{small}$ are fractions of that same batch.  We will revise section 2 to note that gradient norms are measured for both batch sizes, but typically this can be done cheaply when training DDP because the small batches gradients are on each node and the large batch gradient can be obtained after synchronisation before the update. Per-example gradient norms are intended to obtain a gradient norm for $B_{small}=1$ without requiring any specific training setup (ie DDP). Figure 1 shows that this would give the most accurate measurement of GNS.\n\nWe also agree that key information is difficult to extract from Figures. As described in our rebuttal to Reviewer 32m2 we will revise the Figure captions to include the important message of each Figure.\n\n### Regarding the role of gradient noise in optimizer performance\n\n> In section 2.3, it is claimed that the variance of gradient norms determines whether SGD outperforms adaptive optimizers. However, subsequent studies … have shown that the performance gap persists even in full-batch scenarios without stochasticity\n\nThank you for referring us to this paper!  Indeed, Kunstner et al. (2023) does strongly suggest that gradient variance is likely not the determining factor for why Adam prevails over SGD in Transformer models.  Our overall point was that our work can support such studies by providing tools to efficiently collect gradient statistics.  However, as you mention, this particular area is perhaps now less well-motivated, and we will revise this part of the related work accordingly.  Thanks again!", "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": "WvvQ3e0iHw", "reviewer": "Reviewer_fwbc", "summary": "The paper proposes a method for efficient computation of per-example gradient norm for the broader usecase of computing gradient noise scale. Further the authors showcase the usecase of gradient noise scale (GNS) in Transformers and showcases that GNS of only normalization layers in transformer models suffices, rather than the total GNS. Using this proposal, the paper proposes a batch size scheduling resulting in training time reduction.", "questions": "Step 4. of Algorithm 1. (what is the exact correction) and is very limited in scope in terms of contribution.", "limitations": "Yes, I appreciate that the authors have adequately addressed the limitations of this work in terms of their empiricial evaluations only on Transformer architectures and not on other similar architectures like RNNs or state-space models, which is possibly left as a future work. Apart from this there are no limitations or any potential negative social impact of their work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "The paper for the most part is well written and has shown useful applications of gradient noise scale in Transformer architectures which can be further extrapolated to other architectures. e.g. state space models.\n\n- Related work section discussion is well motivated to discuss the utility of gradient noise scaling, perhaps a more better way to represent this section would have been to segregate the related work with smaller sections with headers - one discussing practical utility and other definitions used in literature for gradient noise scaling and any recent efforts towards efficient computation for per-gradient norms.", "weaknesses": "Overall, I feel the main contributions of the paper is not well presented in this draft and in many portions the writing is hard to follow and seems disconnected from other sections in the paper.\n\nOn one aspect the authors propose existing methods of efficient per-example gradient norms in Section 2.2 but only a few papers like Li et al. [29] and Goodfellow [22] are mostly discussed. It would be nice to have a much more better discussion of more related work to have a thorough pros and cons discussion.\n\nNotation writing could be improved to a better extent.\nFor instance in Sec 2.2, line 74 which space does B, I, K indicate or is just a placeholder for 3D tensor dimensions.\nSimilarily line 75 $x_{bti}y^{'}_{btk}x_{bui}y^{'}_{buk}$ , how does the index indicate u ?\n\nRelated work discussion on other methods on per-gradient norm efficient computation is missing. For instance https://openreview.net/forum?id=xINTMAvPQA is cited but this one or similar other approaches how they have approached this problem and how the authors proposal differs in that respect , that thorough discussion would be appreciated.\n\n**Minor fixes in Writing**\n\n- Section 2.1: it may be shown [32] -> It may be shown McCandlish et al. [32].  And you may remove McClandish elsewhere.\n- Section 2.1 : I do not see any specific discussion on $B_{big}$ and $B_{small}$. How do they differ? Any references to later section in the paper where they define the size differences.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful review.  Your suggestions will definitely help improve the paper.  It is also gratifying to know that the work is well motivated and the application of per-example gradient norms to Transformers training is useful and clear.\n\n### Regarding revising the related work\n\n>  perhaps a better way to represent this section would have been to segregate the related work with smaller sections with headers\n\nWe thank the reviewer for this suggestion. We will add `\\paragraph` headers to the related work section to make it easier to navigate.\n\n### Regarding clarifying the contributions\n\n> Overall, I feel the main contributions of the paper is not well presented in this draft\n\nWe plan to address this suggestion through a new figure in Section 1 that\nhighlights the unique contributions of our method.  For example, using something\nlike Rebuttal PDF Figure 1 and the figure caption, we will contrast our\n(layerwise) per-example norms with other methods that must aggregate gradients\nat a coarser level based on their specific data-parallel configuration.  We will\nalso do a better job of highlighting the specific contibutions in the\nintroduction, previewing the experimental investiation that unfolds over the\nrest of the paper.\n\nWe will also do a better job of motivating the use of per-example gradient norm\napplications beyond GNS estimation (such as for differential privacy).  We\nshould have also made it more clear that GNS itself is rarely used because it\ntraditionally requires large DDP setups to be useful. With our method, it can be\napplied both on the smallest MNIST experiment or extremely large scale language\nmodel training.\n\n### Regarding per-example norms in related work\n\n> only a few papers like Li et al. [29] and Goodfellow [22] are discussed … it would be nice to have … more related work to have a thorough pros and cons discussion.\n\nThis is a good suggestion, thanks.  While related work in this area is limited,\nfollowing your suggestion, we did discover [\"Efficient Per-Example Gradient\nComputations in Convolutional Neural\nNetworks\"][pe_conv](). This paper should have been\nmentioned in the related work section and we will add it in the final version.\nAs [backpack][] makes clear, this also reduces to a 3D tensor regime, so will be\nequivalent to our method. Li et al. [29] is the only method we know of that\nfocuses directly on per-example norms in 3D tensor regimes, which motivated our\ncomparison in Section 3 (lines 111-148). We agree that a discussion on pros and\ncons would be valuable, for example Goodfellow [22] is optimal in 2D tensor\nregimes, and we will add this in the final work.  Thanks again!\n\n[backpack]: https://github.com/f-dangel/backpack\n\n> discussion on other methods on per-gradient norm efficient computation is missing. For instance https://openreview.net/forum?id=xINTMAvPQA is cited but this one or similar other approaches how they have approached this problem and how the authors proposal differs in that respect\n\nWe agree that a thorough discussion of this comparison would be valuable. At\npresent the paper only provides a short comparison in Appendix A (line 389).\nAppendix A currently only compares the methods in the context of GNS estimation.\nTo improve it, we will list pros and cons of all known per-example gradient norm\nestimation methods [22, 29, [rochette2019efficient][pe_conv],\n[gray2023efficient][]] directly in the final version and discuss their usage in\nGNS estimation.\n\n[pe_conv]: https://arxiv.org/abs/2204.02311\n[gray2023efficient]: https://openreview.net/forum?id=xINTMAvPQA\n\n### Regarding notation and citations\n\n> Notation writing could be improved … for instance in Sec 2.2, line 74 which space does B, I, K indicate or is just a placeholder for 3D tensor dimensions. Similarily line 75 $x_{bti}{y^{'}}_{btk}x_{bui}{y^{'}}_{buk}$ , how does the index indicate u ?\n\nWe agree that the interpretation of $B$, $I$, and $K$ should be included. We\nwill add notes that they correspond to the batch size, input dimension, and\noutput dimension, respectively. Additionally, the $u$ index shares the same\nindex space $u \\in (1, ..., T)$ as $t$, which may cause confusion. We will\nclarify this in the final version.  Good catch!\n\n> Section 2.1: it may be shown [32] -> It may be shown McCandlish et al. [32]. And you may remove McClandish elsewhere.\n\nWe thanks the reviewer for pointing this out, we mechanically applied `\\citet`\nwithout considering how many times \"McCandlish et al.\" would get printed. We\nwill address this in the final version.\n\n> Section 2.1 : I do not see any specific discussion on $B_{big}$ and $B_{small}$. How do they differ? Any references to later section in the paper where they define the size differences.\n\nWe thank the reviewer for pointing out this gap. On line 50 $B_{big}$ and\n$B_{small}$ have an incomplete definition. In the final version we will explain\nthat $B_{big}$ is typically the full batch size and $B_{small}$ is a fraction of\nthat, typically the batch size executed on DDP nodes or during gradient\naccumulation.\n\n### Regarding Algorithm 1\n\n> Step 4. of Algorithm 1. (what is the exact correction) and is very limited in scope in terms of contribution.\n\nWe will expand the discussion in the paper to make the contribution of line\n4 clear. We agree that the correction on line 4 of Algorithm 1 is only\na consequence of backpropagation of losses that have been reduced by a mean; it\nis not strictly part of the per-example gradient norm estimation. Lines 123-128\ndescribe why this correction exists and we found that it was often a source of\nerror in new implementations.\n\nThe form of the correction in Algorithm 1 is because it is safer to apply the\ncorrection to the norm after taking the mean of the squared norm $s_w$ (it would\nbe equivalent to sum the loss and then scale the update gradients, but this is\nmore resource intensive and could introduce loss scale issues).", "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": "rAitpeR0aV", "reviewer": "Reviewer_32m2", "summary": "This paper explores efficient ways of computing the gradient noise scale (GNS) when training transformers. The main practical relevance of the GNS is that it can be used to estimate the critical batch size, where the larger batch sizes become computationally inefficient. The authors discuss different ways of doing this at a fine granularity both time and layerwise. They show that the backwards pass through a layernorm can be modified to estimate their GNS for essentially free and that this predicts the GNS of other layers. The authors then vary the batch size throughout training based on this, showing this can potentially save compute for a given target model performance.", "questions": "Overall I think the core ideas are interesting but this paper would strongly benefit from being resubmitted with improved writing and overall presentation.\n\nSuggestions for improvement:\n- Drop the Einstein notation, make things explicit instead of implicit.\n- Consider rewriting the introduction to gently introduce the GNS, maybe with a diagram and summarize the high level importance of this before diving in. What are the high level ideas needed to understand your contribution? Do you really need to go into all these details upfront?\n- Improve the description of your figures. Take Figure 1 for example, it is very hard to parse. The caption does not summarize the message, it should also emphasize how the dashed and plotted lines are related.\n- For Figures 2 and 3, maybe emphasize more what Li et al is, what the differences are and what you expect to see? Why not normalize the IO costs as well in the comparison (maybe you need some assumptions on the other kernels. Without some sort of reference number it is hard to see how much these IO overheads matter in practice.\n- Figure 5: Again, maybe tell us what is expected? What is the takeaway?\n- Figure 6: It is very unclear what is going on in the two plots on the right. Almost no description is given in the text or the caption of the “EMA alpha values”.\n- Batch Size Schedule: It is not clear how “a batch size schedule that increases linearly with the number of tokens processed to the original batch size” relates to your GNS estimations. Do you set the length of the schedule based on the GNS? Is this done based on prior estimates of the GNS or in an online fashion?\n- Figure 8: This should probably be done using learning rate schedules of different lengths for statements like 2x to be meaningful.", "limitations": "Some limitations discussed, negative societal impact not applicable", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "- The core idea of estimating the GNS in a cheap way based on normalization layers is interesting and relevant to the community.\n- The proposed layernorm based estimation can be performed very cheaply.", "weaknesses": "- Some more work is required to make the proposed method practical due to numerical issues. Although the idea is interesting it would be much more impactful if the kernel worked as a drop-in replacement.\n- The 18% speed improvement claimed in the abstract may be an overclaim due to the numerical issues and presumably only being applicable to certain types of training setups (maybe one GPU doing gradient accumulation rather than distributed setups).\n- The paper is quite hard to follow, despite the relatively simple ideas presented.\n- The use of Einstein notation contributes to this, it is fine for brevity but not clarity.\n- The paper relies too heavily on McCandlish et al 2018. Despite spending significant effort on trying to summarize the relevant portions of this paper, many things are still unclear in the later sections.\n- The figures are hard to interpret, both due to very short captions that do not summarize the high level idea and takeaway, and insufficient labels on the figures themselves.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your detailed and constructive review, and your kind words regarding the appeal and relevance of the paper’s core idea.\n\n### On the clarity of the paper\n\n> The paper relies too heavily on McCandlish et al 2018. Despite spending significant effort on trying to summarize the relevant portions of this paper, many things are still unclear in the later sections.\n\nUnfortunately, despite the known use of GNS as detailed in the Related Work\nsection, there is limited research published on the topic since McCandlish et al.\n(2018). However, GNS estimation is of practical value in our work and we believe\nit will be for the community.\n\n### On Einstein notation\n\n> Drop the Einstein notation, make things explicit instead of implicit.\n\nWe agree that Einstein notation may be unclear. The most straightforward way to\naddress this problem would be to write out the sums explicitly. However, we find\nthe utility of the Einstein notation to be the ability to express that the order\nof the summation is flexible. Another better solution would be to write example\nmatrix operations. For the equation on line 75, we could have written\n$$\nW'_b = X_b {Y'}_b^T, \\quad n_b = ||W'_b||_F^2\n$$\nas one possible contraction. We will add these examples to the paper to make the\nnotation more clear.\n\n> Consider rewriting the introduction to gently introduce the GNS, maybe with a diagram and summarize the high level importance of this before diving in. What are the high level ideas needed to understand your contribution? Do you really need to go into all these details upfront?\n\nThanks for this suggestion. Based on this we have created a diagram to shown in\nour attached pdf in Figure 1. We think this may help prime the reader for the\nconcepts in the paper.\n\n### On Figure messaging\n\n> The figures are hard to interpret, both due to very short captions that do not summarize the high level idea and takeaway, and insufficient labels on the figures themselves.\n\n> Improve the description of your figures. Take Figure 1 for example, it is very hard to parse. The caption does not summarize the message, it should also emphasize how the dashed and plotted lines are related.\n\nThank you for this feedback. We agree that the figure captions should state the\nmessage of the figures. For example, in Figure 1, we will say, \"For the same\nnumber of samples processed, a smaller $B_{small}$ always has a lower standard\nerror (dashed), while the size of the large batch $B_{large}$ does not affect\nthe standard error.\" The solid lines in this figure are the estimate of the GNS\nat each number of samples processed, supposed to illustrate the uncertainty\ncaused by the finite number of samples. To simplify the figure it will be\nclearer to remove the solid lines.\n\n> For Figures 2 and 3, maybe emphasize more what Li et al is, what the differences are and what you expect to see? Why not normalize the IO costs as well in the comparison?\n\nWe agree that it would be beneficial to provide more context for Figures 2 and 3.\nWe will add the following messages to the captions:\n\n- Figure 2: \"...The FLOP cost of Simultaneous per-example gradient norms is strictly dominant to alternative methods (left) and the ratio of this additional cost to the FLOP cost of processing the entire model does not depend on context length (right).\n- Figure 3: \"...The I/O cost of Simultaneous per-example norms is greater for models of 10B parameters or more but approximately equivalent for models of 1B parameters or less, depending on context length.\"\n\nIn addition we have tested normalized I/O costs in the Figure and found that it\nmakes that Figure more clear. We will use this Figure in the updated paper.\n\n> Figure 5: Again, maybe tell us what is expected? What is the takeaway?\n\nTo Figure 5 we will add the message, \"This Figure replicates an experiment from\nMcCandlish et al. (2018) showing how the ratio of $\\epsilon / B$ causes changes\nin the measured GNS but only due to changes in the learning rate. The batch size\ndoes not have the predicted effect.\"\n\n> Figure 6: It is very unclear what is going on in the two plots on the right. Almost no description is given in the text or the caption of the “EMA alpha values”.\n\nWe thank the reviewer for noticing that the description of the EMA on lines\n177-178 does not mention how the alpha value controls smoothing. We will add an\nAppendix making the role of the alpha value and add EMA pseudo code, because\nthere are many ways to implement it.\n\nFor Figure 6 we will add the message, \"The total GNS (black) on the left is\npredicted well by individual layer types as indicated by the correlation\ncoefficients (right), however the type with slope closest to 1 is LayerNorm\n(center),\nonly overestimating the GNS by less than 40% across EMA alpha values.\"\n\n### On batch size scheduling\n\n> Batch Size Schedule: ... Do you set the length of the schedule based on the GNS? Is this done based on prior estimates of the GNS or in an online fashion?\n\nThis is a great question! The batch size schedule we used linearly increased\nduring training because we had observed this trend in the GNS and this is\nillustrated in Figure 3 of our attached pdf. We agree that an automatic online\nbatch size schedule using GNS would be ideal but we didn't want to\novercomplicate the paper. The schedule is supposed to mimic what a practitioner\nwatching the GNS might do; in practice an operator adjusts the batch size\nmanually for expensive runs.\n\n> Figure 8: This should probably be done using learning rate schedules of different lengths for statements like 2x to be meaningful.\n\nWe agree that varying hyperparameters would make this result more robust. At\nthis scale, we observed that the batch size scheduled run was outperforming\na tuned baseline, so we did not pursue this further", "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:
Thank you to all the reviewers for their thoughtful feedback. After reading the
reviews we noted the following points we could address with additional figures:

1. Reviewers 6ptT, fwbc and 32m2 brought attention to the clarity of the work.
2. Reviewers 5nZ5 and 6ptT asked for additional experiments on other models or
   datasets.
3. Reviewer 32m2 asked about the numerical stability, restrictions to DDP and
   requirements for gradient accumulation.

To address the first point we provide Figure 1 in our attached document. This
figure presents a diagram to explain some of the key concepts, ie that we are
dealing with gradient norms that vary across minibatches and between layers.
Additionally, it presents an intuition for why the gradient norms between layer
types may be similar.

To address the second point we provide Figures 2a and 2b in our attached document.
This figure illustrates the results of a 1.3B GPT model trained on
OpenWebText twice from scratch on 8 H100 GPUs. Both runs were configured to
match the 1.3B Chinchilla optimal Cerebras-GPT training run. The GNS was also
gathered at the same time in both runs using the traditional DDP method (small batch gradient
norms gathered on each node). The runs were numerically stable after addressing
issues independent to our method (details in response to Reviewer 32m2).

The first run gathered per-example gradient norms for all layers and ran at 40%
[MFU][] (p.9). Figure 2a repeats the analysis of Figure 6 in the paper, finding
again that the total GNS is well predicted by only the LayerNorm per-example
gradient norms. However, it was found that the slope of the regression is higher
at approximately 2.2. This indicates that it may be necessary to calibrate the
estimate of the GNS intermittently by enabling all per-example gradient norms
for a short time during training, or matching it to the GNS gathered by DDP, if
available.

The second run enabled only the LayerNorm per-example gradient norms and ran at
57% [MFU][]. Using the 2.2x calibration factor the GNS is plotted against the GNS
gathered at the same time by DDP in Figure 2b. The GNS estimates are very close
throughout training. This would allow, for example, continued tracking of the
GNS if we were to move the run to one GPU during training, at which point
tracking the GNS via DDP would not be possible.

[mfu]: https://arxiv.org/abs/2204.02311

### On numerical stability

Thanks to reviewer 32m2 for raising an important point about the numerical
stability of the work that we would like to address to all reviewers to make
sure it is clear that our method does not affect the numerical stability of
training.

> Some more work is required to make the proposed method practical due to numerical issues. Although the idea is interesting it would be much more impactful if the kernel worked as a drop-in replacement.

The numerical issues mentioned in the paper are indeed an issue we spent a lot
of time on. However, they are not due to our method; early in testing we
disabled our code and found the same issues on the main branch of nanoGPT.

Specifically, the following config is sufficient to reproduce the issue:
```
batch_size=8
gradient_accumulation_steps=2
block_size=2048
max_iters=72_000
lr_decay_iters=72_000
warmup_iters=1000
```
It is caused by an interaction with flash attention and AMP in PyTorch.
Disabling flash attention or training float32 resolves it. Some other users have
made issues documenting similar behaviour in the nanoGPT repository, such as
[here](https://github.com/karpathy/nanoGPT/issues/137). We have since resolved
these issues with small architectural changes.

The kernel does work as a drop in replacement and we agree that this is vital.
We hope additional experiments shared in our attached pdf at 1.3B scale on
a distributed setup should make this clear. Complete code will be released with
the camera-ready version of the paper.

> The 18% speed improvement claimed in the abstract may be an overclaim due to the numerical issues and presumably only being applicable to certain types of training setups (maybe one GPU doing gradient accumulation rather than distributed setups).

This is an important concern. As shown in our attached rebuttal pdf, we are able
to train DDP at large scale with this method. Additionally, it does not depend
on gradient accumulation, although that is another way to gather small batch
gradient norms.

Meta-review: This paper provides observations on gradient norms, which lead to removal of some nearly-redundant calculations across layers and thus speedups.  Reviewers had a variety of concerns regarding the rigor and presentation quality, but overall the paper is above bar and I am happy to recommend acceptance.
Decision: Accept (poster)