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
This commit is contained in:
Etherll
GitHub
parent
e20bbeff9a
commit
fb13275787
1 changed files with 180 additions and 1 deletions
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue