1 Introduction↩︎

Pre-training large-scale transformers is hard, for its intensive computation and time-consuming process [1]. To accelerate training, sparsity-based methods have recently emerged as a promising solution, and one of the hardware-friendly sparse patterns is 2:4 sparsity. In a 2:4 sparse matrix, every four consecutive elements contain two zeros. Within a tensor core, a 2:4 sparse matrix multiplication (2:4-spMM) could be 2x faster than its dense equivalent on NVIDIA Ampere architecture GPUs.

Some works use 2:4 sparsity for accelerating training [2]–[5]. However, they mainly target on convolutional neural networks (CNNs) [2], [4], whose architecture, optimizer and training procedure are different from transformers. Whether these 2:4 sparse training methods are capable for transformers remains under-explored. In practice, we find two barriers: 1) Low accuracy. The hyperparameters in some accuracy preserving techniques for transformers vary significantly from that for CNNs, which is ineffective if transplanted directly. Remarkably, simply halving the inner dimensionality of a feed-forward network can also reduce the same amount of computational cost, but provides better performance than most of proposed 2:4 sparse training methods. 2) Inefficiency. All previous works on 2:4 training stay on simulation, and do not provide actual acceleration results. Besides, they don’t focus on other key operations beyond matrix multiplication that affect the practical time cost, such as overheads of pruning and activation functions. They usually lead to substantial mismatches between simulation and actual acceleration performance.

In this work, we aim to propose an end-to-end acceleration method for pre-training transformers based on 2:4 sparsity. Here are our major contributions:

2 Related Work↩︎

Existing sparsity-based methods can be classified into two categories: accelerating inference and accelerating training. For training acceleration, they can be further grouped by whether 2:4 sparsity is involved.

