Progressive Mixed-Precision Decoding for Efficient LLM Inference

Dear Reviewer FUCY,

We thank you for your time and your detailed feedback. Please check our reponses below.

Q1.1: replacement of repetition of concepts with details of scheduler model/training design.

Thank you for the feedback. We will revise the paper to make it concise and include scheduler model/training details in the main paper.

Q1.2: high level stats/evals on how different tasks/models perform when changing the precision in the decode/prefill phase.

We provide experiments with high-precision decoding and low-precision prefilling for Vicuna-7B on MT-Bench.Note that this analysis is also included in our response to Reviewer RDtY

In the paper, we also presented relevant quantitative results of more models and tasks in Figure 3, Table 1, and FIgure 6. We will add include more discussions in the revised paper.

Q2: LLM benchmarks around Math, reasoning and especially long context.

We added gsm8k for Math/Reasoning and LongBench for long context evaluation. Note that for other models we tested, given their small size and limited context window length adopted by edge-friendly models, they did not perform well on Math and long-context tasks even in full precision.

Q3: more statistics on evaluation like average decode length

We provide the average decode length for three tasks and three models here. The stat is in format {precision: average decode length}. For example, {'3': 39.0, '2': 65.4} means 3-bit model is used to generate 39 tokens and 2-bit model is used to generate 65.4 tokens on the average.

CNN/DM:

Dialoguesum:

IWSLT:

Q4: effect of different N and suggestion for the choice of N.

N is a user-defined hyperparameter. As discussed in Section 4.3, larger N leads to larger search space, while leading to potentially larger bitwidth reduction or higher performance. The optimal N depends on both task and model. We empirically found that N=4 is a reasonable choice that balances the performance and search complexity.

We conducted an ablation study to evaluate the impact of N on the maximum bitwidth reduction while maintaining a performance drop of less than 2% as compared to high-precision model. The results show that N=4 achieves a good balance, enabling significant bitwidth reduction while meeting performance requirements for most tasks. Note that for Dialoguesum, solely using the low-precision model during decoding achieves lossless performance so the maximum bitwidth reduction is already achieved at N=2. For IWSLT, when N <= 3, a precision switch step that allows for performance drop within 2% couldn’t be found, so the high-precision model is used throughout, leading to no drop in both bitwidth and performance.

Q5: Inconsistency of bitwidths specified in the table and caption for phi-1.5.

Thank you for pointing this out. This is indeed a typo. It should be 3/4 bit in both the table and caption. We will correct it in the revised paper.

Dear Reviewer TCS8,

We thank you for your time and your detailed feedback. Please check our reponses below.

Q1: Comparison between PMPD-Learned and PMPD-Static in Table 1.

Thanks for this insightful question. The key distinction is that PMPD-Static requires in-domain validation data, while PMPD-Learned operates in a zero-shot manner. The performance gap between these approaches can be attributed to two factors: Domain Shifts: While we used the C4 dataset to train the PMPD-Learned scheduler for better cross-task generalization, significant distribution shifts were observed between this seed dataset and the downstream tasks we evaluated on. Consistent Performance Gap Across Samples for a Single Task: Within individual tasks, we observed that the performance gaps between models of different precisions remained consistent across samples. This consistency suggests that a static schedule per task can effectively approximate the optimal precision scheduling policy, making PMPD-Static particularly effective.

Q2: Including speedup with respect to 3 and 4-bit baselines to the Abstract.

Thank you for the suggestion. We will include the number in the Abstract in the revised paper.

Q3: Is PMPD lossless? How does it compare to unquantized fp16 models?

PMPD introduces a precision allocation scheme that is independent of the underlying post-training quantization (PTQ) method. Whether it is lossless depends on the PTQ method it uses rather than its precision allocation scheme. In the table below, we provide further comparisons to the fp16 models, demonstrating that PMPD generally preserve the fp16-level accuracy in most scenarios.

Regarding MT-Bench results (Table 2): as we used the fp16 model outputs as reference answers, the Rouge-L and BertScore metrics for the fp16 model could become unreasonably high, making it unsuitable as a baseline for this specific evaluation.

We will include more fp16 comparisons in the revised paper.

Q4: Comparisons to more baselines.

Please refer to our answer to Reviewer A5Dp for more discussions about the comparisons to other baselines, e.g. GPTQ.

