From 0d70391f9bcc8e5a14f8b8633e29e3014d862cd2 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Mon, 8 Sep 2025 15:47:32 -0500 Subject: [PATCH 1/2] GptAttention turn training off during inference (#3289) --- unsloth/models/llama.py | 2 ++ unsloth/models/vision.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index f978060c9c..09bb3e04eb 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2986,6 +2986,7 @@ class FastLlamaModel: _for_inference(m) m = m.model _for_inference(m) + model.eval() # to turn off training on modules deeper in # Since transformers 4.53, must turn off explicitly for module in model.modules(): @@ -3030,6 +3031,7 @@ class FastLlamaModel: _for_training(m) m = m.model _for_training(m) + model.train() # to turn on training on modules deeper in # Since transformers 4.53, must turn on explicitly for module in model.modules(): diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index cce6554d52..3c71543a09 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -772,6 +772,7 @@ class FastBaseModel: _for_inference(m) m = m.model _for_inference(m) + model.eval() # to turn off training on modules deeper in # Since transformers 4.53, must turn off explicitly for module in model.modules(): @@ -823,6 +824,7 @@ class FastBaseModel: _for_training(m) m = m.model _for_training(m) + model.train() # to turn on training on modules deeper in # Since transformers 4.53, must turn on explicitly for module in model.modules(): From 6f2228d1089a1418d33458bb9e448865a8769b87 Mon Sep 17 00:00:00 2001 From: andrewor14 Date: Mon, 8 Sep 2025 18:07:50 -0400 Subject: [PATCH 2/2] Add support for QAT full fine-tuning (#3238) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Summary:** Following https://github.com/unslothai/unsloth/pull/2976, which adds support for QAT + LoRA, this PR adds support for QAT during full fine-tuning. See the [torchao QAT README](https://github.com/pytorch/ao/blob/main/torchao/quantization/qat/README.md) for more details. Current QAT schemes supported are: ``` fp8-int4, targeting the torch.ops.fbgemm.f8i4bf16_shuffled kernel fp8-fp8, targeting the torch.ops.fbgemm.f8f8bf16_rowwise kernel ``` **Test Plan:** https://gist.github.com/andrewor14/048b5c1bd01b7fa23c53913856a8ef9f Full fine-tuning Llama3.1-8B with and without QAT on `yahma/alpaca-cleaned` for 1 epoch: - Batch size = 16 (no grad accum) - Learning rate = 4e-5 - Quantization scheme = fp8-int4 Wikitext perplexity: - QAT improved perplexity by 19.2% compared to regular fine-tuning - QAT's int4 quantized model even outperformed the bf16 baseline - Regular int4 quantized model (without QAT) was significantly worse than the bf16 baseline ``` ==> unsloth_model_full_baseline_output/eval_float.log <== | | |none | 0|word_perplexity|↓ |9.8446|± | N/A| ==> unsloth_model_full_baseline_output/eval_quantized.log <== | | |none | 0|word_perplexity|↓ |11.4595|± | N/A| ==> unsloth_model_full_qat_fp8-int4_output/eval_quantized.log <== | | |none | 0|word_perplexity|↓ |9.2336|± | N/A| ``` Fibonacci test: - Both bf16 baseline and int4 quantized models correctly identified 13 as the next number - QAT quantized model was more succinct in its response - No substantial differences here ``` ### Instruction: Continue the fibonnaci sequence. ### Input: 1, 1, 2, 3, 5, 8 ==> unsloth_model_full_baseline_output/eval_float.log <== ### Response: The next number in the Fibonacci sequence is 13.<|end_of_text|> ==> unsloth_model_full_baseline_output/eval_quantized.log <== ### Response: The next number in the Fibonacci sequence is 13.<|end_of_text|> ==> unsloth_model_full_qat_fp8-int4_output/eval_quantized.log <== ### Response: 13<|end_of_text|> ``` --- unsloth/models/_utils.py | 33 ++++++++++++++++++++++++++++++++ unsloth/models/llama.py | 41 ++-------------------------------------- unsloth/models/loader.py | 20 +++++++++++++++++++- 3 files changed, 54 insertions(+), 40 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index bcde34bb9e..597ed0244b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1528,3 +1528,36 @@ def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, m pass return loftq_config + + +def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.Module: + """ + Transform a model for Quantization-Aware Training (QAT) during fine-tuning. + + On a high level, this means fake quantizing the base (frozen) model during 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. + + QAT can be optionally combined with LoRA fine-tuning to for additional throughput improvement. + For more details: https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700 + """ + from torchao.quantization import ( + Float8DynamicActivationInt4WeightConfig, + Float8DynamicActivationFloat8WeightConfig, + PerRow, + quantize_, + ) + from torchao.quantization.qat import QATConfig + filter_fn = None + if qat_scheme == "fp8-int4": + group_size = 128 + base_config = Float8DynamicActivationInt4WeightConfig() + 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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 09bb3e04eb..4143db93b8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -21,6 +21,7 @@ from ._utils import * from ._utils import patch_unsloth_smart_gradient_checkpointing from ._utils import __version__ from ._utils import move_to_device +from ._utils import _prepare_model_for_qat from torch.nn.functional import scaled_dot_product_attention from transformers import __version__ as transformers_version from unsloth_zoo.utils import Version, _get_dtype @@ -115,45 +116,6 @@ torch_nn_functional_softmax = torch.nn.functional.softmax 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) @@ -1870,6 +1832,7 @@ class FastLlamaModel: disable_log_stats = False, unsloth_vllm_standby = False, num_labels = None, + qat_scheme = None, **kwargs, ): os.environ["UNSLOTH_USE_NEW_MODEL"] = "0" diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 25d03f188f..b1844a1472 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -13,6 +13,7 @@ # limitations under the License. from ._utils import ( + _prepare_model_for_qat, is_bfloat16_supported, is_vLLM_available, HAS_FLASH_ATTENTION, @@ -110,6 +111,7 @@ class FastLanguageModel(FastLlamaModel): random_state = 3407, max_lora_rank = 64, disable_log_stats = True, + qat_scheme = None, *args, **kwargs, ): # Login to allow private models @@ -120,7 +122,7 @@ class FastLanguageModel(FastLlamaModel): login(token = token) except: pass - if load_in_8bit or full_finetuning: + if load_in_8bit or full_finetuning or qat_scheme is not None: return FastModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, @@ -139,6 +141,7 @@ class FastLanguageModel(FastLlamaModel): return_logits = False, # Return logits fullgraph = True, # No graph breaks use_exact_model_name = use_exact_model_name, + qat_scheme = qat_scheme, *args, **kwargs, ) pass @@ -521,6 +524,7 @@ class FastModel(FastBaseModel): whisper_language = None, whisper_task = None, unsloth_force_compile = False, + qat_scheme = None, *args, **kwargs, ): if token is None: token = get_token() @@ -558,6 +562,13 @@ class FastModel(FastBaseModel): ) pass + if qat_scheme is not None and not full_finetuning: + raise ValueError( + "Specifying `qat_scheme` in `FastLanguageModel.from_pretrained(...)` is only " + "compatible with `full_finetuning=True`. If you wish to use QAT with LoRA, " + "please pass in `qat_scheme` in `FastLanguageModel.get_peft_model(...)` instead." + ) + old_model_name = model_name if not use_exact_model_name: model_name = get_model_name(model_name, load_in_4bit) @@ -921,6 +932,13 @@ class FastModel(FastBaseModel): # Patch it as well! model = FastBaseModel.post_patch_model(model, use_gradient_checkpointing, trust_remote_code = trust_remote_code) pass + + # Apply QAT 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 + return model, tokenizer pass pass