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

Merge branch 'main' into nightly

Browse source
This commit is contained in:
Daniel Han 2025-04-19 19:46:17 -07:00
commit 9bdeab9427
4 changed files with 216 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

169
tests/saving/test_unsloth_save.py Normal file
View file

@ -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,
)

11
unsloth/models/loader.py
View file

@ -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,
)

36
unsloth/models/vision.py
View file

@ -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,

12
unsloth/save.py
View file

@ -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,