To summarize, PMPD introduces a precision allocation scheme that is independent of the underlying post-training quantization (PTQ) method. In our experiments, we primarily used SqueezeLLM as our PTQ approach. Direct comparisons with other PTQ methods (such as GPTQ [4] and QuIP [7]) would primarily reflect differences in the underlying quantization approaches—a comparison already reported in the SqueezeLLM paper [2], which demonstrated SqueezeLLM's superior performance compared to GPTQ [4] and QuIP [7].

We also tried to include comparisons to other baselines [6] suggested by the reviewer, but did not manage to find any public implementation.

We will include more discussions about the baselines in the revised paper.

[2] SqueezeLLM: Dense-and-Sparse Quantization. ICML 2024.

[4] GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers

[6] Fast and Efficient 2-bit LLM Inference on GPU: 2/4/16-bit in a Weight Matrix with Asynchronous Dequantization.

[7] QuIP: 2-Bit Quantization of Large Language Models With Guarantees.

Q5: Quantization method for Baseline-l and Baseline-h

We used SqueezeLLM with Dense-And-Sparse decomposition (DNS) as the baseline. Since PMPD also utilized SqueezeLLM (without DNS), this ensures a fair comparison. In this setting, our bitwidth reduction is solely attributed to PMPD’s precision allocation policy, which effectively removes redundancy during the decoding process.

Q6: Precision range used for PMPD shown on Table 1.

For Vicuna-7B, we used 3bit and 2bit. For all other models (MobileLlama, Phi-1.5, Zephyr-3B), we used 3bit and 4bit. This is because larger models like Vicuna-7B have lower quantization error and its 3bit version is almost lossless as suggested in SqueezeLLM[2].

Q7: Memory overhead of precision switching.

Please refer to our answer to Reviewer A5Dp for more discussions about the implementation and efficiency of precision switching.

In summary, as described in Sec. 2.1, PMPD leverages AnyPrecision LLM [3], which enables efficient storage of M different precision models (1 to M-bit) using only M bits, while preserving model accuracy. This approach ensures PMPD requires no additional memory overhead. Furthermore, precision switching incurs negligible latency since no extra weight loading is required.

We adopted AnyPrecision LLM to store models of different precisions, as detailed in their paper. AnyPrecision LLM enables storing up to n models of different precisions using only n bits, without compromising the performance of any model. Hence, our method does not incur additional memory overhead. For instance, when using both 3-bit and 2-bit models, the memory usage remains equivalent to that of a single 3-bit model.

Q8: PMPD-Static with the same bitdiwth outperforms Baseline-l with 2bits

For that specific experiment in Table 1, PMPD-Static uses 3bit for prefilling and 2bit only for decoding. Hence, the average bitwidth for decoding is 2. As shown in Figure 6, using a higher precision for prefilling increases the performance without incurring noticeable latency overhead.

[2] SqueezeLLM: Dense-and-Sparse Quantization. ICML 2024.

[3] Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs. ICML 2024

Dear Reviewer A5Dp,

We thank you for your time and your detailed feedback. Please check our reponses below.

Q1: Implementation and cost of precision switching.

As described in Sec. 2.1, PMPD leverages AnyPrecision LLM [3], which enables efficient storage of M different precision models (1 to M-bit) using only M bits, while preserving model accuracy. This approach ensures PMPD requires no additional memory overhead. Furthermore, precision switching incurs negligible latency since no extra weight loading is required.

To verify this, we measured the decoding latency of Vicuna-7B under two scenarios:

1. Uniform 2-bit precision for 100 decoding steps
2. Mixed precision: 3-bit for 1 step followed by 2-bit for 99 steps

The results, shown in the table below, confirm that the overhead from precision switching is negligible.

We thank the reviewer for coming up with experiments to strengthen our work. We will further clarify the implementation details and cost of the precision switching in the revised paper.

Q2: Result on GQA models.

As suggested by the reviewer, we performed further experiments on Llama3 that leverages GQA attentions, using static schedulers. The results confirm PMPD's ability to effectively reduce bitwidth while preserving model performance, validating its compatibility with different attention mechanisms.

Q3: Comparisons with more baselines

PMPD introduces a precision allocation scheme that is independent of the underlying post-training quantization (PTQ) method. While we primarily utilize SqueezeLLM as our PTQ method, PMPD is compatible with other approaches like GPTQ.