Sparsity for Inference Acceleration. Early methods include one-shot pruning ([6]–[9]. Later methods ([10]–[12]) suggest using dynamic sparse training (DST). Particularly, [11] proposes sparse-refined straight-through estimator (SR-STE) for 2:4 inference. Iterative magnitude-based pruning (IMP) methods [13]–[15], originated from the winning lottery ticket theory [16], [17], can also be viewed as a DST approach. All these methods only speedup the forward pass. They are insufficient to accelerate training.

2:4 Semi-Structured Sparsity for Training Acceleration. Accelerating training by 2:4 sparsity is hard, because both the forward and backward passes need to be accelerated. On some GPUs involving sparse tensor cores, 2:4-spMMs perform 2x faster than dense GEMMs [9], [18]. In light of this, [2] firstly proposes a transposable N:M mask to accelerate both output activations and input gradients computation in backward pass. [19] improve transposable mask to bi-directional mask (Bi-Mask) to further boost mask diversity. To accelerate calculating weight gradient via 2:4-spMM, an unbiased minimum-variance estimator (MVUE) is introduced [5]. In addition, [20] also achieve fully sparse training of CNNs using spatial similarity. However, all these works do not report end-to-end training speedups on 2:4 sparse tensor cores, and they are built for CNNs. Practical 2:4 training acceleration on transformers has not been reported so far.

Other Structured Sparsity for Training Acceleration. Structured sparsity means channel-wise pruning to dense networks. For instance, training a large model and then compressing it to be thinner or shallower seems effective [21], [22], given a fixed accuracy requirement. However, it’s not memory-efficient due to the larger model’s redundancy. In addition, low-rank adaption proves to be an effective method to reduce fine-tuning costs [23], but it can’t accelerate the pre-training.

3 Preliminary↩︎

In this section, we first present the mathematical formulations of dense training and fully sparse training. Afterward, we revisit the related methods which are helpful to achieve fully sparse training with 2:4 sparsity, including SR-STE [11], transposable N: M mask[2], and MVUE[5].

3.1 Dense Training↩︎

3.1.0.1 Problem Formulation

Dense training solves an optimization problem \(\min _{\boldsymbol{\mathbf{w}}} \mathcal{L}(\boldsymbol{\mathbf{w}})\), where \(\mathcal{L}\) is a loss function, \(\boldsymbol{\mathbf{w}}\in \mathbb{R}^D\) is the collection of dense weights of all layers, flattened to a vector. The loss is optimized by gradient descent optimization algorithms such as SGD, Adam [24] and AdamW [25].

3.1.0.2 GEMMs of a Linear Layer in Dense Training

In each training step, a single linear layer performs three general matrix multiplications (GEMMs): \[\begin{align} \label{f:FWD} \boldsymbol{\mathbf{Z}}=\boldsymbol{\mathbf{X}}\boldsymbol{\mathbf{W}}^{\top},~~~\nabla_{\boldsymbol{\mathbf{X}}}=\nabla_{\boldsymbol{\mathbf{Z}}} \boldsymbol{\mathbf{W}},~~~\nabla_{\boldsymbol{\mathbf{W}}}=\nabla_{\boldsymbol{\mathbf{Z}}}^\top \boldsymbol{\mathbf{X}}, \end{align}\tag{1}\] where \(\boldsymbol{\mathbf{X}}, \boldsymbol{\mathbf{W}}\) and \(\boldsymbol{\mathbf{Z}}\) are input activations, weights, and output activations, with shape \(\boldsymbol{\mathbf{X}}, \nabla_{\boldsymbol{\mathbf{X}}} \in \mathbb{R}^{p \times q}\), \(\boldsymbol{\mathbf{W}}, \nabla_{\boldsymbol{\mathbf{W}}} \in \mathbb{R}^{r \times q}\), and \(\boldsymbol{\mathbf{Z}}, \nabla_{\boldsymbol{\mathbf{Z}}} \in \mathbb{R}^{p \times r}\). Here, the three GEMMs computes output activations, input activation gradients, and weight gradients, respectively. Without loss of generality, we assume the input \(\boldsymbol{\mathbf{X}}\) to be a 2D matrix rather than a 3D tensor. In the feed-forward networks of a transformer, this can be done by simply flattening the input tensors’ first two axes, i.e., axes of batch size and sequence length.

3.2 Fully Sparse Training with 2:4 Sparsity↩︎

GEMMs can be accelerated with structured sparsity methods. Particularly, 2:4 sparsity [9] is a semi-structured sparsity pattern supported on NVIDIA Ampere architectures. A 2:4 sparse matrix partitions its elements into groups of four numbers, where each group has exactly two zeros. Depending on the direction of partition, there are row-wise 2:4 sparse matrix and column-wise 2:4 sparse matrix (Appendix 9.1). With such sparsity, a GEMM \(\boldsymbol{\mathbf{C}}=\boldsymbol{\mathbf{A}}\boldsymbol{\mathbf{B}}\) can be accelerated by 2x with the 2:4-spMM kernel if either \(\boldsymbol{\mathbf{A}}\) is row-wise 2:4 sparse, or \(\boldsymbol{\mathbf{B}}\) is column-wise 2:4 sparse (Appendix 9.3).

To accelerate training, each GEMM in Eq. (1 ) should have one 2:4 sparse operand. In general, weights and output activation gradients are selected to be pruned due to relatively lower pruning-induced loss [5]. That is, \[\label{f:PFWD} \boldsymbol{\mathbf{Z}}= \boldsymbol{\mathbf{X}}S_{wt}( \boldsymbol{\mathbf{W}}^{\top} ),\tag{2}\] \[\label{f:PBWD-Actv} \nabla_{\boldsymbol{\mathbf{X}}} = \boldsymbol{\mathbf{Z}}S_w(\boldsymbol{\mathbf{W}}),\tag{3}\] \[\label{f:PBWD-Weight} \nabla_{\boldsymbol{\mathbf{W}}} = S_z(\nabla_{\boldsymbol{\mathbf{Z}}}^{\top} ) \boldsymbol{\mathbf{X}}.\tag{4}\] In Eq. (2 )-(4 ), \(S_{wt}, S_{w}\), and \(S_{z}\) represent the pruning functions of \(\boldsymbol{\mathbf{W}}^{\top},\boldsymbol{\mathbf{W}}\), and \(\nabla_{\boldsymbol{\mathbf{Z}}}^\top\). They take a dense matrix as input, and outputs a 2:4 sparse matrix. By intuition, a pruning function should pick out the 2 elements with the max magnitudes in the adjoining 4 elements and zero out the rest. With hardware support, computing Eq. (2 )-(4 ) can be theoretically 2x faster than computing Eq. (1 ). This method use 2:4-spMMs for all matrix multiplications in forward and backforward propagation, so we call it fully sparse training (FST). Note that Eq. (4 ) contains a straight-through estimator (STE). We will later explain it in Section 3.

3.2.0.1 Transposable Masks

[2] suggest that a weight matrix and its transpose can be simply pruned by multiplying binary masks, i.e., \[\begin{align} S_{wt}( \boldsymbol{\mathbf{W}}^{\top} )=\boldsymbol{\mathbf{W}}^{\top} \odot \boldsymbol{\mathbf{M}}_{wt},~~~S_{w}( \boldsymbol{\mathbf{W}})=\boldsymbol{\mathbf{W}}\odot \boldsymbol{\mathbf{M}}_{w}, \end{align}\] where \(\boldsymbol{\mathbf{M}}_{wt}, \boldsymbol{\mathbf{M}}_{w} \in \{0,1\}^{p \times q}\) are 2:4 sparse, and \(\odot\) is element-wise product. To utilize 2:4-spMM, the two binary masks should be mutually transposable: \[\label{f:transmask} \boldsymbol{\mathbf{M}}_{wt} = \boldsymbol{\mathbf{M}}_{w}^{\top},\tag{5}\] which they call as transposable masks (same as our defination in 5.1). In this manner, the backward pass share the same sparse weight matrix with the forward pass. In addition, the authors propose a 2-approximation method for generating such masks with claimed low computational complexity.

3.2.0.2 Minimum-Variance Unbiased Estimator

[5] propose to calculate the 2:4 sparse masks of neural gradients by MVUE, i.e., \[\label{f:mvue} S_z(\nabla_{\boldsymbol{\mathbf{Z}}}^{\top} ) = \operatorname{MVUE}(\nabla_{\boldsymbol{\mathbf{Z}}}^{\top}).\tag{6}\] Compared to the commonly used minimum square error estimation, MVUE guarantees unbiasedness and minimizes the variance of the sparsified gradients, which is more favorable for promoting the convergence of training.

3.3 Optimization Strategies for Sparse Training↩︎

The optimization of a sparse network is difficult as it has non-differentiable pruning functions. The optimization objective can be formulated as \(\min _{\mathbf{w}} \mathcal{L}( \mathbf{\tilde{w}})\). The network makes prediction with a sparse weight vector \(\mathbf{\tilde{w}}=\mathbf{m}(\mathbf{w}) \odot \mathbf{w}\), where the mask \(\mathbf{m}(\mathbf{w}) \in \left\{ 0,1 \right\}^{D}\) is the concatenation of masks for each layer. If a layer is not sparsified, then the corresponding mask is an all-one matrix. Computing the gradient is tricky since the mask \(\boldsymbol{\mathbf{m}}\) is dynamically computed based on the dense weight \(\boldsymbol{\mathbf{w}}\): by chain rule we have \(\nabla_{\mathbf{w}} \mathcal{L}(\mathbf{\tilde{w}})=\frac{\partial \tilde{\boldsymbol{\mathbf{w}}}}{\partial \boldsymbol{\mathbf{w}}}\nabla_{\mathbf{\tilde{w}}} \mathcal{L}(\mathbf{\tilde{w}}),\) where \(\frac{\partial \tilde{\boldsymbol{\mathbf{w}}}}{\partial \boldsymbol{\mathbf{w}}}\) is a Jacobian matrix. However, \(\tilde{\boldsymbol{\mathbf{w}}}\) is not differentiable with \(\boldsymbol{\mathbf{w}}\) since it includes a non-differentiable mask-computing-function \(\boldsymbol{\mathbf{m}}(\cdot)\) in it. Thus, it takes some skills to estimate the gradients and update the parameters.

3.3.0.1 STE

As \(\tilde{\boldsymbol{\mathbf{w}}}\) is an approximation of \(\boldsymbol{\mathbf{w}}\), a straight-through estimator (STE, [26]) directly passes the gradient of \(\tilde{\boldsymbol{\mathbf{w}}}\) to \(\boldsymbol{\mathbf{w}}\): \[\nabla_{\mathbf{w}} \mathcal{L}(\mathbf{\tilde{w}}) \gets \nabla_{\mathbf{\tilde{w}}} \mathcal{L}(\mathbf{\tilde{w}}).\]

3.3.0.2 SR-STE

There is a problem with STE: only a portion of the weights in a layer participate in the forward calculation, but all the weights receive gradients. This indicates that the gradients associated with masked weights¹ might be inaccurate. To suppress those inaccurate gradients, [11] proposes sparse-refined straight-through estimator (SR-STE) which adds a decay term when updating: \[\begin{align} \label{eq:sgd} \mathbf{w}_{t} \gets \mathbf{w}_{t-1}-\gamma(\nabla_{\mathbf{w}} \mathcal{L}_t({\mathbf{\tilde{w}}_{t-1}}) + \lambda_W (\overline{\mathbf{m}(\mathbf{w}_{t-1})}) \odot \mathbf{w}_{t-1}), \end{align}\tag{7}\]

where \(\gamma\) stands for the learning rate, \(\lambda_W\) is the decay factor, and \(\overline{\mathbf{m}(\mathbf{w}_{t-1})}\) denotes the logical not operation of \(\mathbf{m}(\mathbf{w}_{t-1})\). This decay term alleviates the change of weight mask. With SR-STE, the optimization target becomes \[\label{target} \min _{\mathbf{w}} \mathcal{L} (\mathbf{\tilde{w}})+\tfrac{\lambda_W}{2} \Vert \mathbf{w} \odot \overline{\mathbf{m}(\mathbf{w})}\Vert_2^2.\tag{8}\]

4 Accuracy Preserving Techniques↩︎

While the methods reviewed in Sec. 3 can successfully perform FST on small-scale models such as ResNets and DenseNets, it is not yet clear whether they can be directly applied to pre-train large transformers. It is challenging for FST to preserve the accuracy of dense training, since the weights and masks need to be learned jointly, which is a non-differentiable, combinatorial optimization problem. Moreover, unlike inference acceleration methods, FST has no pre-trained dense model to start with. In this section, we propose three practical techniques to improve the convergence of FST for transformers: transformer-specific masked decay and dense fine-tuning.

4.1 Flip Rate: Stability of Training↩︎

Inspired by previous work [11], [15], we define a “flip rate” to measure how frequently the mask vector changes after one optimizer step. This metric could be used to monitor whether the network connection is stable during training.

[15] suggest that a sparse neural network acts differently in different training phases. In the early phase of training, it eagerly explores different connection modes, which means the masks vector change rapidly over time. Later, the masks gradually become stable, and the network turns itself to fine-tune weight values. In terms of flip rate, we hypothesize that

A healthy training process comes with the flip rate \(r_t\) rising at the beginning of training and then gradually fading to \(0\).

We measure flip rate change for dense training, STE and SR-STE with different \(\lambda_W\) in 1. For dense training, we compute the flip rate by pruning the dense weight in each iteration, despite the pruned weight is never used for training. In terms of flip rate, dense training is healthy: its \(r_t\) exactly increases first before declines. If a training process consistently has higher flip rate than dense training, which we call as “flip rate explosion”, it may suffer from a loss in final accuracy due to unstable training. In practice, STE suffers from a flip rate explosion, while SR-STE takes effect by “freezing” masks of weights: by adding a decay term, it decrease the number of flips. This inhibition effect is related to the decay factor of SR-STE: the larger \(\lambda_W\) is, the stronger the inhibition of flips is, and the smaller flip rate goes.

In this section, all methods we propose involve our ultimate principle: the peak of the curve should be sufficiently high to fully explore different connection modes, and the tail should be sufficiently low for the optimization process to converge.

4.2 Transformer-Specific Masked Decay↩︎

Based on our insights on flip rate, we propose a method to suppress the frequent change of masks during FST for transformers, which we call masked decay. Unlike Eq. (7 ) which imposes regularization directly on weights, we propose to add masked decay on gradients, i.e., \[\begin{align} \label{eq:adam-way} \mathbf{g}_t & \gets \nabla_{\mathbf{w}} \mathcal{L}_t(\mathbf{\tilde{w}}_{t-1}) + \lambda_W (\overline{\mathbf{m}(\mathbf{w}_{t-1})} \odot \mathbf{w}_{t-1}). \end{align}\tag{9}\] On SGD, applying the decay term on weights and on gradients are equivalent, while on popular optimizers like Adam and AdamW they are not. Specifically, Adam updates weights by \[\begin{align} \label{long} \boldsymbol{\mathbf{w}}_t \gets \boldsymbol{\mathbf{w}}_{t-1}-\frac{\gamma (\beta_1 \boldsymbol{\mathbf{u}}_{t-1}+(1-\beta_1)\boldsymbol{\mathbf{g}}_t)}{(1-\beta_1^t)(\sqrt{\hat{\boldsymbol{\mathbf{v}}}_t}+\epsilon)} \end{align}\tag{10}\] where \(\boldsymbol{\mathbf{u}}\) and \(\boldsymbol{\mathbf{v}}\) are the first and second order momentum of \(\boldsymbol{\mathbf{w}}\). Compared to Eq. (7 ), the masked decay regularization term in Eq. (9 ) would be later normalized by \(\sqrt{\hat{\boldsymbol{\mathbf{v}}}_t}+\epsilon\) in Eq. (10 ), before it is subtracted from weights. In this way, each dimension receives a different intensity of decay (“masked decay”). More specifically, weights with larger gradients get smaller decay intensity, and vice versa.

In FST, we periodically prune weights by their magnitudes. STE may cause the network to fall into such “dilemma points”, where a portion of pruned weights and unpruned weights have nearly the same total L1 norm. Thus, the network consistently oscillate between two possible masks \(\boldsymbol{\mathbf{m}}_1\) and \(\boldsymbol{\mathbf{m}}_2\), and it is unlikely to jump out the dilemma point itself. This process can be observed by the flip rate explosion of STE. On these occasions, we argue that an evenly masked decay applied on weights is insufficient to save the training from such “traps”. The weights don’t differentiate themselves after an update, so masks may oscillate back. By normalizing the weight gradients with \(\sqrt{\hat{\boldsymbol{\mathbf{v}}}_t}+\epsilon\), our masked decay amplifies the regularization strength for the dimension with smaller gradient, pushing it towards zero. Then, the regularized dimension can no longer compete with other dimensions. So we effectively break the tie and push the training process out of the trap, towards a “healthier” state.

The comparison results between our masked decay defined in Eq. (9 ) and the conventional counterpart in Eq. (7 ) are shown in 2. Results show that applying masked decay on weights takes no effect to inhibit flip rate explosion of STE while applying on gradients works well enough.

4.3 Fast Decay Factor Determination↩︎

The determination of the decay factor \(\lambda_W\) in Eq. (9 ) is non-trivial: if \(\lambda_W\) is excessively large, then the “peak” of the flip rate curve is not high enough; if \(\lambda_W\) is too small, the “tail” of the curve is not low enough. Both do not provide a healthy training process. Besides, we find that \(\lambda_W\) values for CNNs and other small-scale networks differ significantly from those for transformers, while on transformers, optimal \(\lambda_W\) can span up to three orders of magnitude (1).

-0.1in

As pre-training large transformers is costly, grid searching for \(\lambda_W\) with the final accuracy is impractical, so it is vital to determine a feasible \(\lambda_W\) as quickly as possible. To quickly determine \(\lambda_W\), here we propose a test-based method:

4.4 Dense Fine-Tuning↩︎

To better improve accuracy, we suggest using a “dense fine-tuning” procedure at the end of training. Formally, we select a switch point \(t_s\). FST is performed while \(t \leq t_s\), and dense training is switched to if \(t > t_s\).

4.4.0.1 Why Choose Dense Fine-Tuning Instead of Dense Pre-training?

Some recent works like STEP [3] utilize dense pre-training rather than dense fine-tuning, which means a dense network is initially trained for a period of time before being switched to a sparse one. However, we argue that dense pre-training is meaningless in our FST process. As described in 4.1, the peak of the flip rate curve should be sufficiently high to explore connection modes, so what matters most to the flip rate is the magnitudes of weights, which are the key to determine if connections are built or demolished. In this regard, both FST and dense pre-training are capable of delivering proper gradient magnitudes, so dense pre-training is a waste. The precise gradients are generally more necessary in the later stages of training, where the flip rate of the dense network comes to its tail. 3 visualizes the loss curve of pre-training BERT-base, where dense pre-train obtains nearly the same result as the naive SR-STE method. From this, we propose the following insight:

If dense pre-training of \(t_\alpha\) steps provides slight improvement of accuracy, then moving the \(t_\alpha\) dense steps to the end gives far more improvement than dense pre-training.

As for the specific position of the switch point in training, STEP [3] suggests that the dense pre-training occupy \(10\%\) to \(50\%\) of the total steps. Likewise, we determine that our dense fine-tuning takes up the last \(1/6\) of total steps for balance training efficiency and accuracy.

5 Training Acceleration Techniques↩︎

For transformers, the forward pass of our FST involves pruning weights in feed-forward networks with transposable 2:4 masks and then performing normal forward propagation. During the backward propagation in our FST, the gradients of input activations and weight gradients in FFNs are derived by Eq. (3 ) and (4 ), respectively. Note that we also utilize MVUE to prune gradients of output activations, i.e., Eq. (6 ). Compared to dense training, our FST replaces all the GEMMs in FFNs with 2:4-spMMs that theoretically perform 2x faster than their dense counterparts on GPUs within sparse tensor cores.

In addition to speeding up the most time-consuming GEMMs in FFNs, there are three major operations that also have non-negligible impacts on training speed:

Below, we show our methods to accelerate these operations.

5.1 Fast Computation of Transposable Masks↩︎

5.1.0.1 Problem Formulation

We aim to find such a mask matrix \(\boldsymbol{\mathbf{M}}\in \{0,1\}^{r \times q}\) for every \(\boldsymbol{\mathbf{W}}\in \mathbb{R}^{r \times q}\) in the FFN layer that 1) each adjoining \(4 \times 4\) block contains 8 non-zero positions; each row and column in the block occupies 2 non-zero elements exactly (see in 9.2); 2) \(\max _{\boldsymbol{\mathbf{M}}} \Vert\boldsymbol{\mathbf{M}}\odot\boldsymbol{\mathbf{W}}\Vert_1\). Then \(\boldsymbol{\mathbf{M}}\) would be our targeting transposable mask.

