  Recent innovations on hardware (e.g. Nvidia A100) have motivated learning N:M
structured sparsity masks from scratch for fast model inference. However,
state-of-the-art learning recipes in this regime (e.g. SR-STE) are proposed for
non-adaptive optimizers like momentum SGD, while incurring non-trivial accuracy
drop for Adam-trained models like attention-based LLMs. In this paper, we first
demonstrate such gap origins from poorly estimated second moment (i.e.
variance) in Adam states given by the masked weights. We conjecture that
learning N:M masks with Adam should take the critical regime of variance
estimation into account. In light of this, we propose STEP, an Adam-aware
recipe that learns N:M masks with two phases: first, STEP calculates a reliable
variance estimate (precondition phase) and subsequently, the variance remains
fixed and is used as a precondition to learn N:M masks (mask-learning phase).
STEP automatically identifies the switching point of two phases by dynamically
sampling variance changes over the training trajectory and testing the sample
concentration. Empirically, we evaluate STEP and other baselines such as ASP
and SR-STE on multiple tasks including CIFAR classification, machine
translation and LLM fine-tuning (BERT-Base, GPT-2). We show STEP mitigates the
accuracy drop of baseline recipes and is robust to aggressive structured
sparsity ratios.
