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

Merge branch 'main' of https://github.com/unslothai/unsloth

Browse source
This commit is contained in:
Daniel Han 2025-09-08 17:15:57 -07:00
commit 4ae5db3287
4 changed files with 58 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

43
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"
@ -2986,6 +2949,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 +2994,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():

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

2
unsloth/models/vision.py
View file

@ -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():