diff --git a/tests/saving/test_normalize_tied_weights_keys.py b/tests/saving/test_normalize_tied_weights_keys.py new file mode 100644 index 0000000000..d18c9877ea --- /dev/null +++ b/tests/saving/test_normalize_tied_weights_keys.py @@ -0,0 +1,116 @@ +"""Unit tests for the tied-weights-keys coercion used by unsloth.save. + +Regression for the NemotronH save / GGUF-export crash: transformers >= 5 +``save_pretrained`` reads ``_tied_weights_keys.keys()`` and raises on the legacy list +form. Exercised on tiny module trees, no model download. +""" + +import pytest +import torch + +from unsloth.save import ( + _coerce_tied_weights_keys_to_dict, + _normalize_tied_weights_keys_for_save, + _restore_tied_weights_keys, +) + + +def _build_tree(): + root = torch.nn.Module() + mixer = torch.nn.Module() + root.add_module("mixer", mixer) + return root, mixer + + +def test_list_becomes_dict_and_restores(): + root, mixer = _build_tree() + mixer._tied_weights_keys = ["q_proj.weight", "o_proj.weight"] + originals = _coerce_tied_weights_keys_to_dict(root) + assert mixer._tied_weights_keys == { + "q_proj.weight": "q_proj.weight", + "o_proj.weight": "o_proj.weight", + } + _restore_tied_weights_keys(originals) + assert mixer._tied_weights_keys == ["q_proj.weight", "o_proj.weight"] + + +def test_tuple_and_set_become_dict(): + root, mixer = _build_tree() + root._tied_weights_keys = ("lm_head.weight",) + mixer._tied_weights_keys = {"q_proj.weight"} + _coerce_tied_weights_keys_to_dict(root) + assert root._tied_weights_keys == {"lm_head.weight": "lm_head.weight"} + assert mixer._tied_weights_keys == {"q_proj.weight": "q_proj.weight"} + + +def test_empty_containers_become_dict(): + root, mixer = _build_tree() + root._tied_weights_keys = [] + mixer._tied_weights_keys = () + _coerce_tied_weights_keys_to_dict(root) + # transformers skips only None; an empty list still hits .keys(). + assert root._tied_weights_keys == {} and mixer._tied_weights_keys == {} + + +def test_none_and_existing_dict_are_left_unchanged(): + root, mixer = _build_tree() + root._tied_weights_keys = None + original = {"a.weight": "b.weight"} + mixer._tied_weights_keys = original + originals = _coerce_tied_weights_keys_to_dict(root) + assert root._tied_weights_keys is None + assert mixer._tied_weights_keys is original # untouched, not rebuilt + assert originals == [] # nothing to restore + + +def test_model_without_modules_method_does_not_raise(): + class NoModules: + pass + + assert _coerce_tied_weights_keys_to_dict(NoModules()) == [] + + +def test_decorator_coerces_during_save_then_restores(): + root, mixer = _build_tree() + mixer._tied_weights_keys = ["lm_head.weight"] + seen = {} + + @_normalize_tied_weights_keys_for_save + def save(model): + seen["keys"] = dict(model.mixer._tied_weights_keys) + return "ok" + + assert save(root) == "ok" + # Dict form was visible to the save, list form restored afterwards. + assert seen["keys"] == {"lm_head.weight": "lm_head.weight"} + assert mixer._tied_weights_keys == ["lm_head.weight"] + + +def test_decorator_restores_on_exception(): + root, mixer = _build_tree() + mixer._tied_weights_keys = ["lm_head.weight"] + + @_normalize_tied_weights_keys_for_save + def save(model): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError): + save(root) + assert mixer._tied_weights_keys == ["lm_head.weight"] + + +def test_decorator_finds_model_in_kwargs_and_positional(): + # unsloth_save_model / unsloth_generic_save pass model= as a keyword; the gguf path + # binds it as the first positional (method ``self``). Both must be coerced. + for call in (lambda f, r: f(model = r), lambda f, r: f(r)): + root, mixer = _build_tree() + mixer._tied_weights_keys = ["w.weight"] + captured = {} + + @_normalize_tied_weights_keys_for_save + def save(model): + captured["dict"] = isinstance(model.mixer._tied_weights_keys, dict) + + call(save, root) + assert captured["dict"] is True + assert mixer._tied_weights_keys == ["w.weight"] diff --git a/unsloth/save.py b/unsloth/save.py index a6cc665d3e..0e8c8cdc2d 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -44,6 +44,7 @@ import json import shutil import pickle import gc +import functools from transformers.models.llama.modeling_llama import logger from .kernels import fast_dequantize, QUANT_STATE, get_lora_parameters_bias import subprocess @@ -489,6 +490,62 @@ def _qwen3_5_vlm_state_dict_for_save(state_dict): return remapped_state_dict +def _coerce_tied_weights_keys_to_dict(model): + """Coerce each module's legacy list/tuple/set ``_tied_weights_keys`` to dict form, + returning ``[(module, original), ...]`` for the caller to restore. + + transformers >= 5 ``save_pretrained`` reads ``_tied_weights_keys.keys()``, so a model + still declaring it as a list (e.g. NemotronH) crashes mid-save. + """ + originals = [] + try: + modules = list(model.modules()) + except Exception: + return originals + for module in modules: + keys = getattr(module, "_tied_weights_keys", None) + if isinstance(keys, (list, tuple, set)): + try: + module._tied_weights_keys = {k: k for k in keys} + originals.append((module, keys)) + except Exception: + pass + return originals + + +def _restore_tied_weights_keys(originals): + """Undo _coerce_tied_weights_keys_to_dict.""" + for module, keys in originals: + try: + module._tied_weights_keys = keys + except Exception: + pass + + +def _normalize_tied_weights_keys_for_save(save_fn): + """Coerce legacy list-form ``_tied_weights_keys`` to dict for the duration of a save, + then restore: transformers >= 5 re-ties from the dict's *values*, so a persisted + ``{k: k}`` self-map would no-op a later resize/re-tie. ``model`` is the first positional + arg (bound-method ``self``) or the ``model=`` keyword. + """ + + @functools.wraps(save_fn) + def wrapper(*args, **kwargs): + model = kwargs.get("model") + if model is None and args: + model = args[0] + if model is None: + model = kwargs.get("self") + originals = _coerce_tied_weights_keys_to_dict(model) if model is not None else [] + try: + return save_fn(*args, **kwargs) + finally: + _restore_tied_weights_keys(originals) + + return wrapper + + +@_normalize_tied_weights_keys_for_save @torch.inference_mode def unsloth_save_model( model, @@ -2092,6 +2149,7 @@ def push_to_ollama(tokenizer, gguf_location, username: str, model_name: str, tag print("Successfully pushed to ollama") +@_normalize_tied_weights_keys_for_save def unsloth_save_pretrained_gguf( self, save_directory: Union[str, os.PathLike], @@ -2968,6 +3026,7 @@ def save_to_gguf_generic( return metadata +@_normalize_tied_weights_keys_for_save @torch.inference_mode def unsloth_generic_save( model,