As described in Eq. (5 ), both a transposable mask itself and its transposition conform to the format of 2:4 sparsity. Previous 2-approximation algorithm [2] consists of two steps: sort elements, and pick elements out of the array. They claim that the procedure has less computational complexity. However, in practice, the sorting and picking process contains too many jumps in its control flow, and may be fatal to modern GPU architecture. To make full use of the GPUs’ parallel computation capability (SIMD and SIMT), we convert the transposable mask-search process into a convolution operation which traverse all the masks to obtain the optimal one in three steps:

Notably, step (1) is executed offline. Step (2) and (3) are frequently performed during FST. The workflow of our method is shown in 5. Compared to the 2-approximation algorithm, our method is up to about 5 times faster (2).

-0.1in

5.2 Acceleration of Gated Activation Functions↩︎

Activation functions with gated mechanisms are widely used in transformers such as GLM [28] and LLaMA [29]. Typical gated activation functions involve SwiGLU and GEGLU. The bottleneck of such activation functions is that the gate operations easily incur GPU L2 cache miss. Take GEGLU as an example: \(\operatorname{GEGLU}(\boldsymbol{\mathbf{X}},\boldsymbol{\mathbf{U}},\boldsymbol{\mathbf{V}},\boldsymbol{\mathbf{b}}, \boldsymbol{\mathbf{c}}) = \operatorname{GELU}(\boldsymbol{\mathbf{X}}\boldsymbol{\mathbf{U}}^{\top} +\boldsymbol{\mathbf{b}}) \odot (\boldsymbol{\mathbf{X}}\boldsymbol{\mathbf{V}}^{\top} +\boldsymbol{\mathbf{c}})\), where \(\boldsymbol{\mathbf{X}}\in \mathbb{R}^{p \times q},\boldsymbol{\mathbf{U}}, \boldsymbol{\mathbf{V}}\in \mathbb{R}^{r \times q},\boldsymbol{\mathbf{b}},\boldsymbol{\mathbf{c}}\in \mathbb{R}^{r}\). In practice, this function is composed of three steps:

