PaperlayerView exact revision on arXiv

Source-linked AI summary

SLoPe: Double-Pruned Sparse Plus Lazy Low-Rank Adapter Pretraining of LLMs

Mohammad Mozaffari, Amir Yazdanbakhsh, Zhao Zhang, Maryam Mehri Dehnavi

arXiv:2405.16325v3cs.LGcs.AI

TL;DR

Sparse pretraining can reduce model accuracy, while dense fine-tuning can forfeit sparse inference savings. SLOPE adds late low-rank adapters and double-prunes sparse weights; it reports faster training and inference with reduced memory footprints for billion-parameter LLMs.

  • Problem

    Sparse masks can reduce accuracy, and dense fine-tuning used to compensate can disable sparse methods' inference-time memory and compute savings.

    Source: §1 · p. 2

  • Method

    SLOPE combines a double-pruned N:M backward pass with low-rank adapters introduced only in the final 1% of pretraining.

    Source: §1 · p. 2

  • Results

    1.25× training and 1.54× inference speedups, with memory footprints reduced by up to 0.63× and 0.61×, respectively, on LLMs with billions of parameters.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SLOPE improves pretraining and inference times while reducing memory footprint with negligible impact on model performance.

    Source: §4 · p. 11

    Memory reduction is slightly below theoretical expectation because other model components, including layer norms and dense parameters, also use memory.

    Source: §3.1 · p. 8

Abstract

from arXiv · show

We propose SLoPe, a Double-Pruned Sparse Plus Lazy Low-rank Adapter Pretraining method for LLMs that improves the accuracy of sparse LLMs while accelerating their pretraining and inference and reducing their memory footprint. Sparse pretraining of LLMs reduces the accuracy of the model, to overcome this, prior work uses dense models during fine-tuning. SLoPe improves the accuracy of sparsely pretrained models by adding low-rank adapters in the final 1% iterations of pretraining without adding significant overheads to the model pretraining and inference. In addition, SLoPe uses a double-pruned backward pass formulation that prunes the transposed weight matrix using N:M sparsity structures to enable an accelerated sparse backward pass. SLoPe accelerates the training and inference of models with billions of parameters up to $1.25\times$ and $1.54\times$ respectively (OPT-33B and OPT-66B) while reducing their memory usage by up to $0.63\times$ and $0.61\times$ for training and inference respectively.

1 Introduction

SLOPE targets accuracy loss from structured sparse pretraining without relying on dense fine-tuning that erases sparse inference savings. It combines double-pruned backward passes with low-rank adapters used only in the final 1% of pretraining, alongside optimized kernels.

1 Introduction

  • Structured sparsity can reduce computation and memory, but constrained patterns can lower accuracy, while adaptive masks may require dense-weight storage and extra training work.Dense fine-tuning can also eliminate sparse methods' inference-time memory and compute savings.Source: §1 · pp. 1–2
  • Its double-pruned backward pass first transposes an N:M-pruned weight matrix, then applies another N:M pruning, allowing broader patterns and reducing mask-search overhead.The authors say this improves model quality and supports an accelerated sparse backward pass.Source: §1 · p. 2
  • SLOPE adds non-zeros as low-rank adapters during only the final 1% of pretraining, aiming to restore model capacity without storing a dense inference model.The authors report faster convergence than adding the same number of learnable parameters to sparse weights.Source: §1 · p. 2
  • 1.25× training and 1.54× inference speedups, with memory footprints reduced by up to 0.63× and 0.61×, respectively, from optimized CUDA kernels.The speedups and memory reductions are reported for LLMs with billions of parameters.Source: §1 · p. 2

2 Sparse plus low-rank pretraining of LLMs

This section describes SLOPE's N:M sparse training formulation, lazy low-rank adapters, and kernel optimizations. Double pruning supports sparse computation in both passes; low-rank terms are added late, and runtime optimizations target kernel-call and shape-related overheads.

