**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|>
```
Summary:
Previously the test was not ran correctly and the save to local path is not tested
this PR added support for that and tries to test properly
Note: `python tests/saving/test_unsloth_save.py` doesn't run test
Test Plan:
pytest tests/saving/test_unsloth_save.py -k test_save_torchao
Reviewers:
Subscribers:
Tasks:
Tags:
* Update test_qwen3_grpo.py to correct function call
This test file uses the incorrect name for the function, which is gradient_checkpointing_disable(), not disable_gradient_checkpointing().
I copied the line from test_llama32_sft.py - I'm not sure if this actually is required, just wanted it consistent for when other people like me test this and have no clue what they're doing when it throws an exception.
* Update blackwell/test_qwen3_grpo.py
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* Update mistral.py, showed flag to not call cut cross entropy
* Update mistral.py, made it so if its not equal to zero
* Update unsloth/models/mistral.py
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Summary:
Allow users merge the LoRA weights and then do a post training quantization with torchao
Usage:
```
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
torchao_config = Int8DynamicActivationInt8WeightConfig()
model.save_pretrained_torchao(
save_path,
tokenizer=tokenizer,
torchao_config=torchao_config,
)
```
Test Plan:
python tests/saving/test_unsloth_save.py
Reviewers:
Subscribers:
Tasks:
Tags:
