From fb13275787451bd67bba88e546c94e4636f353e2 Mon Sep 17 00:00:00 2001
From: Etherll <61019402+Etherll@users.noreply.github.com>
Date: Sun, 17 May 2026 14:02:36 +0300
Subject: [PATCH] fix(sentence_transformer): resume PEFT checkpoints under
 sentence-transformers >= 5.4 (#5454)

Saves the base config.json next to adapter_config.json when checkpointing PEFT-wrapped sentence-transformer models, and overrides SentenceTransformerTrainer._load_from_checkpoint to load adapter weights via set_peft_model_state_dict and rebuild aux modules (Pooling, Normalize, Dense) from modules.json with strict type and path validation. Patches only activate on Unsloth-managed Transformer modules so non-Unsloth pipelines fall through to upstream behaviour.

Fixes https://github.com/unslothai/unsloth/issues/5373
---
 unsloth/models/sentence_transformer.py | 181 ++++++++++++++++++++++++-
 1 file changed, 180 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py
index c53e3a7a81..aafc416221 100644
--- a/unsloth/models/sentence_transformer.py
+++ b/unsloth/models/sentence_transformer.py
@@ -527,6 +527,45 @@ This sentence-transformers model was finetuned and converted to GGUF format usin
 
 
 class FastSentenceTransformer(FastModel):
+    @staticmethod
+    def _save_base_config_for_processor_resume(config, output_path):
+        """
+        sentence-transformers >= 5.4 reloads Transformer modules through
+        AutoProcessor. Tokenizer-only checkpoint roots make AutoProcessor fall
+        back to AutoConfig, so PEFT adapter checkpoints still need the base
+        config.json next to adapter_config.json.
+        """
+        if config is None or not getattr(config, "model_type", None):
+            return
+        if hasattr(config, "save_pretrained"):
+            config.save_pretrained(output_path)
+        elif hasattr(config, "to_json_file"):
+            config_path = os.path.join(output_path, "config.json")
+            config.to_json_file(config_path)
+
+    @staticmethod
+    def _patch_transformer_module_save_config(transformer_module, base_config = None):
+        transformer_module._unsloth_st_managed = True
+        if base_config is not None and getattr(base_config, "model_type", None):
+            transformer_module._unsloth_base_config = base_config
+
+        if getattr(transformer_module, "_unsloth_save_config_patched", False):
+            return transformer_module
+
+        original_save = transformer_module.save
+
+        def _save_with_base_config(self, output_path, *args, **kwargs):
+            original_save(output_path, *args, **kwargs)
+            FastSentenceTransformer._save_base_config_for_processor_resume(
+                getattr(self, "_unsloth_base_config", None), output_path
+            )
+
+        transformer_module.save = types.MethodType(
+            _save_with_base_config, transformer_module
+        )
+        transformer_module._unsloth_save_config_patched = True
+        return transformer_module
+
     @staticmethod
     def _read_pooling_mode(model_name, token):
         """
@@ -1157,6 +1196,9 @@ class FastSentenceTransformer(FastModel):
                 config_keys.append(config_key)
         transformer_module.config_keys = config_keys
         transformer_module.save_in_root = True
+        FastSentenceTransformer._patch_transformer_module_save_config(
+            transformer_module, getattr(model, "config", None)
+        )
 
         if hasattr(model, "config"):
             model.config.tokenizer_class = tokenizer.__class__.__name__
@@ -1644,6 +1686,9 @@ class FastSentenceTransformer(FastModel):
             st_model._dtype = dtype
             st_model._load_in_4bit = load_in_4bit
             st_model.no_modules = False
+            FastSentenceTransformer._patch_transformer_module_save_config(
+                st_model[0], getattr(st_model[0].auto_model, "config", None)
+            )
 
             # Add save methods
             def _save_pretrained_merged(self, save_directory, **save_kwargs):
@@ -2067,6 +2112,9 @@ class FastSentenceTransformer(FastModel):
                     transformer_module.model = peft_model
                 else:
                     transformer_module.auto_model = peft_model
+                FastSentenceTransformer._patch_transformer_module_save_config(
+                    transformer_module, getattr(inner_model, "config", None)
+                )
 
                 # Store compile info for auto-compile at trainer time
                 # torch.compile is deferred until training starts so we can check max_steps
@@ -2121,6 +2169,9 @@ class FastSentenceTransformer(FastModel):
                 transformer_module.model = peft_model
             else:
                 transformer_module.auto_model = peft_model
+            FastSentenceTransformer._patch_transformer_module_save_config(
+                transformer_module, getattr(inner_model, "config", None)
+            )
             return model
         else:
             return FastModel.get_peft_model(
@@ -2235,5 +2286,133 @@ def _patch_sentence_transformer_trainer():
     SentenceTransformerTrainer._unsloth_auto_compile_patched = True
 
 
-# Auto-patch trainer on module import
+def _patch_st_trainer_load_from_checkpoint():
+    try:
+        from sentence_transformers import SentenceTransformerTrainer
+    except ImportError:
+        return
+    if getattr(
+        SentenceTransformerTrainer, "_unsloth_load_from_checkpoint_patched", False
+    ):
+        return
+    if not hasattr(SentenceTransformerTrainer, "_load_from_checkpoint"):
+        return
+
+    _original = SentenceTransformerTrainer._load_from_checkpoint
+
+    def _unsloth_load_from_checkpoint(self, checkpoint_path):
+        try:
+            from peft import PeftModel, load_peft_weights, set_peft_model_state_dict
+        except ImportError:
+            return _original(self, checkpoint_path)
+
+        try:
+            mod0 = self.model[0]
+        except (IndexError, TypeError):
+            return _original(self, checkpoint_path)
+
+        if isinstance(getattr(type(mod0), "auto_model", None), property):
+            inner = getattr(mod0, "model", None)
+        else:
+            inner = getattr(mod0, "auto_model", None)
+        inner = getattr(inner, "_orig_mod", inner)
+
+        if not isinstance(inner, PeftModel):
+            return _original(self, checkpoint_path)
+        if not getattr(mod0, "_unsloth_st_managed", False):
+            return _original(self, checkpoint_path)
+
+        if not any(
+            os.path.isfile(os.path.join(checkpoint_path, fn))
+            for fn in ("adapter_model.safetensors", "adapter_model.bin")
+        ):
+            return _original(self, checkpoint_path)
+
+        adapter_name = getattr(inner, "active_adapter", None)
+        if adapter_name is None and callable(getattr(inner, "active_adapters", None)):
+            adapter_name = inner.active_adapters()
+        if isinstance(adapter_name, (list, tuple, set)):
+            if len(adapter_name) != 1:
+                raise RuntimeError(
+                    "Unsloth: Cannot resume multiple active PEFT adapters."
+                )
+            adapter_name = next(iter(adapter_name))
+        adapter_name = adapter_name or "default"
+        if adapter_name not in getattr(inner, "peft_config", {}):
+            raise RuntimeError(f"Unsloth: PEFT adapter {adapter_name!r} is not loaded.")
+
+        load_result = set_peft_model_state_dict(
+            inner, load_peft_weights(checkpoint_path), adapter_name = adapter_name
+        )
+        unexpected = getattr(load_result, "unexpected_keys", []) or []
+        missing = [
+            x
+            for x in (getattr(load_result, "missing_keys", []) or [])
+            if f".{adapter_name}." in x or x.endswith(f".{adapter_name}")
+        ]
+        if unexpected or missing:
+            raise RuntimeError(
+                "Unsloth: PEFT checkpoint does not match the active adapter "
+                f"(missing={missing[:8]}, unexpected={unexpected[:8]})."
+            )
+
+        modules_json = os.path.join(checkpoint_path, "modules.json")
+        if not os.path.isfile(modules_json):
+            raise RuntimeError("Unsloth: PEFT checkpoint is missing modules.json.")
+        try:
+            with open(modules_json, "r") as f:
+                module_configs = json.load(f)
+        except Exception as e:
+            raise RuntimeError("Unsloth: Cannot parse checkpoint modules.json.") from e
+
+        root = os.path.abspath(os.fspath(checkpoint_path))
+        restored = set()
+        for entry in module_configs:
+            idx = int(entry.get("idx", -1))
+            if idx == 0:
+                continue
+            if idx < 0 or idx >= len(self.model):
+                raise RuntimeError(f"Unsloth: Bad module index in modules.json: {idx}.")
+            module = self.model[idx]
+            module_cls = type(module)
+            saved_type = entry.get("type", "")
+            if saved_type and not saved_type.endswith(f".{module_cls.__name__}"):
+                raise RuntimeError(f"Unsloth: Checkpoint module {idx} type mismatch.")
+            module_path = entry.get("path")
+            module_dir = os.path.abspath(
+                os.path.join(root, os.fspath(module_path or ""))
+            )
+            try:
+                inside_root = os.path.commonpath([root, module_dir]) == root
+            except ValueError:
+                inside_root = False
+            if not module_path or not inside_root or not os.path.isdir(module_dir):
+                raise RuntimeError(
+                    f"Unsloth: Bad checkpoint module path for index {idx}."
+                )
+            if not hasattr(module_cls, "load"):
+                raise RuntimeError(f"Unsloth: Module {idx} cannot be reloaded.")
+            fresh = module_cls.load(module_dir)
+            if not isinstance(fresh, module_cls):
+                raise RuntimeError(f"Unsloth: Module {idx} reload returned wrong type.")
+            # Parameterless modules (Pooling, Normalize) make
+            # next(module.parameters()) raise StopIteration; route through
+            # the SentenceTransformer's device property instead.
+            try:
+                fresh.to(self.model.device)
+            except AttributeError:
+                pass
+            self.model[idx] = fresh
+            restored.add(idx)
+        missing_idx = sorted(set(range(1, len(self.model))) - restored)
+        if missing_idx:
+            raise RuntimeError(
+                f"Unsloth: Checkpoint modules.json is incomplete (missing idx={missing_idx[:8]})."
+            )
+
+    SentenceTransformerTrainer._load_from_checkpoint = _unsloth_load_from_checkpoint
+    SentenceTransformerTrainer._unsloth_load_from_checkpoint_patched = True
+
+
 _patch_sentence_transformer_trainer()
+_patch_st_trainer_load_from_checkpoint()