S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training
This work addresses the problem of efficient training for deep neural networks on Nvidia GPUs, offering an incremental improvement over existing 2:4 sparse pre-training methods.
The paper tackles the optimization difficulties in 2:4 sparse pre-training for deep neural networks, which arise from discontinuous pruning functions, by proposing S-STE, a method that continuously projects weights to be 2:4 sparse and rescales them, achieving results comparable to full parameter models.
Training deep neural networks (DNNs) is costly. Fortunately, Nvidia Ampere and Hopper GPUs can accelerate matrix multiplications twice as fast as a dense equivalent by implementing 2:4 sparsity. However, previous STE-based 2:4 pre-training methods (e.g. STE with hard-thresholding, SR-STE) suffer from optimization difficulties because of discontinuous pruning function. In this study, we comprehensively analyse the bottleneck of traditional N:M sparse training and recognize three drawbacks with discontinuity: incorrect descending direction, inability to predict the amount of descent and sparse mask oscillation. In light of this, we propose S-STE, a simple yet powerful 2:4 training method that contains two parts: to continuously project weights to be 2:4 sparse, and to rescale sparse weights with a per-tensor fixed scaling factor. Besides, we adopt minimum-variance unbiased estimation for activation gradient and FP8 quantization for whole process. Results show that our method surpasses previous 2:4 pre-training recipes and is comparable even with full parameter models. Our toolkit is available at https://github.com/huyz2023/2by4-pretrain.