Scalify: scale propagation for efficient low-precision LLM training

Authors: paul  balanca
sam  hosegood
carlo  luschi
andrew  fitzgibbon
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-05-30 12:08:40

Abstract: Low-precision formats such as float8 have been introduced in machine learning accelerated hardware to improve computational efficiency for large language models training and inference. Nevertheless, adoption by the ML community has been slowed down by the complex, and sometimes brittle, techniques required to match higher precision training accuracy. In this work, we present Scalify, a end-to-end scale propagation paradigm for computational graphs, generalizing and formalizing existing tensor scaling methods. Experiment results show that Scalify supports out-of-the-box float8 matrix multiplication and gradients representation, as well as float16 optimizer state storage. Our JAX
implementation of Scalify is open-sourced at [github.com/graphcore-research/jax-scalify](https://github.com/graphcore-research/jax-scalify).

Main paper: 
 introduction: 
 As the number of parameters in deep learning models have progressively increased through the years, machine learning researchers and engineers have been experimenting with the use of low precision formats for training and inference to improve computational efficiency and reduce memory usage and required bandwidth.Starting with IEEE float16 format (the only 16-bit format supported on GPUs initially), Micikevicius et al. (2017) introduced the idea of loss scaling to improve stability of training and reduce the occurrence of overflow on the backward pass. This concept was later refined with the introduction of NaN and max tracking in the backward graph to dynamically adjust loss scaling depending on the training dynamics. In the landscape of low-precision techniques, loss scaling can be seen as a proto tensor-scaling method, where a global tensor scale is set on the backward gradients for numerical stability, and then compensated when the optimizer update Algorithm 1 JAX training loop using SCALIFY transform (modifications in blue). SCALIFY generalizes previous approaches with end-to-end scale propagation in the neural net computational graph (including optimizer). is applied to the model state. Dean et al. (2012) and Abadi et al. (2016) introduced the bfloat16 format (BF16) as a different path to improve 16bit training stability. Compared to IEEE float16 (FP16), bfloat16 trades off mantissa bits for exponent bits, aligning with float32 exponent width (8 exponent bits compared to the 5 exponent bits of float16). In the rest of this work, we use the common notation E8M7 and E5M10 for respectively designating bfloat16 and float16 formats. Thanks to its higher dynamic range, bfloat16 has the main benefit of being a simple drop-in technique in neural net training, not requiring the complexity of dynamic loss scaling. On the other hand, as presented in Micikevicius et al. (2017) ; Peng et al. (2023) , because of the reduction of mantissa bits, bfloat16 can not be used for master weights, normalization statistics or optimizer state, meaning that for these quantities With the recent emergence of large language models (LLMs) exceeding 100 billion parameters, researchers and ML hardware vendors have explored the use of float8 formats (FP8) in matrix multiplication, further improving energy efficiency, compute time and memory footprint. The literature and practitioners (Wang et al., 2018; Sun et al., 2019; Micikevicius et al., 2022; Kuzmin et al., 2022; Noune et al., 2022) have converged towards the definition and use of two FP8 formats: E4M3 and E5M2 (i.e. respectively using 4 bits and 5 bits width exponent). The format E5M2 is used to represent the dynamic range of gradients in the backward pass, whereas E4M3 is used for activations and weights in the forward pass, that necessitate additional precision. With the introduction of these lower precision formats came the requirement of ad-hoc tensor scaling with more granularity than global loss scaling. Several papers (Noune et al., 2022; Micikevicius et al., 2022; NVIDIA, 2022) have shown that per-tensor scaling of FP8 matrix multiplication (matmul) inputs using the input signals statistics is a reliable technique that maintains the accuracy of training in higher precision. In practice, the implementation of efficient FP8 tensor scaling during training required the introduction of delayed scaling: tensors are scaled using the input signal statistics of the previous micro-batch, allowing scaling and statistics gathering to be done simultaneously. Finally, Peng et al. (2023) recently showed that float16 per-tensor scaling can also be used to reduce the memory footprint of master weights and optimizer state.The main drawbacks of current approaches for FP8 training are the black-box complexity of a fused FP8 matrix multiplication and statistics gathering kernel as well as the additional compute requirement to estimate input tensor statistics at every matrix multiplication (in the forward and backward passes). The latter issue has been partially solved by the introduction of delayed scaling, but at the expense of a more complex model state management. In this work, we aim at simplifying low precision training by generalizing and automating tensor-scaling to the entire computational graph (i.e. forward, backward and optimizer for neural net training). We introduce the SCALIFY transform which fully propagates tensor scaling information in the computational graph, bringing scaled FP8 and FP16 techniques under the same automated paradigm. Generalizing the existing approaches proposed in the literature, we believe the proposed formalization of tensor-scaling has multiple advantages:• Systematic scale propagation decouples matrix multiplication and scaling, allowing a more efficient lowprecision training schemes (reducing tensor statistics gathering);• Integration of FP8 formats as "just another datatype": no ad-hoc black-box custom C++ kernel is required (we present in Algorithm 6 a general linear layer implementation);• Seamless and efficient integration of scaled FP16 weights (Peng et al., 2023) and optimizer state, achieving similar training robustness and accuracy to FP32 and BF16;• Model invariance: the SCALIFY transform does not change the semantics of the computational graph (i.e. in full precision, results are identical);Table 1 presents a summary of the most common lowprecision training methods and how they compare to SCAL-IFY in accuracy and ease of use.We provide an open-source implementation of SCALIFY in JAX (Bradbury et al., 2018) . In the rest of this work, we follow the JAX/XLA naming convention to describe in more details our JAX open-sourced implementation github.com/graphcore-research/jax-scalify. However, the same principles can be applied to any other ML framework (e.g. PyTorch at the ATEN operators and torch.compile level). 
scaled array representation and quantization: 
 This work makes use of a "scaled array" representation, whereby a tensor X is represented by a pair of arrays (X d , X s ) satisfying X = X d • X s , where X d has the same shape and datatype as X and X s is a tensor broadcastable to X d (in many cases, X s is a scalar). In practice, for usability by the ML practitioner, and for automated graph transformation, we represent scaled arrays using an explicit data structure ScaledArray as outlined in Algorithm 2. Note that there is no assumption made in Algorithm 2 on the datatype used in the array. In fact, the purpose of SCALIFY is to support out-of-the-box mixed computational graphs with scaled FP8 and FP16 tensors, generalizing the work of Peng et al. (2023) .In principle, there is also no restriction on the format used for the scale component in X s . Nevertheless, we choose to follow the common practice of using power-of-two 8-bit scaling for the following reasons:• Using general floating point scaling can induce additional loss of accuracy in the several rescaling operations added to the graph by the SCALIFY transform.• Power-of-two rescaling can be implemented very efficiently in accelerated hardware, as a simple add operation on the exponent of floating point numbers.• 8-bit power of two scaling (i.e. E8M0 format) is the standard adopted by the Open Compute Plaftorm (OCP, 2023) for the recent micro-scaling block formats (MXFP4, MXFP6 and MXFP8).As the format E8M0 is not yet directly supported by modern ML accelerated hardware, in our experiments we have chosen to simulate the power-of-two scale using float32 (the latter having 8 exponents bits, which is sufficient to represent any E8M0 value). The additional memory footprint is negligible: one float32 value per tensor.Assuming X is represented in full precision, finding the optimal ScaledArray low-precision representation (q(X d ), X s ) of X is a quantization error problem (where q is the quantization method corresponding to the datatype selected). Using the signal-to-noise ratio (SNR) as the underlying metric, one wants to solve:X s = argmax σ∈E8 s.t. X=σX d E[X 2 ] E[(X -σq(X d )) 2 ] = argmax σ∈E8 s.t. X=σX d E[X 2 d ] E[(X d -q(X d )) 2 ],where E 8 is the set of finite (valid) values in E8M0. In other words, the quantization problem is reduced to optimal quantization of the term X d , under the scale constraint X s ∈ E 8 (i.e. X s being a power-of-two). In practice, we approximate the optimal solution of this problem with a two stage algorithm:1. Quantizing X d under the unit scaling constraint E[X 2 d ] ≃ 1, generating a floating scale tensor X s ;2. Rounding X s to a power of two, and correcting accordingly the quantized tensor q(X d ).For the first step of estimating the optimal q(X d ), we follow the path of the unit-scaling work (Blake et al., 2023) by using E[X 2 d ] ≃ 1 as the sweet spot for maximizing the SNR of X d . In short, it falls into the optimal SNR interval for FP8 and FP16 floating point representations (see Fig. 2 in (Blake et al., 2023) ), while still allowing accurate representation of large outliers (common in LLMs training and inference as shown by Dettmers et al. (2022a) ). In other words, it can properly represent a tensor distribution that is the combination of a main mode (Gaussian or heavy-tailed) and large outliers. Note that Blake et al. (2023) considers the variance of tensors, but the same argument holds on the mean squared error (MSE) using the bias-variance decomposition.Assuming ScaledArray inputs satisfying E[X 2 d ] ≃ 1, the SCALIFY graph transform aims at propagating this property in the computational graph, using a combination of unit-scaling rules for every operation to re-balance the X d and X s terms, and dynamic rescaling using X d tensor statistics.In the case of neural-network (NN) training (see Algorithm 1), estimating the scale in the model state can be done as following:• At initialization, the scale is usually known by the user (e.g. random Gaussian initialization with known variance or zeroed tensors);• At every state update, SCALIFY will propagate a new scale into the updated state. Additionally, the scale can be dynamically re-estimated every N step by gathering statistics on the current state (see dynamic rescale in Algorithm 1).In practice, as presented in Blake et al. (2023) , the probability distribution of weights evolves only slowly during training. As a consequence, with a proper scale initialization of the state, our experiments presented in Section 4 did not show dynamic rescaling of the model state is required. 
scalify transform: 
 The SCALIFY graph tracer propagates ScaledArray inputs (as defined in Algorithm 2) in the computational graph.Algorithm 2 ScaledArray data structure propagated through the graph by the SCALIFY transform @dataclass class ScaledArray: """Scaled array generic representation. data: Main data tensor. scale: Scale tensor Usually scalar and power-of-two. """ data: Array scale: Array def __array__(self) -> Array:"""Reconstruct represented array.""" dtype = self.data.dtype return self.data * self.scale.astype(dtype)In the case no scaled arrays are passed (e.g. omitting the calls to as scaled array in the training loop Algorithm 1), the computational graph after SCALIFY will remain unchanged. On the other hand, passing ScaledArray model state and input data, the example update function will then return an updated ScaledArray model state as well as a ScaledArray loss (which can be trivially converted to a scalar) by propagating scaling in the forward pass, backward pass, and optimizer update.Note that even in the case where all inputs were ScaledArray instances, the computational graph may still contain unscaled tensors (typically broadcasted constants), meaning that SCALIFY must be able to handle properly mixed unscaled-scaled tensors combination. Additionally, as presented in Section 3.3, supporting unscaled arrays is necessary to dynamic and delayed rescaling features.As shown in Algorithm 2, ScaledArray objects are simply a pair of a main data array X d and a scale array X s , broadcastable to the former. In this work, we focus on the case where the scale X s a scalar, but more general shapes can cover channel scaling (e.g. van Baalen et al. ( 2023)) and block scaling formats such as MX (OCP, 2023).As presented above, the SCALIFY tracer should support mixed unscaled/scaled computational graphs. As a consequence, there is a need to explicitly control and represent in the graph the switch between unscaled and scaled representations of an array. For that purpose, we introduce the following two bundling/unbundling primitives to JAX:• set scaling: Always transform the input into a ScaledArray using the scale provided. If the input is already a ScaledArray, the data component is rescaled using the ratio of existing and new scales. • get data and scale: Unbundling an input tensor into its data and scale components. If the input is not a ScaledArray instance, the method directly returns the input paired with a constant scale 1.Defining set scaling and get data and scale with proper no-op semantics when the SCALIFY transform is not used allows full backward compatibility in the model definition: any additional scaling logic added to improve numerical stability (e.g., any dynamic rescaling strategies presented in Section 3.3) does not change the semantic of the computational graph (i.e. the neural net model definition or optimizer update).To propagate scale arrays in the computational graph, SCAL-IFY requires to implement for every basic JAX primitive (or PyTorch ATEN operation) a mathematically equivalent operation on ScaledArray. More formally, if f is a basic primitive with(Y 1 , . . . , Y m ) = f (X 1 , . . . , X n )we need to define a mathematically equivalent method f scaled : A trivial (and universal) way to satisfy the previous property would be to simply un-scale the inputs X 1 , . . . , X n , and compute Y 1 , . . . , Y k by directly calling f , and then rescale the outputs using tensor statistics. Obviously, doing this would erase the main benefits of doing scale propagation through the graph, as the unscaling step would potentially reintroduce underflow or overflow in X i (or require higher floating point precision). Additionally, we aim at keeping SCALIFY computationally light by default, meaning estimating X d statistics is opt-out by default (see Section 3.3 for explicit dynamic rescaling).(Y 1,d , Y 1,s , . . . , Y m,d , Y m,s ) = f scaled (X 1,d , X 1,s , . . . , X n,d , X n,s ) such that if the inputs satisfy unit-scaling E[X 2 i,d ] ≃ 1, then the outputs also verify E[Y 2 j,d ] ≃ 1.As a consequence, we adopt the same unit-scaling strategy as Blake et al. (2023) : define a set of static rules for every primitive operation, relying only on the inputs scale and shape to approximate output scaling. To estimate the former, we employ the same simplifying assumption as Blake et al. (2023) : for every operation f , X 1 , . . . , X n are independent Gaussian tensors. Under this hypothesis, one can estimate the scaling of the outputs Y 1 , . . . , Y k for most common NN operators. We provide in Algorithm 3 the (simplified) implementation of scale propagation for the primitives add and matmul.As presented in Section 2 and in Algorithm 3, the scale value is always rounded to a power-of-two to avoid additional floating point errors. In our experiments, rounding down of the scale is used instead of round-to-nearest (or stochastic rounding), as experimental data tends to show it improves preservation of the unit-scaling property E[X 2 d ] ≃ 1 in the graph.Similarly to automatic differentiation available in ML frameworks, automatic scale propagation using basic primitive decomposition of the graph can be sub-optimal in terms of numerical accuracy and computational efficiency. As a consequence, we also provide a mechanism, custom scale (similar to JAX custom jvp and custom vjp), to define custom scale propagation rules. The custom scale feature happens to be typically useful in two common types of neural-net layers: activation and normalization layers.In the first case, scale propagation through a complex nonlinear function like gelu can be brittle, whereas from the high level semantic point of view of activation functions, it is clear these functions aims at propagating positive values unchanged while clipping negative ones to zero (with a more or less smooth transition between the two regimes). While the L2 norm of the clipped tensor is smaller than the original one (meaning we could potentially assign a smaller scale to the output), we take the position of propagating the same scale to the output in order to avoid overflowing of large outliers. Appendix C.1 presents in more details the implementation of scale propagation in activation functions.Similarly, it is inefficient and unnecessary to perform full scale propagation through a normalization layer (when excluding the affine correction part). The purpose of normalization layer is to provide an output with (approximately) mean 0 and variance 1, and as a consequence, using the scale invariance of the mean-variance normalization, it is more efficient and numerically stable to directly normal-Algorithm 4 Implementation of dynamic rescaling using set scaling and get data and scale primitives. Finally, the custom scale feature allows properly tuning scale propagation and precision in attention layers, whether it is optimized standard softmax attention (Dao et al., 2022) or more recent alternative approaches such as Mamba (Gu & Dao, 2023) . We believe that SCALIFY formalized scale propagation will incentivize ML researchers and engineers publishing new type of NN layers to investigate low precision accuracy and release optimal scale propagation rules.A large literature on low-precision FP16 and FP8 training relies on tensor statistics gathering and dynamic rescaling to keep the training dynamic stable (see FP16 automaticloss scaling (Micikevicius et al., 2017) and FP8 training (Micikevicius et al., 2022; Noune et al., 2022; Peng et al., 2023) ). One main goal of SCALIFY scale propagation is to reduce the amount of dynamic rescaling necessary to integrate into a NN model in order to match higher precision accuracy, leading to more efficient FP8 and FP16 training schemes.One major benefit of scale propagation is to be able to decouple dynamic rescaling of tensors from operations using FP8 low precision in the computational graph (typically matrix multiplication or reduce operations). More specifically, standard FP8 training requires dynamic rescaling of inputs at every matrix multiplication (see TransformerEngine implementation NVIDIA (2022)). SCALIFY scale propagation allows decoupling and finer control of when dynamic rescaling is applied, meaning the overhead of tensor statistics gathering can be reduced. For instance, in the example of Algorithm 1, we decide to update scaling at the end of the entire training iteration, following the optimizer update (with the straight-forward optimization of only doing this rescaling every N micro-batches).As presented in Algorithm 4, the implementation of dy-namic rescaling (or delayed rescaling) can be done directly in pure Python using the two primitives set scaling and get data and scale introduced in Section 2. This highlights another benefit of SCALIFY scale propagation: the rescaling logic appears explicitly in the computation graph instead of being embedded in a low level C++ kernel (NVIDIA, 2022).More specifically, the implementation of Algorithm 4 shows that it is useful to first introduce a generic rebalance scale method operating on ScaledArray: instead of setting the scale, the semantic of this method is to perform a relative correction on the scale.As for set scaling and get data and scale, we note that this method is a simple no-op when applied to unscaled array, keeping backward compatibility when using rescaling inside a model definition.Implementing dynamic rescaling using L2 moment (or any other statistic) is then a simple combination of get data and scale and rebalance scale method. Delayed rescaling can be similarly implemented by decoupling the compute of statistics and rebalancing in Algorithm 4, keeping track of the first one in the model state.Accurate scale propagation through an entire computational graph requires careful handling of special tensors. More specifically, we keep track in the SCALIFY tracer of two categories of tensors:• Unscaled arrays which are broadcasted scalars;• Scaled tensors indifferent to the scale value (i.e. combining 0, Infinite and NaN values);Both type of tensors are commonly generated during the tracing of an ML graph (see for instance the tracing of the method logsoftmax in Appendix B). As they get incorporated with other scaled arrays in scaled operations, the tracer requires this additional metadata to estimate the optimal output scaling.In the first case, broadcasted scalars are commonly generated by the JAX (or PyTorch) graph tracer from scalar constants (with potential intermediate unary operations applied). When combined with scaled array inputs, the SCAL-IFY tracer needs to be able to estimate a scale (or otherwise will throw an error). In the general case, this would require gathering expensive statistics, but in the case of a broadcasted scalar, it is trivial: it simply corresponds to a split between the mantissa and the exponent in the floating point representation.In the second case, scaled tensors invariant to the scale value,  
experiments and results: 
 In this section, we assess the effectiveness of SCALIFY approach for training GPT-style LLMs. Our experiments focus on the training of a GPT2 (Radford et al., 2019 )-like model with 168M parameters on the WikiText-103 dataset (Merity et al., 2017) . The configuration of the decoder-only Transformer (Vaswani et al., 2017) architecture and training hyper-parameters are detailed in Table 2 . Our JAX implementation is heavily inspired by open-source NanoGPT repositories (Karpathy, 2023; Garcia, 2023) .We aim to address three main questions with our GPT2 training experiments:• Is SCALIFY a drop-in replacement of loss scaling strategies for pure FP16 training?• Does FP8 SCALIFY training accuracy match higher precision, with minimal dynamic rescaling introduced?• Can master weights and optimizer state be stored in scaled FP16 using SCALIFY?We summarize the configurations of our main experiments in Table 3 , and present in Figure 1 the profile of training losses. In all experiments, scale propagation is performed end-to-end using SCALIFY, following the simplified training loop presented in Algorithm 1. Our first experiment SCALIFY FP16 #1 shows that SCALIFY can be used as drop-in replacement of loss-scaling for FP16 training. In this setting, the main advantage of SCALIFY is to avoid the addition of another hyper-parameter (loss scaling), or a complex automatic loss scaling strategy (i.e. no dynamic rescaling of activations, gradients or state is required). We believe that even this most simple SCALIFY setting could be of interest to machine learning practioners, as a drop-in replacement of BF16 while benefitting from additional precision brought additional mantissa bits in the IEEE FP16 format.For our experiments #2, #3 and #4, we added dynamic rescaling on the backward pass of LayerNorm layers in every transformer block. The main motivation and justification is the following: LayerNorm layers are implicitly dynamically re-scaling activations on the forward pass (see Appendix C.2, ensuring that ScaledArray inputs to the multi-head-attention and feed-forward residual paths are well calibrated. As a consequence, FP8 matmuls on the forward pass have already accurately scaled inputs. On the other hand, no operation in the computational graph is providing (implicit) dynamic rescaling of gradients, meaning that despite the use of unit-scaling type rules described in Section 3.2, the distribution of the gradients will tend to shift away from the ideal unit-scaling property E[X 2 d ] ≃ 1 in the backward pass graph. As a consequence, it is natural to insert some dynamic rescaling operations on gradients in the backward pass, and we decided to follow a similar strategy to forward pass normalization by adding one in each residual path of the Transformer layer. Algorithm 5 presents the simple implementation of the custom LayerNorm layer with dynamic rescaling of input gradient.With the additional LayerNorm gradient dynamic rescaling, experiment SCALIFY FP16 #2 shows similarly to the work of Peng et al. (2023) that per-tensor scaling FP16 can be used for master weight representation without loss in accuracy, optimizing memory footprint during training.Experiments SCALIFY FP8 #3 and #4 aim to demonstrate that SCALIFY supports as well out-of-the-box low precision FP8 training. In the first one, we use FP8 to speed-up matrix multiplications and reduce memory footprint of the weight gradients. Similarly to the previous FP16 training experiment, matching higher precision training only requires 2 dynamic rescaling in every Transformer layer instead of 18 (all forward and backward matrix multiplications) in common FP8 training approaches such as NVIDIA (2022).Finally, as presented in Peng et al. (2023) , we show that reducing the optimizer state precision to FP16 can also be achieved with SCALIFY, at the cost of adding additional dynamic rescaling of the state gradients. This is in line with the findings of the latter work which also presents a dynamic rescaling strategy of FP8 gradients, called "autoscale", in order to match higher precision accuracy. We believe this last experiments highlights one of the main benefit of SCALIFY in terms of usability and expressibility for ML practitioners: it only requires a 2 lines code change (cast of the optimizer state and 'dynamic rescale' of gradients) to integrate low precision optimizer state representation in a training loop.In our experiments, dynamic rescaling of the model state  
related work: 
 Most large models are trained with mixed precision (Micikevicius et al., 2017; NVIDIA, 2019) where half precision is used for most operations but sensitive operations such as gradient accumulation and softmax operations are maintained in FP32. Generally loss scaling is required to ensure gradients remain stable. Stability issues have also been noted when using FP16 for training large models (Zhang et al., 2022; Brown et al., 2020) . Therefore, bfloat16 is often preferred for its increased dynamic range (Kalamkar et al., 2019; Almazrouei et al., 2023) and block scaled data formats have been developed to increase dynamic range while maintaining low precision (Rouhani et al., 2023) . It is generally required to keep the optimizer state in high precision. However, there has been work reducing this to 8 bits by employing blockwise quantization methods (Dettmers et al., 2022b) . Some recent methods have gone further to reduce mixed precision down to FP8 for many operations (Peng et al., 2023; NVIDIA, 2022) , however often complex heuristics are required to determine when low precision can be used or overhead is introduced by empirically rescaling tensors to maintain stability. The work most closely related to ours is Unit Sclaing (Blake et al., 2023) which makes similar assumptions about how scale propagates. Unit scaling uses these assumptions to change the model, whereas SCALIFY maintains separate scaling factors alongside a consistently-scaled representation, leaving the model definition unchanged.Many methods have been developed to quantize LLMs for interference (Bondarenko et al., 2021; Frantar et al., 2023) . These methods require varying degrees of finetuning or knowledge distillation to achieve similar performance to their high precision baselines and there is usually a tradeoff between degree of quantization and inference performance. Quantization to 8 bits (Xiao et al., 2024 ) can be achieved with little drop in accuracy and several methods have been developed to quantize more aggressively to 4 bits (Lin et al., 2024; Zhao et al., 2024) or beyond (Kim et al., 2024) . Although these extremely low precision methods generally lead to accuracy loss compared to FP16 baselines (Wu et al., 2023) , they allow extremely efficient deployment of large models. Efficient hardware support is generally only available for quantization to powers of 2. Nonetheless, work has been done to alleviate this issue by developing kernel design schemes with provide support for arbitrary quantization bit widths (Xia et al., 2024) . 
conclusion: 
 In  
 
 Appendix: Independently of per tensor scaling, the recent introduction of FP8 formats E4M3 and E5M2 requires a more general definition of the Linear layer used in neural net models. More specifically, FP8 matrix multiplications differ in two ways from classic 16-bit and 32-bit matrix multiplications:• Mixed inputs: FP8 hardware supports matrix multiplication between mixed E4M3 and E5M2 inputs, on the contrary to 16-bit matmuls (i.e. no mixed FP16/BF16);• Higher precision output: accumulation is performed in higher precision than FP8 in machine earning hardware, meaning FP8 matmuls output by default higher precision than FP8 (FP16 or BF16 typically).This more general setting is a consequence of the machine learning research (Noune et al., 2022; Peng et al., 2023 ) on low-precision training which has demonstrated that neural net activations and backward gradients have different statistical distribution, and thus, require different FP8 formats (E4M3 and E5M2 respectively).As a consequence, we present in Algorithm 6 a more general definition of a Linear layer, supporting any kind of compute format. As presented in the code, it only requires a more granular parametrization of the data types used in the forward activation and weight, and the backward gradient, which can be done using casting operators cast on forward and cast on backward that act respectively only the forward or backward passes.As mentioned above, Algorithm 6 general Linear layerGraphcore, UK. Correspondence to: Paul Balanca <paulb@graphcore.ai>.Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML2024).Submission and Formatting Instructions for the WANT@ICML 2024 


Decision: I recommend acceptance.

To encourage more researchers to use Scalify, I encourage Authors:
1. Discuss and show more results in the future, especially the large model runs.
2. Discuss precision w.r.t Attention mechanisms (standard Softmax Attention, and other new attentions, e.g. Linear Attention, Mamba etc, as these linear attentions are more sensitive with precision issues)
3. Discuss if Scalify can handle training/evaluation in case of state sharding (especially with tensor parallel) without too much extra effort.