ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add support for QAT full fine-tuning (#3238)

Browse source 
**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|>
```
This commit is contained in:
andrewor14 2025-09-08 18:07:50 -04:00 • committed by GitHub
commit 6f2228d108
3 changed files with 54 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

33
unsloth/models/_utils.py
View file

@ -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

41
unsloth/models/llama.py
View file

@ -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"

20
unsloth/models/loader.py
View file

@ -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