PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models

Thanks for your valuable feedback.

Q1: We conduct extra experiments on LLaMA-2 7B under 1:4 sparsity.

We also conduct extra experiments on LLaMA-2 7B under 3:4 sparsity.

The results demonstrates that PermLLM can further improve the pruning performance under 1:4 and 3:4 sparsity.

Q2: In Table 2, we list the better results between $\text{PermLLM}_{Wanda}$ and $\text{PermLLM}_{RIA}$ due to the page limitation.

Here, we list more detailed results. And we will list the complete results in the final version.

For LLaMA-3.1 8B,

For LLaMA 7B,

For LLaMA-2 7B,

For OPT 6.7B

For LLaMA 13B,

For LLaMA-2 13B,

These results demonstrate that PermLLM is able to enhance pruning performance regardless of the one-shot pruning method applied.

Q3: Yes, we will include Pruner-Zero and MaskLLM for comparison in the final version.

For LLaMA 7B under 2:4 sparsity,

For LLaMA-2 7B under 2:4 sparsity,

As the results show, PermLLM can outperform Pruner-Zero on certain tasks even when using Wanda (which is originally much worse than Pruner-Zero) as the pruning metric. Pruner-Zero is designed to automatically search for the optimal pruning metric. Therefore, PermLLM can be integrated with Pruner-Zero, resulting in $\text{PermLLM}_{Pruner-Zero}$, which utilizes the pruning metric identified by Pruner-Zero. In this way, PermLLM can further improve the performance of Pruner-Zero.

As for MaskLLM, they use a blended training set and the training overhead is significant, which is different from the experimental settings by us and the baselines we used. We implement MaskLLM using 128 samples from C4 for training, resulting in a perplexity exceeding 5k on LLaMA-3.1 8B and Qwen-2.5 7B models. This demonstrates that the performance of MaskLLM heavily depends on sufficient training data, which is consistent with the findings presented in Figure 4 of their paper. If available, one can perform PermLLM first to find a relative good channel order, which will be beneficial for subsequent mask learning by MaskLLM.

Thank you once again for your valuable suggestions to enhance the quality of our paper!

Thank you for taking the time and effort to review our paper.

W1: We conduct additional experiments about $\text{PermLLM}_{Wanda}$ on LLaMA-2 7B to show the trade-off between performance and training cost with varying block size.

The training time is ~2h for block size=32, ~2.5h for block size=64 and ~6h for block size=128. A larger block size offers a greater optimization space. However, this increased space leads to longer exploration and convergence times. We choose block size=64 as the default configuration because it can achieve good balance between pruning performance and training efficiency.

W2: As the benchmarks we used in the evaluation are diverse, so it is extremely challenging to achieve optimal results in every benchamark. Thus, we believe the evaluation results are good enough to demonstrate the superiority of the proposed PermLLM.

Q1: We perform extra experiments about $\text{PermLLM}_{Wanda}$ on LLaMA-2 7B with different calibration data: (a) 128 samples from Pile[1], (b) 128 samples from Wikitext2 [2]and (c) 128 samples from C4 [3]. The learned permutation performs consistently well across different datasets, it indicates robustness of PermLLM.

We also use different size of calibration dataset (C4) to check the performance of PermLLM.

The performance of PermLLM improve gradually with increased size of calibration data, which also demonstrate the robustness of the proposed methods.

Thank you again for your helpful advice!

References

[1] Gao, Leo, et al. "The pile: An 800gb dataset of diverse text for language modeling." arXiv preprint arXiv:2101.00027 (2020).

[2] Merity, Stephen, et al. "Pointer sentinel mixture models." arXiv preprint arXiv:1609.07843 (2016).

[3] Raffel, Colin, et al. "Exploring the limits of transfer learning with a unified text-to-text transformer." Journal of machine learning research 21.140 (2020): 1-67.

We appreciate the time and effort the reviewer has invested in evaluating our work.

W1: While permutation training adds wall-time, it enables significant performance improvements in pruning. Unlike naive one-shot pruning, which often sacrifices performance, PermLLM achieves a better accuracy, especially on the updated LLMs, such as LLaMA-3.1 and Qwen 2.5.

W2: We conduct additional ablation study about block size and we provide a discussion about adaptive block, which can be found in Q1.