Different from the dense model, where output activations are row-major matrices, in FST, the output activations are column-major matrices. This property results in the third step being extremely time-consuming if conventionally \(\boldsymbol{\mathbf{Z}}\) is accessed along the row dimension. To illustrate, 6 shows that in a column-major matrix \(\boldsymbol{\mathbf{Z}}\), accessing along the column accords with memory layout. Thus, adjacent elements loaded into the GPU cache can be probably hit. By contrast, accessing along the row does not fully utilize the efficiency of GPU cache. In light of this, we carefully implement a GEGLU kernel where elements are accessed along the column dimension. In this way, GEGLU is performed 5 times faster than the naive counterpart (3).

-0.1in

5.3 Other Implementation Details↩︎

Reducing updating frequency. In fact, we find that a 2:4 mask doesn’t change a lot after one optimization step, and it is not necessary to update a mask frequently. For the sake of efficiency, we update the transposable masks of weights every \(l\) optimizer steps. We usually take \(l=40\) in practice.

Utilities. For 2:4-spMMs, we use CUTLASS [30] and cuSPARSElt [9]. Other GPU kernels are implemented in Triton, including a transposable mask search kernel, a pruning kernel, an MVUE kernel, a GEGLU kernel, and a masked decay kernel.

6 Experiments↩︎