2.1 Double-pruned backward pass

  • Double pruning lets SLOPE exploit weight sparsity in forward and backward computations, where pruning along both weight dimensions supports the required reduction dimensions.The method focuses on weight tensors because their sparsification affects storage and serving, and initial findings indicate less impact on end-to-end quality than activation pruning.Source: §2 · pp. 2–3, §2.1 · p. 3
  • The authors state that lossy materialization of pruned weights has convergence properties; double pruning adds 12.5%, 9.375%, and 3.39% zeros for 1:2, 2:4, and 2:8 patterns, respectively.They also state that dynamically altering the column-wise mask does not harm optimizer convergence, while the training policy keeps the initialized mask fixed.Source: §2.1 · pp. 3–4

2.2 Lazy low-rank adapters

  • SLOPE augments the doubly-pruned matrix with a low-rank approximation to mitigate potential quality loss while retaining sparse matrix multiplication in forward and backward passes.The rank r controls trade-offs among memory footprint, computational efficiency, and model quality.Source: §2.2 · pp. 4–5
  • Low-rank adapters converge faster than sparse weights in the authors' experiments, so SLOPE uses them only during the final 1% of training iterations to reduce training operations.The paper calls this schedule lazy low-rank adapters.Source: §2.2 · p. 5

2.3 Sparse kernels

  • The sparse training algorithm initializes sparse formats for weights and their transpose, masks gradients for pruned weights, and updates both sparse representations during optimization.The described implementation assumes weight decay and includes gradient scaling for half-precision backward computation.Source: §2.3 · p. 5

2.4 SLOPE runtime optimization

  • Tiling upsample tensors into square matrices addresses a drop in sparse-multiplication speedup around hidden dimension 4000; this yields 12% faster inference and 4% faster training with SLOPE.The authors tune tile size for best performance.Source: §2.4 · pp. 5–6
  • To reduce low-rank adapter overhead, SLOPE concatenates the downsample tensor with sparse weights and uses fused matrix multiplication and addition, cutting kernel calls and cache access.Together, these optimizations improve end-to-end inference speed by up to 6%.Source: §2.4 · p. 6

3 Experimental results

Experiments assess SLOPE’s memory and speed performance across model scales and its accuracy on smaller pretrained models. Results indicate that sparse-model accuracy can improve with late low-rank adapters, while the first blocks are more sensitive to pruning than later blocks.

3 Experimental results

  • Accuracy evaluation focuses on models up to 774M parameters, while speed and memory tests span models from 2.6B to 66B parameters.For cuSPARSELt, speedup rises with larger matrices until its 2× capacity; low-arithmetic-intensity adapter multiplications add relative overhead.Source: §3 · p. 6, §3.1 · pp. 6–7

3.1 End-to-end speedup and memory saving: pretraining and inference

  • Training memory use falls 68%, while inference memory use falls 54%.The measured memory savings are slightly below theoretical expectations because other components, including layer norms and dense parameters, also use memory.Source: §3.1 · pp. 7–8

3.2 Pretraining accuracy results

  • SLOPE achieves lower perplexity than Wanda and Extended SR-STE, and higher zero-shot accuracy than Extended SR-STE; lazy adapters also benefit static and dynamic training.A perplexity gap between sparse and dense models remains.Source: §3.2 · p. 9
  • Adapters at ranks 4, 16, and 64 improve final accuracy across tasks, with higher ranks trading increased computation for better performance.Using adapters for only the final 1% of iterations is sufficient to recover pretraining accuracy.Source: §3.2 · p. 10
  • Increasing sparsity reduces accuracy across tasks, and pruning the first 12 blocks causes a larger drop, especially on GLUE; Wanda performs worse than SLOPE in these settings.The results indicate that the first blocks are more sensitive to sparsity, so later blocks can be sparsified more aggressively.Source: §3.2 · p. 11
  • SLOPE sustains competitive quality while pruning all model modules, including self-attention and MLP modules.Source: §3.2 · p. 11

4 Conclusion

SLOPE combines static N:M sparsity, a double-pruned backward pass, and lazy low-rank adapters to improve efficiency with negligible impact on model performance. The section relates these design choices to dynamic-mask overhead and the scaling behavior of low-rank adapters.

