A Proximal Operator for Inducing 2:4-Sparsity

Thank you very much for thoroughly checking our theoretical claims and helping to improve our presentation. We will carefully take your considerations into account and will revise Section 3.3 around Corollary 6 to ensure it is easier to follow.

We want to briefly reply to some of your comments and hope that this helps for clarification. In particular we ran a new ablation study on the relevance of calibration samples that we will add to the paper.

There is no ablation study on the number of calibration samples

Thank you, this is an important point and was also raised by reviewer 5gi2. We ran a similar ablation study to wanda on the 3B model. Except when using only a single sample, our proposed prox+GD outperforms all other methods when using the same amount of calibration data. Please check our reply to 5gi2 for the full results.

In table 1, results of prox without GD do not provided.

Indeed this is on purpose, we do not want to propose prox without masked gradient as a method. This would mean that the optimization stops as soon as we have a 2:4 mask, without allowing enough iterations for the remaining weights to optimally adjust to this mask.

Why do you adopt exponential schedules for instead of others?

First note (Section 3.2) that having a regularizer that grows exponentially guarantees that the regularizer eventually is large enough such that the solution to the proximal operator will be exactly 2:4 sparse, which we need to ensure finite runtime. In other words: without making assumptions about the problem space, the regularizer eventually needs to be scheduled towards infinity. Furthermore, as we can empirically see from Figure 5 in the appendix our exponential schedule leads to a good balance of runtime and accuracy. We believe that at the initial steps of the algorithm it is important to not increase the regularizer too quickly, as this results in committing to a certain sparsity pattern too early. On the other hand once the sparsity pattern is mostly established we can drive the sparsification faster. This is well done by an exponential schedule.

Figure 2 shows that the proposed method is more effective when the correlation between input features is high. What does such correlation look like for LLMs? Is there any real-world task to prove your proposed method has a prominent advantage?

Please also see our reply to reviewer fgi2 ("Performance on 70B scale"). We are in fact struggling to fully nail down what exactly it is that changes between the model scales. On small models using off-diagonal information seems to be relevant, whilst on large models it does not help significantly. Whilst currently the research community (and us) are focusing on pruning open source Llama models, it is conceivable that on other future models, the off-diagonal information will again be more relevant even at larger scales.

Thank you very much for checking the correctness of all proofs for our theoretical claims and your positive assessment. We want to briefly answer two comments:

It would be better to show some real speedup for 2:4 sparsity in LLMs.

We kindly ask you to check our response to reviewer 8YCw where we provide more context on speedups -- the code for our benchmark is attached below. We believe that realizing the full speedup of 2:4 is an important effort orthogonal to our contribution which focuses on the accuracy.

Minor typo: Is there some specific reasons to use 'WandA' instead of 'Wanda'?

Yes, we thought it is more instructive to use Weights and Activations as abbreviation. However, we realized that the authors themselves used Wanda / wanda. So we will change it to their notation for consistency.

We are happy to answer any further questions that might come up.

---

---

2:4 sparsity benchmark with pytorch

import pandas as pd
import torch
import time
from torch.sparse import to_sparse_semi_structured

shapes = {
    "8b_qkv_proj": (4096 + 1024 + 1024, 4096), # assuming fused QKV
    "8b_o_proj": (4096, 4096),
    "8b_up_gate_proj": (14336 * 2, 4096), # assuming fused Up/Gate
    "8b_down_proj": (4096, 14336),

    "70b_qkv_proj": (8192 + 1024 + 1024, 8192), # assuming fused QKV
    "70b_o_proj": (8192, 8192),
    "70b_up_gate_proj": (28672 * 2, 8192), # assuming fused Up/Gate
    "70b_down_proj": (8192, 28672),
}

