RoSTE: An Efficient Quantization-Aware Supervised Fine-Tuning Approach for Large Language Models

Table B: Training Time and Training Memory Consumption. Our server of 8 $\times$ A100 has a total GPU memory of 320GB.

Table C.1: Extended Experiment Results on Qwen 2.5 0.5B model.

Table C.2: Extended Experiment Results on Qwen 2.5 7B model.

RoSTE combines activation rotation with SFT but does not include comparisons or discussions with the following rotation-based quantization methods ...

We included DuQuant in the new experiments on Qwen 2.5 models in the above Table C.1, C.2. Despite our efforts spent on hyperparameter tuning (including learnable weight clipping, activation clipping, epoch, block size), we found that DuQuant suffers from a huge performance degradation in W4A4KV4 and W4A8KV4 setup.

I observed that for Llama3.1 8B, on certain evaluation benchmarks (e.g., TruthfulQA), RoSTE performs worse than QuaRot. I recommend providing analysis and insights into the reasons behind this discrepancy.

We suspect that this discrepancy is due to that the two approaches use different objectives: QuaRot is not adapted to the SFT loss; RoSTE directly performs training on the SFT loss and SFT dataset. It is expected that QuaRot's performance distribution will shift further away from the full-precision SFT model. We also remark that none of the quantized model can outperform the full-prec models in the TruthfulQA benchmark.

Could you provide more analysis on the training cost of RoSTE compared to traditional PTQ methods?

Details on the training cost can be found in Table B. Moreover, we illustrate in figure on the training cost-accuracy trade-offs of RoSTE and other SOTA methods.

Why do the PTQ baselines perform poorly in Experiment 1? Is this due to the Pythia 6.9 model or the SFT dataset?

We suspect that the poor performance is due to the fine tuned Pythia models.

• In the original papers, the PTQ baselines (QuaRot, SpinQuant, DuQuant) are only proposed for evaluation on pre-trained Llama models. As such, we speculate that an adaptive rotation configuration is necessary to achieve good performance for LLMs other than the Llama family.
• Most of these methods were not evaluated on fine-tuned benchmarks. We speculate that fine tuned models may introduce certain features that are not addressed in existing PTQ methods.

Please include measurements of inference speedup and memory usage.

Since our model's structure (cf. Fig 4) is equivalent to QuaRot, the inference time comparison can be referred to [Ashkboos et al., 2024b] for an extensive measurement on the inference speedup and inference memory consumption. E.g., Fig 7 therein shows that 4-bit quantized models with Hadamard rotation has 4x speedup compared to full-precision models, Fig 4 also shows significant advantages on time-to-first-token and peak memory saving. We will include these references in the revision.

We sincerely appreciate the reviewer for acknowledging the strength of our approach. Below we summarize our response to your concerns.

it is worth noting that the theoretical claims depend on certain assumptions (such as the interpolation condition and specific properties of the Gram matrix) that could limit their applicability in some scenarios.

We agree. However, we point out that the analysis in Sec 4 aims only at giving theoretical insights for the design of RoSTE algo, rather than ensuring convergence for STE training on LLMs, which is an open problem as modern LLM involves complex architecture. Despite the added assumptions, we retain essential elements of RoSTE training with linear layer featuring a rotation matrix in activations and weights, and the derived results gives motivation for RoSTE to use quantization error in the lower level objective of (11). To our best knowledge, this is the first analysis that provides insights on the convergence of STE training with rotation. Our derived bound echoes with the observation in Table 2 and Figure 3, where RoSTE empirically outperforms STE without rotation on quantization error and downstream task accuracy; also see the new experiments in Table A.1, A.2, A.3, A.4 in the response to Rev. c8mS.

Can you provide fine-tuning results of smaller models using Tulu-3 SFT dataset?

Since Tulu 3 is originally proposed for Llama 3.1 base models, extending the application of Tulu-3 SFT dataset to other families of base models would require an extensive fine-tuning on the hyperparameters. Due to limited time, we cannot produce additional results on Tulu-3 dataset. We will include additional results in the revision.

It seems that there are far more cases where a model that has already been fine-tuned needs to be quantized. Can the QA-SFT technique be practically applied in various situations?

It is correct that there are many unquantized but fine-tuned model that are publicly available. The proposed RoSTE can be applied on them through either (i) knowledge distillation, e.g., LLM-QAT in [Liu et al., 2023], or (ii) treating them as initialization using the original fine-tuning dataset. Note that all our numerical results indicate that existing PTQ methods may fail to maintain the fine-tuned accuracy after quantization. Our work thus demonstrate to practitioners that a good QA-SFT strategy (e.g., RoSTE) can greatly improve the quality of quantized & fine-tuned model.

Table A.1: Results on Pythia 1B model.

Table A.2: Results on Pythia 6.9B model.

Table A.3: Results on Qwen 2.5 0.5B model.

Table A.4: Results on Qwen 2.5 7B model.

The authors should either provide stronger empirical evidence to justify their claim or acknowledge the limitations of their theoretical analysis in real-world LLM scenarios.

We agree with your points. However, we point out that the analysis in Sec 4 aims at giving theoretical insights for the design of RoSTE algo, rather than ensuring convergence for STE training on LLMs, which is an open problem to our best knowledge as modern LLM involves complex architecture. Despite the limitations raised, we retain essential elements of RoSTE with a linear layer featuring a rotation matrix in activations and weights, and the derived results motivated RoSTE to minimize quantization error in the lower level objective (11). The derived bound echoes with our Table 2 and Fig 3. We will emphasize these limitations in the revision.

primarily evaluates RoSTE on Pythia (1B/6.9B) and Llama (8B) ...

We extended our experiments to the latest Qwen 2.5 models - see Table A.3, A.4 on accuracy of 4-bit quantized 0.5B / 7B models. RoSTE delivers the top quantized model vs the baselines.