Our choice of SqueezeLLM with uniform bitwidth as the baseline is methodologically sound, as it isolates the impact of our precision allocation strategy. Direct comparisons with other PTQ methods (such as GPTQ and AWQ) would primarily reflect differences in the underlying quantization approaches—a comparison already reported in the SqueezeLLM paper [2], which demonstrated SqueezeLLM's superior performance compared to GPTQ [4] and AWQ [5].

To showcase the generalizability of our method, we have already included results with GPTQ in Appendix A.4, Table 6 of our original submission.

To demonstrate PMPD's versatility, we further include a comparison below showing that GPTQ underperforms both SqueezeLLM and PMPD.

For efficiency, since the underlying PTQ method we used in Experiment is SqueezeLLM, here is the comparison with GPTQ based on results from SqueezeLLM[2].

[2] SqueezeLLM: Dense-and-Sparse Quantization. ICML 2024.

[3] Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs. ICML 2024

[4] GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers

[5] AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

Dear Reviewer RDtY,

We thank you for your time and your detailed feedback. Please check our reponses below.

Q1: quantitative analysis about the optimal bitwidth allocation between prefilling and decoding stages.

We performed experiments comparing different bitwidth configurations for prefilling and decoding stages using MT-Bench. The fp16 model outputs were used as the reference answer to compute the Rouge-L score:

These findings complement our quantitative analysis presented in Fig. 6 (Section 5.3), which we will further elaborate in the revised paper.

The results demonstrate that, given the same total bitwidth budget, allocating higher bitwidth to prefilling (3-bit prefill, 2-bit decode) significantly outperforms the alternative configuration (2-bit prefill, 3-bit decode). This validates our proposed approach. Additionally, as discussed in our paper, our high-precision-prefilling strategy is also more hardware-efficient since prefilling operations are compute-bound [1].

[1] LLM Inference Unveiled: Survey and Roofline Model Insights. Yuan et al. 2024

Q2: Bitwidth search for prefilling did not consider that the decoding precision was decreasing.

Our research demonstrated that higher precision during prefilling enhanced both performance and hardware efficiency. Given this observation, there may be different approaches to prefilling bitwidth search. In our paper, we chose a sequential optimization solution and used the prefill stage bitwidth as the initial point for decoding as it is simple yet empirically effective. However, we acknowledge that more advanced methods could be explored. The reviewer suggests an expanded search space that could incorporate non-uniform bitwidths in the decoding stage, an interesting direction for future investigation. We appreciate this insightful suggestion and will explore it in the future research.

Q3: The Task-Agnostic Learned Scheduler did not make use of the generated tokens.

We agree with the reviewer that ideally we should also consider the generated tokens for optimal switch point search. We currently only used the input prompt tokens for two practical reasons: i) predicting a precision at every decoding step introduces extra latency, which may offset the benefits of model quantization. In contrast, our input-based approach maintains constant scheduler overhead. ii) our empirical results showed that input prompt information alone yielded effective results.

That being said, we acknowledge that incorporating generated content into scheduler decisions is a promising future direction that we will explore in the future research.

Q4.1: Results for models larger than 7B.

As mentioned in our Introduction (Sec. 1), PMPD mainly targets edge and mobile applications. We agree with the reviewer that it would be interesting to see if larger models could also benefit from PMPD. We conducted additional experiments using Llama3-8B and longchat-16k-13B models, with a static scheduler for PMPD given the availability of validation data. These new results demonstrate our method's generalizability across model scales.

Q4.2: End-to-end GPU speedup data.

As PMPD built on SqueezeLLM [2] and Any-precision LLM [3], we evaluated PMPD's GPU performance primarily on linear layers, following common practices in the literature [2][3]. To measure end-to-end latency improvements compared to fp16 models, we tested using single-batch prompts with a context length of 1024 and recorded Self-CUDA-time as in our paper. Our results demonstrate that PMPD consistently outperforms uniform quantization baselines across various models and GPU architectures.

While the overall speedup is more modest than when evaluating linear layers alone, it is expected since both PMPD and the baselines share common fp16 operations, such as self-attention, which are often time-consuming. Notably, NPUs present a more promising platform than GPUs for maximizing PMPD's performance benefits on edge devices, as they better support custom algorithmic operations.

[2] SqueezeLLM: Dense-and-Sparse Quantization. ICML 2024.

[3] Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs. ICML 2024

Q5: The vague usage of maximum context length of LLMs (CL)

We thank the reviewer for pointing out the vague use of the term. We will revise the paper as suggested.