Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 81 additions & 0 deletions tests/saving/test_quant_method_none_normalization.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
"""CPU-only regression for the quant-method normalization loops in save.py.

`unsloth_save_pretrained_gguf` and `save_to_gguf_generic` each normalize the
`quantization_method` list, mapping a ``None`` element to ``"q8_0"``. The mapping
used to call ``quant_method.lower()`` as the first statement of the loop, so a
``None`` element (e.g. ``quantization_method=[None]`` or ``["q4_k_m", None]``)
raised ``AttributeError: 'NoneType' object has no attribute 'lower'`` and the
``elif quant_method is None`` branch was unreachable dead code.

The loop is inline inside two heavy functions (importing unsloth needs
unsloth_zoo / a GPU), so - like test_is_gpt_oss_detection.py - we extract just the
loop source via ``ast`` and exec it against sample inputs. That exercises the real
source: it fails on the old ordering and passes once ``None`` is handled first.
"""

from __future__ import annotations

import ast
from pathlib import Path

import pytest

SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
SAVE_SRC = SAVE_PY.read_text(encoding = "utf-8")
SAVE_TREE = ast.parse(SAVE_SRC, filename = str(SAVE_PY))

# The target functions and the list variable each one appends the normalized method to.
TARGETS = (
("unsloth_save_pretrained_gguf", "quantization_methods"),
("save_to_gguf_generic", "new_quantization_methods"),
)


def _func(tree, name):
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == name:
return node
raise AssertionError(f"function {name!r} not found in {SAVE_PY.name}")


def _quant_loop(func_name):
# The quant-normalization `for` loop iterates `quantization_method`; grab its source.
func = _func(SAVE_TREE, func_name)
for node in ast.walk(func):
if (
isinstance(node, ast.For)
and isinstance(node.iter, ast.Call)
and isinstance(node.iter.func, ast.Name)
and node.iter.func.id == "enumerate"
and isinstance(node.iter.args[0], ast.Name)
and node.iter.args[0].id == "quantization_method"
):
return node
raise AssertionError(f"quant-normalization loop not found in {func_name}")


def _run_loop(func_name, out_var, quantization_method):
# exec just the extracted loop against a given input, returning the appended methods.
loop_src = ast.get_source_segment(SAVE_SRC, _quant_loop(func_name))
namespace = {out_var: [], "quantization_method": quantization_method}
exec(loop_src, {"__builtins__": __builtins__}, namespace)
return namespace[out_var]


@pytest.mark.parametrize("func_name, out_var", TARGETS)
def test_none_element_maps_to_q8_0(func_name, out_var):
# A bare None inside the list must map to q8_0, not raise AttributeError.
assert _run_loop(func_name, out_var, [None]) == ["q8_0"]


@pytest.mark.parametrize("func_name, out_var", TARGETS)
def test_none_mixed_with_strings(func_name, out_var):
# None resolves to q8_0 while sibling string methods are still normalized (lowercased).
assert _run_loop(func_name, out_var, ["Q4_K_M", None]) == ["q4_k_m", "q8_0"]


@pytest.mark.parametrize("func_name, out_var", TARGETS)
def test_string_methods_unchanged(func_name, out_var):
# The fix must not alter behavior for the ordinary string inputs.
methods = ["not_quantized", "fast_quantized", "quantized", "Q8_0"]
assert _run_loop(func_name, out_var, methods) == ["f16", "q8_0", "q4_k_m", "q8_0"]
14 changes: 8 additions & 6 deletions unsloth/save.py
Original file line number Diff line number Diff line change
Expand Up @@ -2926,15 +2926,16 @@ def unsloth_save_pretrained_gguf(
"Unsloth: quantization_method can only be a string or a list of strings"
)
for i, quant_method in enumerate(quantization_method):
quant_method = quant_method.lower()
if quant_method is None:
quant_method = "q8_0"
else:
quant_method = quant_method.lower()
if quant_method == "not_quantized":
quant_method = "f16"
elif quant_method == "fast_quantized":
quant_method = "q8_0"
elif quant_method == "quantized":
quant_method = "q4_k_m"
elif quant_method is None:
quant_method = "q8_0"
quantization_methods.append(quant_method.lower())

try:
Expand Down Expand Up @@ -3727,15 +3728,16 @@ def save_to_gguf_generic(
"Unsloth: quantization_method can only be a string or a list of strings"
)
for i, quant_method in enumerate(quantization_method):
quant_method = quant_method.lower()
if quant_method is None:
quant_method = "q8_0"
else:
quant_method = quant_method.lower()
if quant_method == "not_quantized":
quant_method = "f16"
elif quant_method == "fast_quantized":
quant_method = "q8_0"
elif quant_method == "quantized":
quant_method = "q4_k_m"
elif quant_method is None:
quant_method = "q8_0"
new_quantization_methods.append(quant_method.lower())
else:
new_quantization_methods.append(quantization_type.lower())
Expand Down