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

Remove unused code and simplify quantization handling in CLI

Browse source 
Removes unused `TestUnslothTrainingArguments` test class. Simplifies GGUF quantization logic in CLI by eliminating redundant list wrapping and intermediate variable. Removes deprecated `forced_merged_4bit` save method choice from CLI (kept `merged_4bit_forced` alias in save.py for backward compatibility). Fixes `UnslothTrainingArguments` to not store unused `embedding_learning_rate` attribute.
This commit is contained in:
Can 2026-01-17 09:27:08 +03:00 • committed by Daniel Han
commit 7f594b9457
4 changed files with 8 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

16
tests/test_asft.py
View file

@ -1472,21 +1472,5 @@ class TestASFTTrainerComputeLoss:
assert loss_mock.call_args.kwargs["original_model"] is None
class TestUnslothTrainingArguments:
"""Tests for UnslothTrainingArguments."""
def test_embedding_learning_rate_is_set(self):
"""Test embedding_learning_rate is stored on the args object."""
from unsloth import trainer as trainer_module
with patch.object(
trainer_module.TrainingArguments, "__init__", return_value = None
) as base_init:
args = trainer_module.UnslothTrainingArguments(embedding_learning_rate = 0.01)
assert args.embedding_learning_rate == 0.01
assert base_init.called
if __name__ == "__main__":
pytest.main([__file__, "-v"])

18
unsloth-cli.py
View file

@ -190,13 +190,8 @@ def run(args):
if args.save_model:
# if args.quantization_method is a list, we will save the model for each quantization method
if args.save_gguf:
quantization_methods = (
args.quantization
if isinstance(args.quantization, list)
else [args.quantization]
)
if len(quantization_methods) > 1:
for quantization_method in quantization_methods:
if isinstance(args.quantization, list):
for quantization_method in args.quantization:
print(
f"Saving model with quantization method: {quantization_method}"
)
@ -212,18 +207,17 @@ def run(args):
quantization_method = quantization_method,
)
else:
quantization_method = quantization_methods[0]
print(f"Saving model with quantization method: {quantization_method}")
print(f"Saving model with quantization method: {args.quantization}")
model.save_pretrained_gguf(
args.save_path,
tokenizer,
quantization_method = quantization_method,
quantization_method = args.quantization,
)
if args.push_model:
model.push_to_hub_gguf(
hub_path = args.hub_path,
hub_token = args.hub_token,
quantization_method = quantization_method,
quantization_method = args.quantization,
)
else:
model.save_pretrained_merged(args.save_path, tokenizer, args.save_method)
@ -438,7 +432,7 @@ def build_parser():
"--save_method",
type = str,
default = "merged_16bit",
choices = ["merged_16bit", "merged_4bit", "forced_merged_4bit", "lora"],
choices = ["merged_16bit", "merged_4bit", "lora"],
help = "Save method for the model, default is 'merged_16bit'",
)
save_group.add_argument(

2
unsloth/save.py
View file

@ -290,7 +290,7 @@ def unsloth_save_model(
"if you're planning to do multiple saves.\n"
"If you are certain, change `save_method` to `merged_4bit_forced`."
)
elif save_method in {"merged_4bit_forced", "forced_merged_4bit"}:
elif save_method == "merged_4bit_forced":
save_method = "merged_4bit"
save_pretrained_settings = dict(locals())

2
unsloth/trainer.py
View file

@ -141,7 +141,7 @@ except:
class UnslothTrainingArguments(TrainingArguments):
def __init__(self, embedding_learning_rate: float = None, *args, **kwargs):
self.embedding_learning_rate = embedding_learning_rate
embedding_learning_rate = embedding_learning_rate
super().__init__(*args, **kwargs)