In this section, we validate the proposed training speedup methods on several transformers, including BERT [31], GPT-2 [32], Transformer-base for machine translation [33], and DeiT [34]. For BERT, we use Cramming [35] to pre-train a 16-layer BERT model with the sequence length of 512 on the C4 dataset [36]. For GPT-2, we use nanoGPT [37] to pre-train GPT-2 124M, 355M, 774M, and 1.5B on OpenWebText [38]. Both BERT and GPT-2 models are estimated on GLUE [39]. For DeiT [40], we pre-train DeiT-tiny on ImageNet-1K dataset [41]. Besides, we use fairseq [42] to train Transformer-base on the WMT 14 En-De dataset [43] and measure the BLEU [44] score of the trained model.

Of note, we use \(n\) to denote the length of sequences, \(d\) to denote the input and output dimensions of each transformer block, \(d_{ff}\) to denote the inner dimensions of the FFNs in each transformer block, \(h\) to denote the number of heads, and \(N\) to denote the micro-batch size on each device.

6.1 Accuracy Results↩︎

To investigate the effect of different 2:4 sparse training methods, we pre-train a sparse BERT-base model on the C4 dataset using two sparse training methods: STEP [3] and Bi-Mask [19]. Besides, we also pre-train a dense BERT-base and a ‘Half’ BERT-base for comparison. Of note, ‘Half’ denotes a smaller yet still dense BERT-base model. To create Half model, we simply reduce the \(d_{ff}\) of each FFN layer in the original BERT-base by half while maintaining the original value of \(d\). Theoretically, this adjustment halves the floating operations (FLOPs) of the original BERT-base as well. Except for the FFN layers, the shapes of the rest layers remain unaltered.