REPEATS = 10
WARMUPS = 10
BATCH_SIZES = [1, 16, 32, 8192, 8192*2, 8192*4]
results_df = pd.DataFrame(columns=["Layer", "Batch Size", "Speedup"])
with torch.no_grad():
    for matrix in shapes:
        for bs in BATCH_SIZES:
            out_featues, in_features = shapes[matrix] 

            # Create input and dense weight
            dense_weight = torch.Tensor([0, 0, 1, 1]).tile((out_featues, in_features//4)).half().cuda()
            x = torch.rand(in_features, bs).half().cuda()
            sparse_weight = to_sparse_semi_structured(dense_weight)

            # Function to benchmark
            def benchmark_matmul(weight, x):
                total_time = 0
                torch.cuda.synchronize()
                for _ in range(WARMUPS):
                    y = torch.mm(weight,x)

                for _ in range(REPEATS):
                    torch.cuda.synchronize()
                    start = time.time()
                    # torch.cuda.synchronize()
                    y = torch.mm(weight,x)
                    torch.cuda.synchronize()
                    end = time.time()
                    total_time += (end - start)
                return total_time / REPEATS

            # Run benchmarks
            dense_time = benchmark_matmul(dense_weight, x)
            sparse_time = benchmark_matmul(sparse_weight, x)

            # print(f"Dense time per run: {dense_time * 1000:.3f} ms")
            # print(f"Sparse time per run: {sparse_time * 1000:.3f} ms")

            print(f"Matrix: {matrix}, bs: {bs}, Speedup: {dense_time / sparse_time:.2f}x")
            results_df = pd.concat([results_df, pd.DataFrame({"Layer": matrix, "Batch Size": bs, "Speedup": dense_time / sparse_time}, index=[0])], ignore_index=True)

results_df.to_csv("sparse_matmult_results.csv", index=False)

Thank you very much for checking all our theoretical statements and the review of our paper. Below we provide a discussion on other optimization objectives as suggested in your review and some data and considerations on speedups.

Optimization objectives beyond squared loss

First we recall the motivation for using the squared loss. This builds on Hassibi & Stork (1992) "Second order derivatives for network pruning: Optimal Brain Surgeon", Section 2. If the Assumption A1: "the model is trained to convergence" holds, then by a Taylor expansion of the loss, a squared loss of the "global" hessian is the first relevant order. Note, however, that this global Hessian scales quadratically with the number of total parameters, which is completely impossible to handle for LLMs. Now we can make a second assumption A2: "suppose that the Hessian is block-diagonal, where each block contains the parameters of a single linear layer."

Combining assumptions A1 & A2 we have that in first relevant order the local squared loss is the best objective to optimize -- which is what we consider. Furthermore, a linear squared loss allows for efficient optimization with theoretical guarantees. It would be non-trivial to (efficiently) apply our proximal optimization to other objectives.

That said, neither assumptions A1 and A2 are fully met in practice.

Recent work of Dong et al (ICML'24) "Pruner-Zero: Evolving Symbolic Pruning Metric from scratch for Large Language Models" makes a systematic study of which information is relevant for pruning. However, their approach requires significantly more computational resources, as it requires end-to-end gradient information. Furthermore, they consider only a static pruning approach, i.e., no gradual weight updates during pruning. Unfortunately, their codebase does not support the llama-3 series of models and we were unable to compare directly.

Speedup

NVIDIA reports 1.1x to 1.2x speedups with 2:4 sparsity on ResNext-101 model (https://tinyurl.com/uhd2rhvb). However, we are not able to independently verify their results. Unfortunately, we were also unable to measure realized speedups using the current vLLM package with a 2:4 sparse checkpoint. We ran a Pytorch benchmark on an A100 GPU. Note that Pytorch support is in itself experimental. Notably, we observed significant slowdowns in many memory-bound workloads. However, we do see speedups up to 1.6x on some shapes which is consistent with what wanda reported (their Table 5). We anticipate that as the hardware and software stack matures, these speedups will improve.

We hope this is not considered a shortcoming of our work. Fully realizing the speedup is an important effort orthogonal to our contribution.

[Benchmark script in the response to reviewer fybz]

Thank you for the detailed review and confirming that our theoretical claims and the proposed method make sense. We understand that your reluctance comes from missing ablations on calibration samples as well as the unconvincing performance on the current 70B scale Llama models. We provide a new ablation on calibration samples below and discuss the 70B performance.

We hope that our conceptual contributions as well as the strong empirical improvements over existing SOTA on smaller models and across calibration samples can offset your concerns on the 70B model. Please let us know if there is anything else we can elaborate on.

Ablation study on the effect of calibration samples:

First let us clarify that for our experiments in the submission we were using the exact same number of calibration samples for all methods (we will clarify this at the start of Section 4.3). Additionally we now ran a new ablation on the calibration samples for the 3B model (we consider the same number of samples as in Table 17 of the wanda paper). We make three main observations:

• a/ wanda (which uses only the diagonal of the Hessian) converges quickly as known from prior work.
• b/ methods that use masked gradient updates can consistently improve their validation perplexity when using more calibration data. However, with diminishing returns (note that the number of samples is increased exponentially here).
• c/ Except when using only a single sample, our proposed prox+GD outperforms all other methods when using the same amount of calibration data.

We are thankful for this suggestion and include those results in the updated paper, as they provide valuable further insights.

---

---

Performance on 70B scale:

We fully agree with your concern that for the Llama-3.1 70B models our method despite using more resources does not improve performance. We thus would recommend to use wanda for those models. While we aimed to be transparent about this, we will make it even clearer to not misguide researchers interested in this model particularly.

Now for the question of why it does not work well on 70B models:

Notice that we can qualitatively define a hierarchy between the methods by how much information of the Hessian matrix is used. Wanda uses only the diagonal entries (see also Theorem 1), SparseGPT uses already the full Hessian, and prox makes even more heavier use of the Hessian (see our toy experiment of Figure 2). We observe that on the smaller models, SparseGPT does quite consistently outperform Wanda --> It seems that the off-diagonal elements of the Hessian are relevant. It is those cases, where doubling-down on using Hessian information (through prox or masked gradient updates) further improves performance. Now on the other hand, on the 70B scale using plain Wanda seems close to optimal -- neither SparseGPT, prox, nor masked gradient updates improve there relevantly. This seems to indicate that the behavior on this scale is largely determined by the diagonal entries of the Hessian matrix -- in which case wanda (by our Theorem 1) is optimal. This is also an interesting discussion and will add it to the paper as well.