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

Support saving locally in model.save_pretrained_torchao (#3263)

Browse source 
Summary:
Previously the test was not ran correctly and the save to local path is not tested
this PR added support for that and tries to test properly

Note: `python tests/saving/test_unsloth_save.py` doesn't run test

Test Plan:
pytest tests/saving/test_unsloth_save.py -k test_save_torchao

Reviewers:

Subscribers:

Tasks:

Tags:
This commit is contained in:
Jerry Zhang 2025-09-03 17:51:33 -07:00 • committed by GitHub
commit 43d1f1fc2d
2 changed files with 16 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

18
tests/saving/test_unsloth_save.py
View file

@ -26,6 +26,7 @@ model_to_test = [
save_file_sizes = {}
save_file_sizes["merged_16bit"] = {}
save_file_sizes["merged_4bit"] = {}
save_file_sizes["torchao"] = {}
tokenizer_files = [
"tokenizer_config.json",
@ -168,7 +169,7 @@ 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)
@pytest.mark.skipif(importlib.util.find_spec("torchao") is None, reason="require torchao to be installed")
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("/", "_"))
@ -178,6 +179,7 @@ def test_save_torchao(model, tokenizer, temp_save_dir: str):
save_path,
tokenizer=tokenizer,
torchao_config=torchao_config,
push_to_hub=False,
)
# Check model files
@ -193,13 +195,13 @@ def test_save_torchao(model, tokenizer, temp_save_dir: str):
# 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
save_file_sizes["torchao"][model.config._name_or_path] = total_size
print(f"Total size of merged_4bit files: {total_size} bytes")
# merged_16bit tests are not running yet, so we can't test this for now
# TODO: enable this after merged_16bit is fixed
# assert total_size < save_file_sizes["merged_16bit"][model.config._name_or_path], "torchao files are larger than merged 16bit files."
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
# Check config to see if it is quantized with torchao
config_path = os.path.join(save_path, "config.json")
with open(config_path, "r") as f:
config = json.load(f)
@ -207,9 +209,11 @@ def test_save_torchao(model, tokenizer, temp_save_dir: str):
assert "quantization_config" in config, "Quantization config not found in the model config."
# Test loading the model from the saved path
# can't set `load_in_4bit` to True because the model is torchao quantized
# can't quantize again with bitsandbytes
loaded_model, loaded_tokenizer = FastModel.from_pretrained(
save_path,
max_seq_length=128,
dtype=None,
load_in_4bit=True,
load_in_4bit=False,
)

7
unsloth/save.py
View file

@ -2116,6 +2116,7 @@ def unsloth_push_to_hub_gguf(
pass
if fix_bos_token:
logger.warning(
"Unsloth: ##### The current model auto adds a BOS token.\n"\
"Unsloth: ##### We removed it in GGUF's chat template for you."
@ -2511,7 +2512,7 @@ def unsloth_save_pretrained_torchao(
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
`push_to_hub` (bool): whether to push the checkpoint to huggingface hub or save locally
"""
# first merge the lora weights
arguments = dict(locals())
@ -2550,7 +2551,9 @@ def unsloth_save_pretrained_torchao(
# 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
else:
model.save_pretrained(save_directory, safe_serialization=False)
tokenizer.save_pretrained(save_directory)
pass
for _ in range(3):
gc.collect()