Given that efficiency is a primary motivation... quantify the trade-offs between accuracy improvements and the increased computational cost.

We agree, but would clarify that there are two aspects of efficiency wrt computational cost - inference and training, which we will detail below:

• Inference efficiency: Since our model's structure (cf. Fig 4) is equivalent to QuaRot which implements weight and activation quantization with Hadamard rotation, the inference time comparison can be referred to [Ashkboos et al., 2024b] for an extensive measurement on the inference speedup and inference memory consumption. E.g., Fig 7 therein shows that 4-bit quantized models with Hadamard rotation has 4x speedup compared to full-precision models, their Fig 4 also shows significant advantages on time-to-first-token and peak memory saving.
• Training efficiency: As we restricted the search space of rotation matrices to Hadamard matrix, the theoretical training complexity should be on par with full-param SFT. This is confirmed in the additional Table B of the response to Reviewer mQhc. We acknowledge that the training complexity can still be higher than parameter-efficient methods such as QLoRA. For a comprehensive comparison, we refer to figure on the accuracy-vs-training-time tradeoff for SOTA methods. Observe that RoSTE achieves the best accuracy at a moderate cost in training complexity.

We will include the above references in the revision.

the paper does not provide detailed benchmarks on key efficiency metrics ...

We apologize for the misunderstanding caused by ambiguous wording of the title. We have referred to "efficiency" in the context of deployment for trained quantized model, where RoSTE delivers models with low latency inference and memory consumption. Nevertheless, in the training phase, RoSTE is on par with full-param SFT as seen in Table B in the response to Rev mQhc.

Memory Efficiency: ...

We agree partially with your points. While RoSTE maybe less memory efficient than methods such as QLoRA, we observe that RoSTE has a similar training cost in compute and memory usage as full-param SFT (cf. Table B in response to Rev. mQhc). The overhead with adaptive rotation is insignificant due to the small outer iteration no.

Our goal lies in achieving optimal performance at deployment by tuning fully quantized models with inference speedup and downstream task accuracy. As a comparison, note that Table 10 of [arxiv/2404.14047] shows models such as QLoRA suffer from degraded inference speed due to extra unquantized LoRA layers.

From a practical standpoint, our experiments on <8B models is relevant to on-device deployment where LLMs of such scales are common. E.g., the fine tuned 4-bit Pythia 1B model by RoSTE has greatly improved performance over PTQ on the full-precision base model. This will be useful for applications that prioritize inference latency.

For 70B+ models, RoSTE is still feasible and can potentially improve model quality. First, a rough estimate suggests that using RoSTE to training Llama3 70B under Adam require ~900GB of GPU memory (see here), which is implementable on a 8xH200 cluster. Second, recent studies such as [arxiv/2407.11062] showed that using QAT on 70B+ models with 1xA100 is possible with a simple block processing technique. Such technique can be straightforwardly combined with RoSTE.

Limited Learning Space: ...

This is a good point. Searching for optimal rotation like SpinQuant may enhance the model quality, yet our design to limit the learning space to $\\{I, H\\}$ for each layer is grounded on the tradeoff of complexity vs performance:

• Quantization Error: As shown in Fig 3, using Hadamard matrices suffices to reduce the quantization error and optimize the lower-level objective in (11);
• Training: extending the search space to all rotation matrices leads (11) to become a mix of non-smooth and manifold optimization problems. Handling it requires substantially higher training complexity.
• Inference: Hadamard matrices are efficient for inference [Ashkboos et al., 2024b], where it supports hardware acceleration and has low memory footprint due to its integer structure. In contrast, optimal rotation requires real-valued rotation matrices that impose significant overheads during inference.

Our design choice is supported by results in Table 2 & 3: RoSTE outperforms learnt rotations (SpinQuant) & no-rotation (STE).

Empirical Evaluation: ...

We extended our experiment results to QLoRA in the above Table A.1, A.2, A.3, A.4 in the response to Rev. c8mS. Notice that QLoRA under-performs RoSTE in terms of quantized model accuracy on downstream tasks, potentially due to the limited learning space of low-rank adaptation.

Missing analysis of memory usage and training time.

We included details on training time and training memory usage in Table B in the response to Rev. mQhc. Notice that RoSTE has a similar training cost as STE.

Scalability to larger models (70B+) is not demonstrated or discussed.

As mentioned in the first response, scaling RoSTE to 70B+ models is feasible through either

• Scaling up computation capability.
• Applying block processing to RoSTE.

While we do not have the computation resource nor time to experiment with 70B+ models at the moment, our codebase is available for the community to run further experiments.

Some experiments use outdated ... Broader eval needed ...

We extended our experiment results to Qwen 2.5 models (released on Sep 25 2024). By Table A in the above, we illustrate the accuracy of 4-bit quantized Qwen 2.5 7B models. RoSTE remains to produce the top performance quantized model among the baselines. Now, our evaluation set covers three families of LLMs (Pythia, Qwen, Llama). Our coverage is broader than existing literature such as QuaRot and SpinQuant, where only Llama models were evaluated.

Novelty is limited ...

We respectfully disagree. From a methodology perspective, our work is the first to employ adaptive rotation with SFT to improve the accuracy of quantized models. Such design is well grounded in theory from an optimization perspective.

Besides innovation in methodology, we observe a research gap where prior works mainly considered pre-trained zero-shot evaluation, while the performance of quantized fine-tuned models remains unexplored. Our results open the door to a new benchmark for approaches that directly optimize the SFT objective on quantized models. We further illustrate the existing research gap in figure on accuracy-vs-training-time tradeoffs. Note that RoSTE achieves better accuracy than existing methods on SFT benchmark with a moderate training overhead.