Low-Rank Quantization-Aware Training for LLMs

We report WikiText-2 perplexity (lower is better) and Zero-shot accuracy over 6 tasks (higher is better). Note that as listed in Table C8 in the appendix, we found the best range estimators to be L^3.5, min-max, L^4, L^3.5 for w4pc, w4g128, w3pc, w3g128, respectively. We will also include this ablation study in the paper.

Below you can find OmniQuant results for LLaMA-1 and LLaMA-2 model families, which we obtained using their public open-sourced code base (https://github.com/OpenGVLab/OmniQuant).

We used the provided checkpoints by the authors and ensured that the evaluation uses exactly the same lm_eval version (0.4.2). With this, we matched the WikiText-2 perplexity numbers reported by OmniQuant closely as well as the FP16 zero-shot baseline reported in our paper.

We report zero-shot accuracy over 6 tasks (higher is better). *Uses asymmetric quantization.

As we can see, LR-QAT is consistently outperforms OmniQuant across all models and settings (except W4g128 in which they are on par) and the gap becomes progressively bigger as we approach more extreme 2-bit quantization.

We will include these results in Table 4 in the updated revision of the paper. In addition to that, we will also generate and include the remaining missing OmniQuant results for LLaMA-3 & Mistral, and for weight-activation quantization and include them in Tables 4 & 5, respectively.

Q1

$s_0$ is the frozen value of the scale at initialization (i.e., right after the range estimation).

Q2

Please see response for W4 regarding the relationship between fixed-point and integer representation.

Q3

To make sure we understand your question correctly, do you mean that during the backward pass we would pretend that $s \cdot \text{clip}(\text{round}(\Phi_0 + AB),\ldots)$ is $s \cdot (\Phi_0+AB)$ , i.e. ignore both clipping and rounding?

If so,

• in terms of memory, no, as we are already making sure that we are not storing intermediate results of quantization using gradient checkpointing (see Figure 1).
• in terms of computation, there might be some minor savings when computing gradients w.r.t $A$, $B$ since now we will be ignoring not only rounding (by using STE assumption, L218) but also clipping. On the other hand, this could potentially lead to less accurate gradients depending on the range of $\Phi_0+AB$.

Q4

Please see response to Q3.

Q5

Please note that the size of the fine-tuning dataset itself is not the issue, as we are using a general pre-training dataset, but rather the number of iterations we used for this study.

As we focus on efficient QAT techniques, we did not explore significantly longer fine-tuning. Indeed, if we would consider significantly longer fine-tuning, then applying scaling laws would be a very interesting direction.

Q6

Please see the response to W1.

We thank the reviewer for their thoughtful and useful feedback. Please, find our replies to specific concerns below.

W1: Limited Comparison in Initialization

We agree that RTN is not necessarily an optimal quantization scheme, though it is the most common initialization choice in QAT literature.

Please note that our RTN baseline is significantly stronger than the ones from related work (e.g., OmniQuant, LLM-QAT), because we set the ranges based on minimizing the $L^p$-norms between quantized and unquantized weights and use the best performing $L^p$-norm based on a holdout Wikipedia validation set (see the details in L819-823 and Table C8).

We did perform an additional ablation study comparing different range initializations with RTN for LR-QAT and PEQA applied to LLaMA-2 7B (best, as reported in the paper vs. using min-max initialization), which we included in the subsequent comment.

Note that it is rather uncommon to use some of the learnable rounding techniques (such as GPTQ, AdaRound) as initialization for the end-to-end QAT setting. At the same time, we see the value in exploring the combination of our method with techniques like OmniQuant as a preprocessing step. We will further extend our ablation study to include OmniQuant for the next revision.

W2: Lack of Complete Results in Tables

We agree that having more than a few missing results in the tables makes it more difficult to compare results across methods and limits the insight into performance in certain configurations.

Please note that we put a significant effort in reimplementing and reporting results for LSQ and PEQA, which we deem the most relevant baselines given that they are also QAT-based techniques. We also made sure to account for initialization, number of training examples, hyperparameter tuning etc., to ensure fairness in this comparison. In addition to that, we report results across 5 models, 6 distinct weight-only and 3 weight and activation quantization settings and compare against numerous baselines, which is substantially more than in many related works.

Note that a good fraction of missing results corresponds to LLaMA-3 and Mistral model families, for which there were no reported results in most related work. In addition to that, unfortunately, some prior works were not consistent in reporting results across all model families and settings (e.g., OmniQuant reports WikiText-2 perplexity but not zero-shot accuracy for W-only results and vice versa for W4A4 setting).

While we agree that authors should put in a reasonable effort to compare fairly to related work, which may include to reimplement and/or run some new studies with existing methods, we do not think that it is reasonable to expect from authors to always reimplement or rerun all relevant baselines for all settings. We do not think that it is reasonable to expect from authors of new papers to always reimplement and/or run all the missing baselines.

Having said that, we agree to generate and include the missing results from OmniQuant, as it is the strongest PTQ baseline that consistently outperforms other PTQ techniques (GPTQ, AWQ, RTN). Please see the below comment for additional zero-shot accuracy results for OmniQuant for LLaMA-1 and LLaMA-2 model families.

W3: Potential for Limited Practical Impact

Please see general response to all reviewers. In short, even with just ~4h runtime (1k steps) we outperform OmniQuant, which for many practical cases we believe is valuable trade-off.

W4: Fixed Point Representation Modeling

This is an interesting observation and suggestion.

Indeed, this is exactly how fixed-point numbers, such as Q4.4, are defined and how we implemented them: the set of numbers representable by Q4.4 is the same as the set of numbers representable by INT8, divided by $2^4$ (or more generally, divided by $2^{8-b}$, where $b$ is the target bitwidth).

Hence $\Phi_0$ can be written as $2^{-(8-b)}\cdot W_{\text{INT8}}$, where $W_{\text{INT8}}$ is the frozen INT8 matrix. Note, $W_{\text{INT8}}$ needs to be frozen in order to get the memory savings, thus the scale used to obtain $W_{\text{INT8}}$ cannot be tuned.

However, by tuning the scaling factor $\tau=2^{-(8-b)}$ but still keeping the $W_{\text{INT8}}$ frozen, we could in principle achieve slightly better accuracy, however, the term $\tau \cdot W_{\text{INT8}}$ would lose its interpretation of being an approximation to the original weights with the fractional part being truncated to $8-b$ bits.

typos and minor issues

All fixed in the updated version, thanks!

We thank the reviewer for their thoughtful and useful feedback. Please, find our replies to specific concerns below.

W1: novelty

Please, find our reply to your concerns regarding the novelty in the general response to all reviewers.

W2: finetuning on different datasets

Please note that one of the core ideas of the proposed approach is that it’s positioned as a general extended pretraining method, as opposed to being strictly a fine-tuning on the downstream task method. The resulting model is a low-bit general pretrained LLM backbone, that can still be utilized for any task afterwards or further combined with LoRA-based techniques or other approaches.

At the same time, we agree that demonstrating the performance by fine-tuning on different datasets will further strengthen the experimental section and enable comparisons to more related work, including the aforementioned QLoRA and QA-LoRA.

We will include results of LR-QAT fine-tuned on Alpaca / FLANv2 datasets in the updated version of the paper.

W3: typos

All fixed in the updated version, thanks!

W4: extended pre-training vs fine-tuning

We agree that the extended pre-training vs task-specific fine-tuning deserves more discussion. We chose to ‘position the paper’ as extended pre-training for our motivation and in order to compare fairly to a wide variety of PTQ and QAT literature.

However, as the reviewer correctly points out, our proposed technical solution (LR-QAT) is not limited to this setup which we deem an additional benefit of it. As mentioned above, we will extend the experimental section to include LR-QAT trained on Alpaca / FLAN-v2 to allow comparison to PEFT literature, including QLoRA and QA-LoRA.

W5: comparison with QA-LoRA

There are indeed some similarities between QA-LoRA and our work, namely:

• Both use low-rank (re)parameterization
• Both target memory- and inference-efficient training

However, there are some significant differences:

• Key difference, restrictions on quantization granularity: QA-LoRA assumes all rows of A must be the same => weight quantization has to be done with group-wise granularity with very small group-size (e.g., 32) otherwise it’s very inflexible, whereas our method works for any quantization granularity out of the box. In particular, QA-LoRA cannot be applied to pre-channel weight quantization, which is a common case for weights.
• Because of the aforementioned restrictions, QA-LoRA also reports results only using asymmetric weight quantization, which of course provides extra degree of flexibility. Note that not all hardware can support asymmetric weight quantization, and even if it does, it introduces extra inference overhead compared to symmetric weight quantization [1]. At the same time, we demonstrate state-of-the-art results for W4 and W3 uniform affine quantization using symmetric scheme.
• The above two points further lead to increased memory usage and compute overhead during inference.
• Finally, note that QA-LoRA is introduced as a downstream task fine-tuning method (like QLoRA). As a consequence of that, this means they do not compare extensively to other PTQ and QAT techniques.

References

[1] Nagel et al., "A white paper on neural network quantization", ArXiV:2106.08295

Dear Reviewer, We thank you for carefully considering our rebuttal and for your valuable feedback. Please, see further clarification below.

Extended pre-training: Of course, we fully agree that from a technical perspective, the difference between extended pre-training and fine-tuning is in practice just switching the dataset (and there is no novelty in that).

We would like to clarify the framing of our method as ‘extended pre-training’ technique. As mentioned, our main motivation is to position our method as a general QAT technique and to compare fairly to a wide variety of PTQ and QAT literature. In contrast to PEFT literature, where it’s prevalent to combine the method together with the downstream task fine-tuning, most PTQ and QAT literature does not do any task fine-tuning (all PTQ literature but also LLM-QAT, LSQ etc). Therefore, we would like to very clearly separate and isolate improvements coming from the novelty of the method from the “pipeline improvements” (which includes a use of different, potentially more advantageous dataset, as well as other things like extra loss terms/knowledge distillation etc.).

As mentioned in our earlier response, we will extend the experimental section to include LR-QAT trained on different datasets (such as Alpaca / FLAN-v2) to allow comparison to PEFT literature, including QLoRA and QA-LoRA. Following your suggestions and this discussion, we will carefully review and adjust the phrasing in the paper to better reflect our intention and the relationship between fine-tuning and extended pre-training.

Novelty: To clarify, the ‘extended pre-training' is not novelty and rather a framing of our work. Though there are clear technical novelties about our approach, including introducing a low-rank reparameterization for QAT that does not lead to any inference time overhead, and the use of a fixed-point downcasting operator, and the use of checkpointing to further decrease memory during training. As mentioned in the general response, we agree that our method is relatively straightforward from an implementation standpoint. However, it is a novel QAT technique that is very effective, easy to use, and is solving the important problem of memory- and inference-efficient LLM quantization. Despite its simplicity, we think it is a novel and valuable contribution to the field.

OmniQuant dataset: Indeed, our OmniQuant numbers are directly taken from their paper as is a common practice in most literature. We do agree that there might be a small difference between fine-tuning on Wikitext-2 vs fine-tuning on SlimPajama.

Since we introduce a new QAT technique, we deem PEQA and full-model QAT (LSQ) as the most relevant baselines. Please note that to ensure fairness, we carefully followed the same experimental setup (including dataset, but also initialization, number of training examples, hyperparameter tuning etc.) for both PEQA and full-model QAT (LSQ). As we included PTQ techniques mainly as additional references, we did not follow the same practice there and used the authors provided results for comparisons.

Given that OmniQuant is our strongest PTQ baseline, we agree that it is valuable to include extra results using the same calibration dataset – SlimPajama, to ensure fairness. Below (in the next comment) you can find OmniQuant results for LLaMA-1/2 model families using SlimPajama (SP) as a calibration set, which we obtained using their public open-sourced code base (https://github.com/OpenGVLab/OmniQuant). We kept the same hyperparameters (number of iterations etc.) as for the original OmniQuant results.

As we can see, WikiText-2 (WT2) perplexity is consistently worse when using SlimPajama as a calibration set (as expected) while zero-shot accuracy stays comparable. We will update the paper with these results and include this detailed comparison in the appendix in the updated version of the paper.

We greatly appreciate your time and consideration and hope these updates address your concerns.

Best, authors

We thank the reviewer for their thoughtful and useful feedback. Please, find our replies to specific concerns below.

W1: referencing Figure 1

That is a very good catch! Indeed, we did not reference Figure 1, and we will make sure to include it in the updated version of the paper.

We agree that a more detailed discussion of the memory breakdown (Figure 1, right) is valuable to the paper as one of our primary goals is to enable memory efficient QAT. We will extend section 5.3 (comparison to QAT) with an additional paragraph discussing this.

W2: a table listing quantization settings

Note that we follow closely prior literature, specifically OmniQuant and LLM-QAT, in our quantization settings and assumptions. In case the two differed (only one activation quantizer in self-attention), we used a superset of both, thus our numbers are comparable in a fair manner to both works.

We do agree that including a comparison table that lists what parts of the model are being quantized together with the rest of quantization settings and assumptions across reported methods will be valuable and will provide a more complete picture. We will include this in the appendix in the updated version of the paper.

W3: usage of training/pretraining/fine-tuning

Thanks for pointing out this inconsistency, we agree that being consistent with the usage of training, pretraining and fine-tuning will enhance the clarity of our paper. We will carefully review and correct those usages across the paper to make it more consistent.

W4,7-9: typos

All fixed in the updated version, thanks!

W5: closest work

Both PEQA and QA-LoRA are close to our work in the sense that both are PEFT methods that aim at reducing the memory footprint during training while also maintaining inference efficiency and improved predictive performance compared to PTQ methods.

Note, however, that QA-LoRA constraints matrix A to have the same rows, meaning that it is only designed to be applicable in the case of block-wise weight quantization with a very small group size (e.g., 32). Because of that, it cannot be applied to neither pre-channel nor groupwise weight quantization with bigger groups sizes (64-512), which are common settings for weights.

At the same time, neither PEQA nor our proposed methods have any constraints on the quantization granularity and can equally well be applied to per-channel and even per-tensor weight quantization regimes.

Hence, we consider PEQA to be the closest work and that’s why we decided to reimplement and compare against this baseline. We will update the related work section to make this clearer.

W6: PEQA trade-offs

Just like LR-QAT, PEQA is also memory- and inference efficient method. However, PEQA has significantly fewer degrees of freedom compared to our method, leading to a subpar predictive performance. At the same time, the accuracy is still consistently better than for PTQ methods. Therefore, we see PEQA in Table 1 as follows:

Thank you authors for taking into account the suggested changes and providing clarifications to other raised issues. I look forward to seeing the revisions as part of the final manuscript. I am raising my score from 6 to 8.

We thank the reviewer for their thoughtful and useful feedback. Please, find our replies to your concerns regarding the novelty and the runtime comparison in the general response to all reviewers.