All the pre-trained models are measured on GLUE benchmark (WNLI excluded). Surprisingly, 4 shows that despite having identical FLOPs, the 2:4-sparse BERT-base trained with STEP and Bi-Mask shows inferior average scores compared to the Half model. The Half model attains an average score of 77.9 on GLUE tests, while STEP and Bi-Mask only reach 77.7 due to the weaknesses in MRPC, QNLI, and STSB. By comparison, BERT-base trained in our proposed training method achieves 79.6 on GLUE, which significantly outperforms other sparse training methods and is comparable with the dense baseline, i.e., 79.8.

-0.1in

-0.1in

Besides, we pre-train the GPT-2 series models with the proposed method. Table 5 shows that our method achieves lossless average scores compared with dense baselines for model sizes of 124M, 350M, and 775M. Similarly, Transformer-base and DeiT-tiny trained with our method also reach comparable results as dense training, as seen in Table 6.

Ablation study. We aim to investigate the effect of masked decay, MVUE and dense fine-tuning introduced in 4.2, 3.2, and 4.4. The 16-layer BERT-base is used for ablation study. Results in Table 7 show that: 1) The dense fine-tuning procedure helps to improve accuracy by 2 points at most on GLUE; 2) MVUE leads to insignificant, controllable accuracy loss; 3) By combining all these techniques together, 2:4 sparse training for transformers achieves comparable accuracy results as dense training.

