From f3ab8c21afb2d24c74a84ba383c8ff9eb4a1a881 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 26 Aug 2025 04:53:39 -0700 Subject: [PATCH] Support `model.save_pretrained_torchao` (#3111) 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: --- tests/saving/test_unsloth_save.py | 46 ++++++++++++++++++++++++++++ unsloth/save.py | 50 +++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py index 14e8af4f07..22dd38c8d1 100644 --- a/tests/saving/test_unsloth_save.py +++ b/tests/saving/test_unsloth_save.py @@ -3,6 +3,7 @@ import os import shutil import tempfile import pytest +import importlib from unsloth import FastLanguageModel, FastModel @@ -167,3 +168,48 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str): load_in_4bit=True, ) +@pytest.mark.skipif(importlib.util.find_spec("torchao") is None) +def test_save_torchao(model, tokenizer, temp_save_dir: str): + save_path = os.path.join(temp_save_dir, "unsloth_torchao", model.config._name_or_path.replace("/", "_")) + + from torchao.quantization import Int8DynamicActivationInt8WeightConfig + torchao_config = Int8DynamicActivationInt8WeightConfig() + model.save_pretrained_torchao( + save_path, + tokenizer=tokenizer, + torchao_config=torchao_config, + ) + + # Check model files + assert os.path.isdir(save_path), f"Directory {save_path} does not exist." + assert os.path.isfile(os.path.join(save_path, "config.json")), "config.json not found." + + weight_files = [f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors")] + assert len(weight_files) > 0, "No weight files found in the save directory." + + # Check tokenizer files + for file in tokenizer_files: + assert os.path.isfile(os.path.join(save_path, file)), f"{file} not found in the save directory." + + # Store the size of the model files + total_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files) + save_file_sizes["merged_4bit"][model.config._name_or_path] = total_size + + print(f"Total size of merged_4bit files: {total_size} bytes") + + assert total_size < save_file_sizes["merged_16bit"][model.config._name_or_path], "Merged 4bit files are larger than merged 16bit files." + + # Check config to see if it is 4bit + config_path = os.path.join(save_path, "config.json") + with open(config_path, "r") as f: + config = json.load(f) + + assert "quantization_config" in config, "Quantization config not found in the model config." + + # Test loading the model from the saved path + loaded_model, loaded_tokenizer = FastModel.from_pretrained( + save_path, + max_seq_length=128, + dtype=None, + load_in_4bit=True, + ) diff --git a/unsloth/save.py b/unsloth/save.py index 9539b66701..4e7c7c6755 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -2122,6 +2122,55 @@ 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) + model = AutoModelForCausalLM.from_pretrained( + new_save_directory, + torch_dtype="auto", + device_map="auto", + quantization_config=quantization_config, + ) + + 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): # Create the custom directory if it doesn't exist @@ -2597,6 +2646,7 @@ def patch_saving_functions(model, vision = False): 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.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