W3: The memory overhead of PermLLM is minimal. For example, if a weight matrix $\mathbf{W} \in \mathbb{R}^{n\times n}$, PermLLM only need to store the permutation order index $\mathbf{p} \in \mathbb{R}^n$, which is derived from the learned permutation matrix $\mathbf{P}$. For LLaMA-2 7B, the additional memory of the permutation will be 16KB for each decoder block if INT16 is used for index data. Thus, the overall memory overhead introduced by PermLLM is negligible. Moreover, the kernel launch time is already accounted for in the runtime evaluation.

W4: 2:4 sparsity is the primary sparsity pattern recommended by NVIDIA, as it strikes an excellent balance between efficiency and accuracy, making it the key focus of prior research on N:M sparsity. Consequently, our work also primarily focuses on the evaluation and application of this sparsity pattern. We provide additional experimental results on Qwen-2.5 7B under 4:8 sparsity. Due to time limitation, we plan to include more evaluations on 4:8 sparsity in the final version. As the results indicate, PermLLM enhances pruning performance under 4:8 sparsity, demonstrating its strong generalization capability. We also conduct experiments under 1:4 sparsity and 3:4 sparsity, the results can be found in the response of Q1 by Reviewer fovA.

W5: We just follow our baselines (i.e., SparseGPT, Wanda and RIA) for a fair comparison, which all use 128 samples from C4 dataset as the calibration data.

Q1: The results about varying block size are shown as follows:

The training time is ~2h for block size=32, ~2.5h for block size=64 and ~6h for block size=128. A larger block size provides more optimization potential but also results in longer exploration and convergence times. We set the block size to 64 as the default configuration, as it strikes a good balance between pruning performance and training efficiency.

Furthermore, we believe that using an adaptive, per-layer block size could lead to further improvements. By measuring the sensitivity or importance of each layer, larger block sizes could be assigned to the more critical layers, ensuring better preservation of the most important weights. However, as adaptive block size is not the primary focus of this paper, we leave it as a direction for future work.

Q2: As shown in Table 5, partial PermLLM only inserts learnable channel permutation modules to the most sensitive layers, and applies traditional channel permutation method to the rest layers. Partial PermLLM outperforms the traditional channel permutation method but falls short compared to full PermLLM, highlighting a performance gap between the two. Because partial PermLLM leverages learnable channel permutations for the most sensitive layers, it sacrifices a significant portion of the optimization space in model-wide optimization. However, the advantage of partial PermLLM lies in its training efficiency, making it a viable option when computational resources are limited. Moreover, it is important to note that we did not perform any retraining or fine-tuning on the dense layers in any of the experiments, meaning the weight values remained unchanged and fixed. The only component we trained was the permutation matrix, which aims to identify a permutation for generating a better pruning mask for the dense layers.

Q3: We follow the experimental settings used by our baselines (i.e., SparseGPT, Wanda and RIA) for a fair comparison. We think C4 is an out-of-domain dataset as the evaluation benchmarks are diverse and different from C4. We also conduct additional experiments using 128 samples from Pile [1] or 128 samples from Wikitext2 [2]. When using Wikitext2 as the calibration data, the pruned model performs better on Wikitext2 compared to using Pile or C4. Since we use C4 as the calibration data and apply the same dataset across all baselines, we believe the settings are fair and appropriate. And if overfitting the calibration data, we believe the results will be much worse on the evalution benchmarks than the existing results.

Q4: Yes, PermLLM remains effective when applied after 8bit GPTQ, and the results are shown as follows. Here, we perform pruning after quantization and we believe an interesting future work will be pruning- and quantization- aware permutaion learning, which can concurrently find a suitable permutation for improved pruning and quanzation effect.

Q5: Since NVIDIA GPUs are the dominant platform for research, our efforts primarily focus on developing custom kernels optimized for NVIDIA GPUs, as has been the case with most previous work. Just as unstructured sparsity does not achieve significant acceleration on GPUs, this has not stopped researchers from extensively exploring it. Similarly, N:M sparsity was initially proposed by NVIDIA, so it is reasonable that it achieves better adaptation on NVIDIA GPUs. Even if acceleration is currently limited to NVIDIA GPUs, there is still a great deal of ongoing research focused on optimizing N:M sparsity.

The permutation kernel we developed is not strictly limited to NVIDIA GPUs, as the operation itself is essentially just reordering a matrix. If hardware such as AMD MI-series or Intel Gaudi supports parallel operations, we can adapt the kernel accordingly to make it compatible with those platforms, extending the applicability of our approach.