6.2 Speedup Results↩︎

The training acceleration techniques proposed in Section 5 are evaluated using GPT-2 models and on RTX3090 GPUs. FP16 mixed precision training is used on all models. The practical speedups of a single FFN layer, a single transformer block, and the entire network, compared to their respective dense counterparts, are reported. All the measured datum contain both forward and backward propagation.

Feed-forward network layers. For a single FFN layer, we fix \(n=2048\) and change \(d\). Results in 7 show that a FFN layer can be accelerated up to  1.5x faster than its corresponding dense layer. The speedup effect fades as the batch size becomes larger due to GPU memory bottleneck.

Transformer block. We measure the acceleration ratio of a transformer block when \(n=512,1024,2048\). Results in 7 show that in most cases, a transformer block can be accelerated to 1.2x faster via 2:4 sparsity.

-0.1in

End-to-end acceleration. Finally, we test the practical speedups of training GPT-2 models. Results in 8 show that our training method conducts up to 1.2x faster than the dense training on a single RTX3090.

7 Conclusions↩︎

In this study, we are the first to propose accelerating the pre-training of transformers by 2:4 sparsity. We analyze the limitations of previous 2:4 training methods, including the impropriety in choosing positions and determining values of the masked decay factor, speed bottleneck incurred by computing transposable masks and gated activation functions. We propose a series of techniques to tackle them. Our training method is validated on DeiT, BERT, Transformer and GPT-2 models. In particular, we have attained 1.2x end-to-end training acceleration for the GPT-2 774M model without losing its accuracy.

