diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py new file mode 100644 index 0000000000..14e8af4f07 --- /dev/null +++ b/tests/saving/test_unsloth_save.py @@ -0,0 +1,169 @@ +import json +import os +import shutil +import tempfile +import pytest + +from unsloth import FastLanguageModel, FastModel + +model_to_test = [ + # Text Models + "unsloth/tinyllama", + "unsloth/tinyllama-bnb-4bit", + "unsloth/Qwen2.5-0.5B-Instruct", + "unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit", + "unsloth/Phi-4-mini-instruct", + "unsloth/Phi-4-mini-instruct-bnb-4bit", + "unsloth/Qwen2.5-0.5B", + # Vision Models + "unsloth/gemma-3-1b-it", + "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit", + "unsloth/Qwen2.5-VL-3B-Instruct-bnb-4bit" +] + +# Variables +save_file_sizes = {} +save_file_sizes["merged_16bit"] = {} +save_file_sizes["merged_4bit"] = {} + +tokenizer_files = [ + "tokenizer_config.json", + "special_tokens_map.json", +] + +@pytest.fixture(scope="session", params=model_to_test) +def loaded_model_tokenizer(request): + model_name = request.param + print("Loading model and tokenizer...") + + model, tokenizer = FastModel.from_pretrained( + model_name, # use small model + max_seq_length=128, + dtype=None, + load_in_4bit=True, + ) + + # Apply LoRA + model = FastModel.get_peft_model( + model, + r=16, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], + lora_alpha=16, + use_gradient_checkpointing="unsloth", + ) + + return model, tokenizer + +@pytest.fixture(scope="session") +def model(loaded_model_tokenizer): + return loaded_model_tokenizer[0] + +@pytest.fixture(scope="session") +def tokenizer(loaded_model_tokenizer): + return loaded_model_tokenizer[1] + +@pytest.fixture +def temp_save_dir(): + dir = tempfile.mkdtemp() + print(f"Temporary directory created at: {dir}") + yield dir + print(f"Temporary directory deleted: {dir}") + shutil.rmtree(dir) + + +def delete_quantization_config(model): + # Since merged, edit quantization_config + old_config = model.config + new_config = model.config.to_dict() + if "quantization_config" in new_config: + del new_config["quantization_config"] + original_model = model + new_config = type(model.config).from_dict(new_config) + while hasattr(original_model, "model"): + original_model = original_model.model + original_model.config = new_config + model.config = new_config + +def test_save_merged_16bit(model, tokenizer, temp_save_dir: str): + save_path = os.path.join(temp_save_dir, "unsloth_merged_16bit", model.config._name_or_path.replace("/", "_")) + + model.save_pretrained_merged( + save_path, + tokenizer=tokenizer, + save_method="merged_16bit" + ) + + # 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." + + # Check config to see if it is 16bit by checking for quantization config + config_path = os.path.join(save_path, "config.json") + with open(config_path, "r") as f: + config = json.load(f) + + assert "quantization_config" not in config, "Quantization config not found in the model config." + + # 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_16bit"][model.config._name_or_path] = total_size + print(f"Total size of merged_16bit files: {total_size} bytes") + + # Test loading the model from the saved path + loaded_model, loaded_tokenizer = FastLanguageModel.from_pretrained( + save_path, + max_seq_length=128, + dtype=None, + load_in_4bit=True, + ) + +def test_save_merged_4bit(model, tokenizer, temp_save_dir: str): + save_path = os.path.join(temp_save_dir, "unsloth_merged_4bit", model.config._name_or_path.replace("/", "_")) + + model.save_pretrained_merged( + save_path, + tokenizer=tokenizer, + save_method="merged_4bit_forced" + ) + + # 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/models/loader.py b/unsloth/models/loader.py index cac5acd838..3cd8508ffa 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -469,10 +469,14 @@ class FastModel(FastBaseModel): return_logits = False, # Return logits fullgraph = True, # No graph breaks use_exact_model_name = False, + auto_model = None, + whisper_language = None, + whisper_task = None, *args, **kwargs, ): if token is None: token = get_token() - + if whisper_language is not None: assert(type(whisper_language) is str) + if whisper_task is not None: assert(type(whisper_task) is str) SUPPORTS_BFLOAT16 = is_bfloat16_supported() if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 @@ -709,7 +713,8 @@ class FastModel(FastBaseModel): # Check if VLM is_vlm = any(x.endswith("ForConditionalGeneration") for x in model_config.architectures) is_vlm = is_vlm or hasattr(model_config, "vision_config") - auto_model = AutoModelForVision2Seq if is_vlm else AutoModelForCausalLM + if auto_model is None: + auto_model = AutoModelForVision2Seq if is_vlm else AutoModelForCausalLM model, tokenizer = FastBaseModel.from_pretrained( model_name = model_name, @@ -727,6 +732,8 @@ class FastModel(FastBaseModel): auto_model = auto_model, use_gradient_checkpointing = use_gradient_checkpointing, supports_sdpa = supports_sdpa, + whisper_language = whisper_language, + whisper_task = whisper_task, *args, **kwargs, ) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index f05cc95d60..d212c224be 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -236,6 +236,8 @@ class FastBaseModel: auto_model = AutoModelForVision2Seq, use_gradient_checkpointing = "unsloth", supports_sdpa = True, + whisper_language = None, + whisper_task = None, **kwargs, ): if model_types is None: @@ -304,7 +306,8 @@ class FastBaseModel: do_forced_float32 = True pass # Stop SDPA for some archs like Pixtral / Mistral3 - kwargs["attn_implementation"] = "sdpa" + if not ("attn_implementation" in kwargs): + kwargs["attn_implementation"] = "sdpa" if not supports_sdpa: print(f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to eager!") del kwargs["attn_implementation"] @@ -352,6 +355,7 @@ class FastBaseModel: # Check if using forced float32 - we load it in bfloat16, then cast to float16! torch_dtype = dtype if do_forced_float32: torch_dtype = torch.bfloat16 + model = auto_model.from_pretrained( model_name, device_map = device_map, @@ -367,12 +371,23 @@ class FastBaseModel: # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name - auto_processor = AutoProcessor if auto_model is AutoModelForVision2Seq else AutoTokenizer - tokenizer = auto_processor.from_pretrained( - tokenizer_name, - padding_side = "right", - token = token, - ) + is_vlm = (auto_model is AutoModelForVision2Seq) + is_whisper = (whisper_language is not None and whisper_task is not None) + auto_processor = AutoProcessor if (is_vlm or is_whisper) else AutoTokenizer + if whisper_language and whisper_task: + tokenizer = auto_processor.from_pretrained( + tokenizer_name, + padding_side = "right", + token = token, + language = whisper_language, + task = whisper_task, + ) + else: + tokenizer = auto_processor.from_pretrained( + tokenizer_name, + padding_side = "right", + token = token, + ) if hasattr(tokenizer, "tokenizer"): __tokenizer = tokenizer.tokenizer # Add padding side as well @@ -469,6 +484,7 @@ class FastBaseModel: modules_to_save = None, init_lora_weights = True, loftq_config = {}, + task_type = TaskType.CAUSAL_LM, temporary_location = "_unsloth_temporary_saved_buffers", **kwargs, ): @@ -492,7 +508,7 @@ class FastBaseModel: finetune_attention_modules = True finetune_mlp_modules = True pass - if target_modules is None: + if target_modules is None or target_modules == "all-linear": target_modules = get_peft_regex( model, finetune_vision_layers = finetune_vision_layers, @@ -503,7 +519,7 @@ class FastBaseModel: else: assert(type(target_modules) in (list, tuple,)) pass - + # Clear deleted GPU items for _ in range(3): gc.collect() @@ -516,7 +532,7 @@ class FastBaseModel: target_modules = target_modules, lora_dropout = lora_dropout, bias = bias, - task_type = TaskType.CAUSAL_LM, + task_type = task_type, ) model = prepare_model_for_kbit_training( model, diff --git a/unsloth/save.py b/unsloth/save.py index b8da9c08d0..e3eece6cca 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -2301,6 +2301,17 @@ def unsloth_generic_save( maximum_memory_usage : float = 0.9, ): if token is None and push_to_hub: token = get_token() + + if save_method == "merged_4bit": + raise RuntimeError( + "Unsloth: Merging into 4bit will cause your model to lose accuracy if you plan\n"\ + "to merge to GGUF or others later on. I suggest you to do this as a final step\n"\ + "if you're planning to do multiple saves.\n"\ + "If you are certain, change `save_method` to `merged_4bit_forced`." + ) + elif save_method == "merged_4bit_forced": + save_method = "merged_4bit" + merge_and_overwrite_lora( get_model_name, model = model, @@ -2309,6 +2320,7 @@ def unsloth_generic_save( push_to_hub = push_to_hub, private = private, token = token, + save_method = save_method, output_dtype = None, low_disk_space_usage = True, use_temp_file = False,