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 <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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"]

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,