Accelerating Transformer Pre-training with 2:4 Sparsity

TL;DR

This work comprehensively investigates the feasibility of accelerating feed-forward networks of transformers in pre-training and devise two techniques to practically accelerate training: to calculate transposable 2:4 masks by convolution, and to accelerate gated activation functions by reducing GPU L2 cache miss.

Abstract

Training large transformers is slow, but recent innovations on GPU architecture give us an advantage. NVIDIA Ampere GPUs can execute a fine-grained 2:4 sparse matrix multiplication twice as fast as its dense equivalent. In the light of this property, we comprehensively investigate the feasibility of accelerating feed-forward networks (FFNs) of transformers in pre-training. First, we define a ``flip rate'' to monitor the stability of a 2:4 training process. Utilizing this metric, we propose three techniques to preserve accuracy: to modify the sparse-refined straight-through estimator by applying the masked decay term on gradients, to determine a feasible decay factor in warm-up stage, and to enhance the model's quality by a dense fine-tuning procedure near the end of pre-training. Besides, we devise two techniques to practically accelerate training: to calculate transposable 2:4 masks by convolution, and to accelerate gated activation functions by reducing GPU L2 cache miss. Experiments show that our 2:4 sparse training algorithm achieves similar convergence to dense training algorithms on several transformer pre-training tasks, while actual acceleration can be observed on different shapes of transformer block apparently. Our toolkit is available at https://github.com/huyz2023/2by4-pretrain.

Figures (10)

• Figure 1: Flip rates change throughout the training of different $\lambda_W$ on Transformer-base. Note that these models utilize an identical learning rate schedule.
• Figure 2: Scatter plots of cumulative flip number and L1 norm gap $g_i$ on every $4 \times 4$ block. All results are selected on Transformer-base, with epoch=20. (a) shows the result of dense model. (b)-(d) shows that of masked decaying on gradients, no decaying, and masked decaying on weights. Also, we do it on purpose to choose an extremely large $\lambda_W$ for SR-STE.
• Figure 3: Applying masked decay on weights takes no effect to inhibit flip rate on BERT-base (compared to applying directly on gradient).
• Figure 4: Dense fine-tuning versus dense pre-training on BERT-base
• Figure 5: Transposable mask search

Theorems & Definitions (1)

• Definition 4.1