From 10d0f25746c109dc9ac119ef8a892a05b383e5da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 3 Sep 2025 15:19:02 -0700 Subject: [PATCH] Update save.py --- unsloth/save.py | 117 +++++++++++++++++++++++++----------------------- 1 file changed, 62 insertions(+), 55 deletions(-) diff --git a/unsloth/save.py b/unsloth/save.py index eace358cbb..62f79c4f5c 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -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