L1: "The reported speed-up (≈1.67×) is far below the theoretical 2× for 2:4 sparsity, indicating residual inefficiencies." $2\times$ speedup is theoretical, and it does not mean everyone can achieve under any cases. SparseGPT and RIA also report speedup under 2:4 sparsity, which can not achieve the theoretical speedup either. Thus, this is not our limitation; it is a limitation of the hardware itself, as well as the gap between theoretical analysis and practical application. The only limitation about runtime is the additional channel permutation operations we introduce. With the developed CUDA kernel, the overhead is minimal.

L2: Our methodology is not restricted to decoder-only transformers. As demonstrated in the paper, it does not rely on any specific components exclusive to LLM. We will open source our framework so that other researchers or users are able to extend it to any architectures. Furthermore, due to the structural similarities across model architectures used by encoder-decoder models or vision transformers, it can be easily adapted for them.

Thanks again for your valuable suggestions to improve the quality of our paper!

References

[1] Gao, Leo, et al. "The pile: An 800gb dataset of diverse text for language modeling." arXiv preprint arXiv:2101.00027 (2020).

[2] Merity, Stephen, et al. "Pointer sentinel mixture models." arXiv preprint arXiv:1609.07843 (2016).

Thanks for your detailed review and valuable feedback.

W1: Thanks for your reminding. We will include the missing references and discuss them in the final version.

(a) Mahajan et al. [1] and [42] propose different methods to achieve a good permutation result for N:M sparsity. However, their major limitations lie in the biased optimization objective (focusing on maximum retained weight importance rather than the output loss between the original model and the pruned model) and the local optimality of the permutation matrix (performing layer-wise optimization, which fails to capture inter-layer dependencies). These issues, as highlighted in the motivation section, limit their ability to further improve performance. To further validate this, we implement the method from [42] combined with Wanda for a direct comparison against the proposed PermLLM. Please refer to Q1 for detailed results. "Prohibitively high computational complexity" in line 46 refers to the solution space complexity of concurrently finding permutations for all layers. For example, if we need to determine two permutations for two linear layers in the model, and there are n possible permutations for layer 1 and m for layer 2, the complexity is n+m when performing layer-wise optimization. However, for model-wise optimization, the complexity increases to n*m. We will rewrite it to avoid misunderstanding.

(b) Lyu et al. [2] employ an $\ell_{1-2}$ penalty to constrain the learnable matrix, ensuring that, upon convergence, each row and column contains exactly one entry equal to 1, with all other entries being zero. They also utilize Sinkhorn normalization to facilitate convergence. The key difference between their approach and ours lies in the method used to achieve convergence to 1 or 0, which we think it is minor: while they rely on the $\ell_{1-2}$ penalty, we adopt an annealing-based method. We believe that two approaches can achieve similar results with appropriate training configurations. Furthermore, the more valuable contributions of PermLLM are its abilities to achieve pruning-aware permutation learning for the first time and efficiently train the permutation matrix, which provides a new pruning framework for further exploration.

W2: RotPruner [5] prunes the model in rotated space. As shown in Fig.2 of their paper, the rotation operation results in more values being 0 or small magnitude. We think it may be more benefical for unstructured pruning as N:M sparisty has more strict constraints.

There are some shortcomings compared with PermLLM: (a) memory overhead: RotPrune requires storing the rotation matrices in memory. For a weight matrix $\mathbf{W} \in \mathbb{R}^{n \times n}$, the corresponding rotation matrix $\mathbf{Q} \in \mathbb{R}^{n \times n}$ has the same size, leading to significant memory usage. In contrast, PermLLM only needs to store a vector $\mathbf{v} \in \mathbb{R}^{n}$ to represent the column permutation order, resulting in much lower memory overhead. (b) runtime overhead: RotPrune requires online computation of rotations, which is less efficient compared to the lightweight permutation operation. Rotation can be seen as an additional dense linear layer, introducing extra computational costs (e.g., Q/K/V_proj vs. CP) in Table 3.

PermLLM can also outperform RotPruner, the results can be found in Q2.

W3: We will revise the paragraph in the final version for improved clarity. The original PyTorch implementation is inefficient, so we develop a custom CUDA kernel to achieve significant speedup. With this optimized kernel, the additional permutation incurs minimal overhead, resulting in a negligible overall impact. Thank you for bringing this to our attention!