4 Conclusion

  • Up to 1.25× pretraining and 1.54× inference speedup; effective memory footprint is reduced by up to 0.63× and 0.61×, respectively.The conclusion attributes these benefits to N:M sparsity, lazy low-rank adapters in both passes, and efficient CUDA kernels.Source: §4 · p. 11

A Comparison with Dynamic Sparsity: SR-STE

  • For the same training budget, SLOPE achieves lower perplexity than SR-STE, whose changing masks spend computation on weights later pruned.The average mask difference from the converged pattern decreases over SR-STE training; Figure 4 represents the resources used to update ultimately pruned weights.Source: A Comparison with Dynamic Sparsity:… · p. 16, Figure 4 · p. 16

B cuSPARSELt Initialization Overhead: Static vs. Dynamic Sparsity

C Low-Rank Adapter Performance: Scaling and Arithmetic Intensity

  • Low-rank adapter speedups fall short of ideal linear scaling, especially at reduced ranks, but the gap narrows for larger matrices as arithmetic intensity increases.The passage attributes the larger-matrix behavior to better tensor-core utilization.Source: C Low-Rank Adapter Performance:… · p. 17, Figure 6 · p. 17

D Efficient low-rank adapter implementation

The implementation combines adapter operations to reduce kernel launches and raise arithmetic intensity.

D Efficient low-rank adapter implementation

  • Adapter concatenation with sparse weights and fusion with the subsequent result addition increase inference speedup by up to 6%.The optimizations reduce kernel calls and increase arithmetic intensity.Source: D Efficient low-rank adapter… · p. 18

E Efficient weight tiling implementation

The section reports that SLOPE retains competitive quality across different module-pruning settings and uses matrix tiling to improve speed.

E Efficient weight tiling implementation

  • Matrix tiling improves training speed by up to 4% and inference speed by up to 12%.The approach splits upsample matrices into multiple square matrices to mitigate effects of dimensions and aspect ratios.Source: E Efficient weight tiling implementation · p. 18

F SLOPE sensitivity to pruning different module in transformer

  • Pruning more transformer modules produces a consistent but slight decrease in model quality, while the results suggest relative insensitivity to which modules are pruned.The reported settings prune both MLP and self-attention modules; Table 9 presents SQuAD and GLUE results.Source: F SLOPE sensitivity to pruning… · p. 18, Table 9 · p. 19

G BERT-Large-Uncased: Pretraining and Downstream Evaluation

The BERT-Large-Uncased evaluation covers two pretraining phases and downstream SQuAD and GLUE benchmarks. Higher sparsity generally raises training loss, while the section also reports slowdowns for Bidirectional Masks relative to dense baselines.

G BERT-Large-Uncased: Pretraining and Downstream Evaluation

H Performance overhead of bidirectional mask

  • Bidirectional Masks incur mask-search and associated overheads that cause significant slowdowns against dense baselines.The experiments use the repository and models from the Bidirectional Masks evaluation.Source: H Performance overhead of… · p. 19

I Sparsity ratio analysis of double-pruned backward pass

The analysis finds that double pruning adds sparsity while generally retaining density close to the original matrix, and that static weight pruning performs better than input pruning. Pruning output gradients leads to divergence after a few iterations.

I Sparsity ratio analysis of double-pruned backward pass

J Sensitivity to the choice of pruning matrix

  • Pruning output gradients leads to model divergence after a few iterations.The authors explain that this pruning introduces errors into the backward pass.Source: J Sensitivity to the choice of pruning… · p. 20
  • Static pruning achieves lower perplexity than dynamic pruning, and pruning weights yields lower perplexity than pruning inputs.The analysis attributes the static-pruning pattern to focusing computation on elements that remain active throughout training.Source: J Sensitivity to the choice of pruning… · p. 20

K Implementation details

