From f89c829fcf3135c3b625b4eebad0fe7e59babfd1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 22 Jun 2026 02:11:46 -0700 Subject: [PATCH] Fix save crash for legacy list-form _tied_weights_keys (NemotronH) (#6540) * Fix save crash for legacy list-form _tied_weights_keys (NemotronH) transformers >= 5 save_pretrained reads module._tied_weights_keys.keys(), which raises 'list' object has no attribute 'keys' for modules that still declare the attribute as a list (e.g. NemotronH backbone.layers.N.mixer.*_proj), crashing GGUF export and merged saves part-way through. Coerce any legacy list/tuple _tied_weights_keys into the dict form transformers 5.x expects, mapping each key to itself. Only the keys are read (as dedup patterns) so behaviour is preserved, and older transformers that iterate the attribute directly see the same keys. The helper is idempotent and best-effort so a save never fails over it. Called from unsloth_save_model, unsloth_save_pretrained_gguf and unsloth_generic_save after tokenizer patching. Adds version-independent unit tests covering list/tuple coercion, dict and None/empty pass-through, idempotency and odd-object tolerance. * Coerce empty/set _tied_weights_keys too transformers only skips _tied_weights_keys when it is None, so an empty list, tuple or set still reaches .keys() and raises the same AttributeError. Coerce every non-dict container (including the empty case and sets) to a dict, and add tests for empty/set inputs. * Tighten comments in tied-weights save fix * Scope tied-weights-keys coercion to the save call Coercing legacy list-form _tied_weights_keys to {k: k} fixed the transformers 5 save crash, but persisted a self-mapping on the live model. transformers 5 re-ties from the dict's values, so a later resize/re-tie would no-op the tie instead of pointing the output weights back at the input embeddings. Replace the in-place mutation with a decorator that coerces before the save and restores the originals afterwards (including on exception), so the save sees the dict form transformers needs while the model keeps its original tie metadata. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim comments to be more succinct --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../test_normalize_tied_weights_keys.py | 116 ++++++++++++++++++ unsloth/save.py | 59 +++++++++ 2 files changed, 175 insertions(+) create mode 100644 tests/saving/test_normalize_tied_weights_keys.py 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,