W4: We conduct additional experiments on LLaMA-2 7B by varying the number of iterations used in Sinkhorn normalization. With the exception of Wikitext2, iter=5 outperforms both iter=0 and iter=10 on the remaining benchmarks. The reason why iter=5 performs better is that introducing some perturbations during the early exploration helps in finding a better solution and avoiding being trapped in a local optimum.

W5: Block size is a critical hyperparameter that balances pruning performance and training efficiency. To explore this trade-off, we conduct additional experiments by varying the block size, which can be found in Q4.

W6: We used SparseSemiStructuredTensor and to_sparse_semi_structured in Pytorch and enabled CUTLASS for sparse operations. Moreover, we utilized Timer utility from torch.utils.benchmark to measure runtimes.

Q1: We apply [42] with Wanda to prune LLaMA-2 7B. The results below demonstrate the superiority of PermLLM. There is no doubt that our motivation is well-founded and our methodology highly effective. The reasons behind PermLLM's superior performance have been discussed in W1.

Q2: In RotPruner, 128 samples from Wikitext2 are used as calibration data. Thus, we also use 128 samples in Wikitext2 as calibration data for PermLLM. The results are shown below:

We didn't find the results for the zero-shot tasks under 2:4 sparsity in their paper. However, for Wikitext-2, PermLLM significantly outperforms RotPruner. Moreover, we believe there is room for both techniques, making it an interesting avenue for further exploration and improvement.

For example, $\mathbf{x}$ = [[-1.0412],[-1.4108],[ 1.0199],[-1.7633],[-0.2013],[ 0.8784],[ 0.2520],[ 1.3273]],

$\mathbf{W}$=[[ 0.5143, 0.5205, 0.6937, 0.0350, -0.4507, 0.3604, 2.4873, 0.9460],[ 0.0755, -0.9656, 1.2620, -0.7471, -2.5672, 1.6677, -1.1052, 0.6418]],

rotation matrix $\mathbf{Q}$=[[-0.1583, -0.5283, 0.4140, -0.0586, 0.2760, -0.0372, 0.1789, -0.6414], [ 0.6179, -0.2853, -0.1028, -0.5905, -0.0206, 0.2548, -0.3318, -0.0461], [-0.2811, -0.3364, -0.3383, -0.2190, 0.4878, 0.2649, 0.3716, 0.4463], [ 0.1433, 0.3674, -0.1850, -0.4794, -0.0594, -0.3628, 0.6207, -0.2449], [ 0.0564, -0.2713, -0.7957, 0.3407, -0.0890, -0.1264, -0.0599, -0.3829], [-0.2781, -0.3765, 0.0066, -0.2865, -0.2223, -0.7095, -0.2565, 0.2831], [-0.3605, -0.1679, -0.0178, -0.1824, -0.7537, 0.4435, 0.1991, -0.0619], [ 0.5321, -0.3844, 0.1894, 0.3723, -0.2377, -0.1262, 0.4749, 0.3110]],

permutation order=[3, 1, 5, 4, 6, 0, 2, 7]. With magnitude pruning, direct 2:4 pruning loss=1.660, rotated 2:4 pruning loss=0.769, rotated and permutation 2:4 pruning loss=0.497. This example highlights the potential of combining rotation and permutation for more effective N:M pruning.

Q3: Based on our analysis in W1, the differences between [2] and permutation learning in PermLLM are minimal. Consequently, with proper usage, the final performance is expected to be similar. We believe this aspect is not particularly critical for the evaluation of PermLLM, as, similar to mask learning in prior pruning works, there are various methods to learn the pruning mask, such as SR-STE [i], decaying-based methods [ii], and others. The more valuable contributions lie in implementing permutation learning to successfully enhance pruning performance and make this approach more practical—achievements that PermLLM fully realizes. At the same time, we hope that our motivation and approach can inspire more researchers to further advance this field. For example, it could be applied to quantization, combined with rotation as you mentioned, or lead to the development of better pruning-aware permutation learning framework.

Q4: We conduct an ablation study about block size to see its impact on final pruning performance and training cost.

The training time is ~2h for block size=32, ~2.5h for block size=64 and ~6h for block size=128. A larger block size offers a greater optimization space. However, this increased space leads to longer exploration and convergence times.

Thank you again for your constructive suggestions!

References

[i] Zhou et al., LEARNING N:M FINE-GRAINED STRUCTURED SPARSE NEURAL NETWORKS FROM SCRATCH, ICLR 2021.

[ii] Kao et al., Training Recipe for N:M Structured Sparsity with Decaying Pruning Mask, ICML2021 workshop.