8 Broader Impact↩︎

Our proposed efficient algorithm can be used to accelerate pre-training large-scale transformers like GLM [28], LLaMA [29], etc. Recently, large transformers have exhibited remarkable efficacy in various fields such as natural language processing, computer vision, and speech recognition. However, the pre-training stage of large transformers is computationally intensive and time-consuming. For instance, pre-training a GPT-4 can span several months, even using a supercomputer equipped with thousands of GPUs. Thus, acceleration approaches are necessary. Our fully sparse training approach of transformers can potentially accelerate the FFN layers of a model by theoretical 2x faster, without loss of accuracy. Thus, it can be potentially used to save energy and reduce carbon footprint. But this work can also be used to accelerate baleful software, like software that generates malicious contents, which may have a negative impact on human society.

langley00?

9 2:4-spMM↩︎

9.1 2:4 Sparsity↩︎

Examples of row-wise 2:4 sparse matrix and column-wise 2:4 sparse matrix are depicted in 8.

9.2 2:4 Transposable Masks↩︎

Example of a transposable mask with 2 blocks is depicted in 9.

9.3 Memory Layout↩︎

The memory layout of different type of matrix multiplications are listed in 9, which explains why output activations and activation gradients are column-major matrices in FST.

-0.1in

10 Workflow↩︎

10.1 Forward and Backward Pass of a Linear Layer↩︎

The main workflow of a single linear layer in FST process is depicted in 10.

11 Training Loss Curve↩︎

For BERT-base and GPT-2, we depict training loss curve in 11.

References↩︎