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 5c22f06a9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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,
)