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

Update save.py

Browse source
This commit is contained in:
Daniel Han 2025-09-03 15:19:02 -07:00
commit 10d0f25746
1 changed files with 62 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

117
unsloth/save.py
View file

@ -2123,58 +2123,6 @@ def unsloth_push_to_hub_gguf(
pass
pass
def unsloth_save_pretrained_torchao(
self,
save_directory : Union[str, os.PathLike],
tokenizer = None,
torchao_config = None,
push_to_hub : bool = False,
):
"""Quantizes the model with torchao and saves a torchao quantized checkpoint
Args
`save_directory`: local folder path or huggingface hub ID when `push_to_hub` is set to True, e.g. `my_model`
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization, full list: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
`push_to_hub` (bool): whether to push the checkpoint to huggingface hub or not
"""
# first merge the lora weights
arguments = dict(locals())
arguments["save_directory"] = save_directory + "-local"
arguments["model"] = self
arguments["tokenizer"] = tokenizer
arguments["push_to_hub"] = False # We save ourselves
arguments["save_method"] = "merged_16bit" # Must be 16bit
del arguments["self"]
del arguments["torchao_config"]
new_save_directory, old_username = unsloth_save_model(**arguments)
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from torchao import quantize_
if torchao_config is None:
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
torchao_config = Int8DynamicActivationInt8WeightConfig()
quantization_config = TorchAoConfig(quant_type=torchao_config)
tokenizer = AutoTokenizer.from_pretrained(new_save_directory)
if HAS_TORCH_DTYPE:
kwargs = {"torch_dtype" : "auto"}
else:
kwargs = {"dtype" : "auto"}
model = AutoModelForCausalLM.from_pretrained(
new_save_directory,
device_map="auto",
quantization_config=quantization_config,
**kwargs,
)
if push_to_hub:
save_to = save_directory
# torchao does not support safe_serialization right now
model.push_to_hub(save_to, safe_serialization=False)
tokenizer.push_to_hub(save_to)
pass
pass
pass
# Corrected function to save LoRA to a custom directory
def save_lora_to_custom_dir(model, tokenizer, save_directory):
@ -2550,6 +2498,65 @@ def unsloth_generic_push_to_hub_merged(
pass
def unsloth_save_pretrained_torchao(
self,
save_directory : Union[str, os.PathLike],
tokenizer = None,
torchao_config = None,
push_to_hub : bool = False,
token : Optional[Union[str, bool]] = None,
):
"""Quantizes the model with torchao and saves a torchao quantized checkpoint
Args
`save_directory`: local folder path or huggingface hub ID when `push_to_hub` is set to True, e.g. `my_model`
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization, full list: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
`push_to_hub` (bool): whether to push the checkpoint to huggingface hub or not
"""
# first merge the lora weights
arguments = dict(locals())
arguments["save_directory"] = save_directory + "-local"
arguments["model"] = self
arguments["tokenizer"] = tokenizer
arguments["push_to_hub"] = False # We save ourselves
arguments["save_method"] = "merged_16bit" # Must be 16bit
del arguments["self"]
del arguments["torchao_config"]
unsloth_generic_save(**arguments)
for _ in range(3):
gc.collect()
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from torchao import quantize_
if torchao_config is None:
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
torchao_config = Int8DynamicActivationInt8WeightConfig()
quantization_config = TorchAoConfig(quant_type = torchao_config)
tokenizer = AutoTokenizer.from_pretrained(arguments["save_directory"])
if HAS_TORCH_DTYPE:
kwargs = {"torch_dtype" : "auto"}
else:
kwargs = {"dtype" : "auto"}
model = AutoModelForCausalLM.from_pretrained(
arguments["save_directory"],
device_map = "auto",
quantization_config = quantization_config,
**kwargs,
)
if push_to_hub:
if token is None and push_to_hub: token = get_token()
# torchao does not support safe_serialization right now
model.push_to_hub(save_directory, safe_serialization = False, token = token)
tokenizer.push_to_hub(save_directory, token = token)
pass
pass
for _ in range(3):
gc.collect()
pass
def not_implemented_save(*args, **kwargs):
raise NotImplementedError("Unsloth: Sorry GGUF is currently not supported for vision models!")
pass
@ -2647,11 +2654,11 @@ def patch_saving_functions(model, vision = False):
if not vision:
if hasattr(model, "config"):
# Counteract tokenizers
model.push_to_hub_merged = types.MethodType(unsloth_generic_push_to_hub_merged, model)
model.save_pretrained_merged = types.MethodType(unsloth_generic_save_pretrained_merged, model)
model.push_to_hub_merged = types.MethodType(unsloth_generic_push_to_hub_merged, model)
model.save_pretrained_merged = types.MethodType(unsloth_generic_save_pretrained_merged, model)
model.push_to_hub_gguf = types.MethodType(unsloth_push_to_hub_gguf, model)
model.save_pretrained_gguf = types.MethodType(unsloth_save_pretrained_gguf, model)
model.save_pretrained_torchao = types.MethodType(unsloth_save_pretrained_torchao, model)
model.save_pretrained_torchao = types.MethodType(unsloth_save_pretrained_torchao, model)
model.push_to_hub_ggml = types.MethodType(unsloth_convert_lora_to_ggml_and_push_to_hub, model)
model.save_pretrained_ggml = types.MethodType(unsloth_convert_lora_to_ggml_and_save_locally, model)
pass