The implementation configures and compresses sparse matrices for cuSPARSELt operations and adds CUDA kernels for gradient masking, sparse addition, and updating sparse weights.

  • cuSPARSELt setup initializes matrix descriptors, prunes and compresses matrices, and searches for an SpMM kernel; repeated matrix multiplications mitigate setup overhead.The setup cost is described as non-negligible but offset by the repetitive operations during training.Source: K Implementation details · p. 21
  • A custom CUDA kernel applies the weight mask to gradients, removing their redundant zeros and reducing storage needs.The dense-format gradients contain 50% extra zeros after applying the same mask as the weights.Source: K Implementation details · p. 21
  • A custom kernel computes sparse linear combinations βA + γB when matrices share a sparsity pattern, supporting optimizer weight decay.Source: K Implementation details · p. 21
  • An optimized kernel copies updated PyTorch weights into cuSPARSELt format so the sparse representation reflects optimizer updates.Source: K Implementation details · p. 21

L Task-specific GLUE results

The section presents task-specific GLUE results, reports that SLoPe and FlashAttention-2 improve speed independently and further together, and situates sparse-plus-low-rank pretraining against prior work.

L Task-specific GLUE results

  • Table 11 presents task-specific results for the eight GLUE classification tasks across the training settings.The paper uses average GLUE score for overall model performance and provides the full task-level breakdown here.Source: L Task-specific GLUE results · p. 21, Table 11 · p. 22

M Integration with Flash Attention

N Additional related work

  • Prior sparse-plus-low-rank work has not examined pretraining, while related fine-tuning methods do not accelerate both passes and add adapters from the start.The passage also notes prior uses in attention heads and pruning after pretraining.Source: N Additional related work · p. 23

O Experiment setup, hyperparameters, compute resources

In the GPT2-Half comparison, SLoPe outperforms the dense model on average, and adding low-rank adapters improves all sparse pretraining methods. Experiments used multiple GPU clusters and specified measurement procedures for speed and memory.

O Experiment setup, hyperparameters, compute resources

  • Accuracy experiments used Narval, Mist, and Lonestar 6; speed and memory tests ran on one A100, with medians from 1,000 samples and five runs, respectively.The experiments used default NVIDIA BERT and FlashAttention GPT hyperparameters; further tuning for sparse pretraining is left for future work.Source: O Experiment setup, hyperparameters,… · p. 23

P Comparison with dense models

  • SLOPE outperforms GPT2-Half on average, whereas dynamic sparse training methods such as SR-STE perform worse; adding low-rank adapters improves accuracy for all sparse pretraining methods.Source: P Comparison with dense models · p. 23

Q Zero-shot GLUE results for GPT

On zero-shot GLUE tasks, SLOPE outperforms SR-STE and GPT2-Half on average; the accompanying table compares GPT models across sparsity methods and LoRA ranks.

Q Zero-shot GLUE results for GPT

R Extended SR-STE and FST implementation details

  • SLoPe is compared with Extended SR-STE for quality at equal parameter counts, while FST comparisons cover speed and memory because FST has approximately 2× more parameters.Source: R Extended SR-STE and FST… · p. 24

T.1 Lemma 2.1

  • Lemma 2.1 derives the expected added sparsity from the expected extra nonzeros in an M-element block, using the binomial distribution of Y.It defines X as added nonzeros and Y as the nonzeros in M consecutive elements, then relates their probabilities to derive the expected sparsity ratio.Source: §T.1 · pp. 25–27

T.2 Theorem 2.2

  • Theorem 2.2's proof models the mask as random, defines its expectation, and uses linearity of differentiation and expectation to obtain the result.The proof begins from the optimization objective and its backpropagated gradient; each mask entry is 1 with probability N/M.Source: §T.2 · pp. 27–28

S Comparison of Depth and Width Pruning

  • Preliminary retraining losses show no significant difference between depth and width pruning, though depth pruning appears to outperform width pruning in some cases.Experiments use LLaMA-2-7B and Gemma-2-2B and Gemma-2-9B; halving the up- and down-sample aspect ratio favors width pruning because self-attention parameters remain intact.Source: S Comparison of Depth and Width Pruning · p. 25, Figure 10 · p. 28
Loading 2405.16325v3…