From a14cf9a7322c6fc6086aeb2632cab7151e1eb363 Mon Sep 17 00:00:00 2001
From: andrewor14 <andrewor14@gmail.com>
Date: Mon, 18 Aug 2025 08:56:35 -0400
Subject: [PATCH] Add support for QAT + LoRA (#2976)

**Summary:** Quantization-aware training (QAT) helps mitigate
quantization degradation by simulating quantization numerics
in high precision during training (fake quantization). This PR
combines QAT with LoRA by applying torchao's QAT support to
the peft model.

See the following for more details:

- torchao QAT: https://github.com/pytorch/ao/blob/main/torchao/quantization/qat/README.md
- torchtune QAT + LoRA: https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700

Current QAT schemes supported are:

```
fp8-fp8, targeting the torch.ops.fbgemm.f8i4bf16_shuffled kernel
fp8-int4, targeting the torch.ops.fbgemm.f8f8bf16_rowwise kernel
```

**Test Plan:**

```
from unsloth import FastLanguageModel

lora_rank = 32

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Qwen3-4B-Base",
    max_seq_length = 2048,
    load_in_4bit = False,
    fast_inference = False,
    max_lora_rank = lora_rank,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = lora_rank,
    target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha = lora_rank*2,
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    qat_scheme = "fp8-fp8",
)

lora.Linear(
  (base_layer): FakeQuantizedLinear(
    in_features=2560, out_features=4096, bias=False
    (activation_fake_quantizer): FakeQuantizer(Float8FakeQuantizeConfig(dtype=torch.float8_e4m3fn, granularity=PerRow(), hp_value_lb=None, hp_value_ub=None))
    (weight_fake_quantizer): FakeQuantizer(Float8FakeQuantizeConfig(dtype=torch.float8_e4m3fn, granularity=PerRow(), hp_value_lb=None, hp_value_ub=None))
  )
  ...
)
```
---
 unsloth/models/llama.py | 49 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 48 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index ae03a685eb..231b17154d 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -16,7 +16,7 @@ import torch
 import gc
 import math
 import functools
-from typing import Optional, Tuple, List, Union
+from typing import Any, Dict, Optional, Tuple, List, Union
 from ._utils import *
 from ._utils import patch_unsloth_smart_gradient_checkpointing
 from ._utils import __version__
@@ -113,6 +113,46 @@ torch_nn_functional_softmax = torch.nn.functional.softmax
 # SDPA has GQA internally
 SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__
 
+
+def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.Module:
+    """
+    Apply QAT + LoRA during fine-tuning.
+
+    On a high level, this means fake quantizing the base (frozen) model during LoRA training.
+    Fake quantization refers to simulating quantization numerics in high precision (e.g. bf16).
+    This helps mitigate quantization degradations when the model is quantized after training.
+
+    For more details: https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700
+    """
+    try:
+        from torchao.quantization import (
+            Float8DynamicActivationFloat8WeightConfig,
+            Float8DynamicActivationInt4WeightConfig,
+            PerRow,
+            quantize_,
+        )
+        from torchao.quantization.qat import QATConfig
+    except ImportError as e:
+        print(
+            "Please install torchao nightly for the latest QAT features:\n"
+            "  pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126"
+        )
+        raise e
+    pass
+    filter_fn = None
+    if qat_scheme == "fp8-int4":
+        group_size = 128
+        base_config = Float8DynamicActivationInt4WeightConfig(group_size=group_size)
+        filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
+    elif qat_scheme == "fp8-fp8":
+        base_config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
+    else:
+        raise ValueError(f"Unexpected QAT scheme {qat_scheme}")
+    pass
+    quantize_(model, QATConfig(base_config, step="prepare"), filter_fn=filter_fn)
+    return model
+pass
+
 # Fix new HF's inference code
 def _fast_prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs,):
     past_key_values = kwargs.get("past_key_values", None)
@@ -2248,6 +2288,7 @@ class FastLlamaModel:
         init_lora_weights   = True,
         loftq_config        = {},
         temporary_location  = "_unsloth_temporary_saved_buffers",
+        qat_scheme          = None,
         **kwargs,
     ):
         if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1":
@@ -2614,6 +2655,12 @@ class FastLlamaModel:
 
         model = _get_peft_model(model, lora_config)
 
+        # Apply QAT + LoRA if specified
+        if qat_scheme is not None:
+            print("Unsloth: Applying QAT to mitigate quantization degradation")
+            model = _prepare_model_for_qat(model, qat_scheme)
+        pass
+
         model._saved_temp_tokenizer = _saved_temp_tokenizer
 
         model = FastLlamaModel.patch_peft_model(model, use_gradient_checkpointing)