From 00fd190d8efa8d6ffd3224e9ea5be299c7a872f7 Mon Sep 17 00:00:00 2001 From: electroglyph Date: Fri, 19 Dec 2025 23:26:35 -0800 Subject: [PATCH] refactor _load_modules, add for_inference to from_pretrained, add transformers 5 code for mpnet, add distilbert patches --- unsloth/models/sentence_transformer.py | 694 ++++++++++++++++--------- 1 file changed, 450 insertions(+), 244 deletions(-) diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 89ac4088bc..0181fc006e 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -18,18 +18,23 @@ import json import os import types from huggingface_hub import hf_hub_download -from sentence_transformers.models import Transformer, Pooling, Normalize from transformers.models.mpnet import modeling_mpnet -from sentence_transformers.util import import_from_string, load_dir_path from typing import Optional import torch from transformers.modeling_outputs import BaseModelOutput from collections import OrderedDict +from transformers.models.distilbert import modeling_distilbert +from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa +import transformers +from packaging.version import Version +from transformers import AutoModel, AutoConfig - -class FastSentenceTransformer(FastModel): +class FastSentenceTransformer(FastModel): @staticmethod def _read_pooling_mode(model_name, token): + """ + Read the pooling mode from the modules.json file if it exists, otherwise return "mean". + """ try: if os.path.exists(model_name) and os.path.exists( os.path.join(model_name, "modules.json") @@ -83,108 +88,361 @@ class FastSentenceTransformer(FastModel): return mode except Exception as e: - print(f"Failed to detect pooling mode: {e}, defaulting to mean pooling.") + print(f"\033[1;33mFailed to detect pooling mode, not a sentence-transformers model. You will have to handle pooling/normalization yourself for inference, but training should be fine.\033[0m") return "mean" # should prolly be done upstream instead of this hackfest here @staticmethod - def _patch_mpnet(): - try: - # add supports_gradient_checkpointing flag - modeling_mpnet.MPNetModel.supports_gradient_checkpointing = True + def _patch_mpnet_v4(): + """ + Patch the MPNetModel to support gradient checkpointing. + Supports transformers 4. + """ + # add supports_gradient_checkpointing flag + modeling_mpnet.MPNetModel.supports_gradient_checkpointing = True - # add _set_gradient_checkpointing method - def _set_gradient_checkpointing(self, module = None, value = True): - if module is None: - module = self.encoder - if isinstance(module, modeling_mpnet.MPNetEncoder): - module.gradient_checkpointing = value + # add _set_gradient_checkpointing method + def _set_gradient_checkpointing(self, module = None, value = True): + if module is None: + module = self.encoder + if isinstance(module, modeling_mpnet.MPNetEncoder): + module.gradient_checkpointing = value - modeling_mpnet.MPNetModel._set_gradient_checkpointing = ( - _set_gradient_checkpointing - ) + modeling_mpnet.MPNetModel._set_gradient_checkpointing = ( + _set_gradient_checkpointing + ) - # patch MPNetEncoder.forward to support checkpointing - # based on: - # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/mpnet/modeling_mpnet.py#L321 - # but head_mask is no longer used in transformers 5.0.0: - # https://github.com/huggingface/transformers/blob/main/src/transformers/models/mpnet/modeling_mpnet.py#L284 - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - output_attentions: bool = False, - output_hidden_states: bool = False, - return_dict: bool = False, - **kwargs, - ): - # backwards compatibility for older transformers versions (4.57.3 and below) - head_mask = kwargs.pop("head_mask", None) - - position_bias = self.compute_position_bias(hidden_states) - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - - for i, layer_module in enumerate(self.layer): - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - # do gradient checkpointing if enabled and training - if getattr(self, "gradient_checkpointing", False) and self.training: - - def create_custom_forward(module): - # bog standard checkpoint - def custom_forward(*inputs): - return module( - *inputs, output_attentions = output_attentions - ) - - return custom_forward - - layer_outputs = torch.utils.checkpoint.checkpoint( - create_custom_forward(layer_module), - hidden_states, - attention_mask, - head_mask[i] if head_mask is not None else None, - position_bias, - use_reentrant = False, - ) - else: - # original code from here on - layer_outputs = layer_module( - hidden_states, - attention_mask, - head_mask[i] if head_mask is not None else None, - position_bias, - output_attentions = output_attentions, - **kwargs, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[1],) + # patch MPNetEncoder.forward to support checkpointing + # based on: + # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/mpnet/modeling_mpnet.py#L321 + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = False, + **kwargs, + ): + position_bias = self.compute_position_bias(hidden_states) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + for i, layer_module in enumerate(self.layer): if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if not return_dict: - return tuple( - v - for v in [hidden_states, all_hidden_states, all_attentions] - if v is not None + # do gradient checkpointing if enabled and training + if getattr(self, "gradient_checkpointing", False) and self.training: + + def create_custom_forward(module): + # bog standard checkpoint + def custom_forward(*inputs): + return module( + *inputs, output_attentions = output_attentions + ) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(layer_module), + hidden_states, + attention_mask, + head_mask[i] if head_mask is not None else None, + position_bias, + use_reentrant = False, ) - return BaseModelOutput( - last_hidden_state = hidden_states, - hidden_states = all_hidden_states, - attentions = all_attentions, + else: + # original code from here on + layer_outputs = layer_module( + hidden_states, + attention_mask, + head_mask[i] if head_mask is not None else None, + position_bias, + output_attentions = output_attentions, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, all_hidden_states, all_attentions] + if v is not None ) + return BaseModelOutput( + last_hidden_state = hidden_states, + hidden_states = all_hidden_states, + attentions = all_attentions, + ) - # assign the patched forward - modeling_mpnet.MPNetEncoder.forward = forward + # assign the patched forward + modeling_mpnet.MPNetEncoder.forward = forward - except Exception as e: - print(f"Unsloth: Failed to patch MPNet for gradient checkpointing: {e}") + @staticmethod + def _patch_mpnet_v5(): + """ + Patch the MPNetModel to support gradient checkpointing. + Supports transformers 5. + """ + # add supports_gradient_checkpointing flag + modeling_mpnet.MPNetModel.supports_gradient_checkpointing = True + + # add _set_gradient_checkpointing method + def _set_gradient_checkpointing(self, module = None, value = True): + if module is None: + module = self.encoder + if isinstance(module, modeling_mpnet.MPNetEncoder): + module.gradient_checkpointing = value + + modeling_mpnet.MPNetModel._set_gradient_checkpointing = ( + _set_gradient_checkpointing + ) + + # patch MPNetEncoder.forward to support checkpointing + # based on: + # https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/models/mpnet/modeling_mpnet.py#L284 + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = False, + **kwargs, + ): + + position_bias = self.compute_position_bias(hidden_states) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # do gradient checkpointing if enabled and training + if getattr(self, "gradient_checkpointing", False) and self.training: + + def create_custom_forward(module): + # checkpoint + def custom_forward(*inputs): + return module(*inputs, output_attentions = output_attentions) + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(layer_module), + hidden_states, + attention_mask, + position_bias, + use_reentrant = False, + ) + else: + # original code from here on + layer_outputs = layer_module( + hidden_states, + attention_mask, + position_bias, + output_attentions, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, all_hidden_states, all_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state = hidden_states, + hidden_states = all_hidden_states, + attentions = all_attentions, + ) + + modeling_mpnet.MPNetEncoder.forward = forward + + @staticmethod + def _patch_distilbert_v4(): + # change kwargs to positional args to be compatible with peft_utils + """ + Patch the forward method of the DistilBertModel to use positional arguments instead of keyword arguments. + Transformers 4 version. + """ + # based on: + # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/distilbert/modeling_distilbert.py#L666 + # original code from here on: + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + head_mask_is_none = head_mask is None + # Prepare head mask if needed + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embeddings = self.embeddings(input_ids, inputs_embeds) # (bs, seq_length, dim) + + if self.config._attn_implementation == "flash_attention_2": + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + else: + if attention_mask is None: + attention_mask = torch.ones(input_shape, device=device) # (bs, seq_length) + + if self.config._attn_implementation == "sdpa" and head_mask_is_none and not output_attentions: + attention_mask = _prepare_4d_attention_mask_for_sdpa( + attention_mask, embeddings.dtype, tgt_len=input_shape[1] + ) + # patch here, change kwargs to positional args: + return self.transformer( + embeddings, + attention_mask, + head_mask, + output_attentions, + output_hidden_states, + return_dict, + ) + + modeling_distilbert.DistilBertModel.forward = forward + + @staticmethod + def _patch_distilbert_v5(): + """ + Patch the forward method of the DistilBertModel to use positional arguments instead of keyword arguments. + Transformers 5 version. + """ + # based on: + # https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/models/distilbert/modeling_distilbert.py#L386 + # original code from here on: + from transformers.masking_utils import create_bidirectional_mask + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + **kwargs, + ): + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + embeddings = self.embeddings(input_ids, inputs_embeds, position_ids) + + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=embeddings, + attention_mask=attention_mask, + ) + + # patch here: unsloth gradient checkpointing hook needs positional arguments + return self.transformer( + embeddings, + attention_mask, + **kwargs, + ) + modeling_distilbert.DistilBertModel.forward = forward + + @staticmethod + def _module_path(model_name, token=None): + """ + Returns the path to the modules.json file or None + """ + try: + if os.path.exists(model_name) and os.path.isdir(model_name): + path = os.path.join(model_name, "modules.json") + return path if os.path.exists(path) else None + else: + try: + return hf_hub_download(model_name, "modules.json", token=token) + except: + return None + except: + return None + + @staticmethod + def _create_transformer_module( + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + ): + """Helper to create and configure a Transformer module.""" + from sentence_transformers.models import Transformer + + transformer_module = Transformer( + model_name, + max_seq_length = max_seq_length, + model_args = {"trust_remote_code": trust_remote_code}, + config_args = {"trust_remote_code": trust_remote_code}, + ) + transformer_module.auto_model = model + transformer_module.tokenizer = tokenizer + transformer_module.do_lower_case = getattr(tokenizer, "do_lower_case", False) + + # sentence-transformers only passes along known keys to model.forward + model_forward_params = list(inspect.signature(model.forward).parameters) + transformer_module.model_forward_params = set(model_forward_params) | { + "input_ids", + "attention_mask", + "token_type_ids", + "inputs_embeds", + } + + # determine max_seq_length if not provided + if max_seq_length is None: + if hasattr(model, "config") and hasattr(model.config, "max_position_embeddings"): + max_seq_length = model.config.max_position_embeddings + elif hasattr(tokenizer, "model_max_length"): + max_seq_length = tokenizer.model_max_length + else: + max_seq_length = 512 + + transformer_module.max_seq_length = max_seq_length + transformer_module.config_keys = ["max_seq_length", "do_lower_case"] + transformer_module.save_in_root = True + + if hasattr(model, "config"): + model.config.tokenizer_class = tokenizer.__class__.__name__ + + return transformer_module @staticmethod def _load_modules( @@ -195,174 +453,70 @@ class FastSentenceTransformer(FastModel): max_seq_length, pooling_mode, trust_remote_code = False, - ): + ) -> tuple[OrderedDict, bool]: + """ + Load modules from modules.json if available, otherwise fallback to hard-coded modules. + + Returns: + tuple[OrderedDict, bool]: (modules, no_modules_json) + """ + from sentence_transformers.util import import_from_string, load_dir_path + from sentence_transformers.models import Pooling, Normalize + modules = OrderedDict() - - # grope around for modules.json - modules_json_path = None - if os.path.exists(model_name) and os.path.exists( - os.path.join(model_name, "modules.json") - ): - modules_json_path = os.path.join(model_name, "modules.json") - else: - try: - modules_json_path = hf_hub_download( - model_name, "modules.json", token = token - ) - except: - pass - - if modules_json_path and os.path.exists(modules_json_path): - with open(modules_json_path, encoding = "utf8") as f: + modules_json_path = FastSentenceTransformer._module_path(model_name, token) + + if modules_json_path: + with open(modules_json_path, encoding="utf8") as f: modules_config = json.load(f) for module_config in modules_config: class_ref = module_config["type"] - name = ( - module_config["name"] - if "name" in module_config - else str(module_config.get("idx", len(modules))) - ) + name = module_config.get("name", str(module_config.get("idx", len(modules)))) - # main module if class_ref == "sentence_transformers.models.Transformer": - transformer_module = Transformer( - model_name, - max_seq_length = max_seq_length, - model_args = {"trust_remote_code": trust_remote_code}, - config_args = {"trust_remote_code": trust_remote_code}, + transformer_module = FastSentenceTransformer._create_transformer_module( + model_name, model, tokenizer, max_seq_length, trust_remote_code ) - transformer_module.auto_model = model - transformer_module.tokenizer = tokenizer - - # move tokenizer do_lower_case to transformer module - transformer_module.do_lower_case = getattr( - tokenizer, "do_lower_case", False - ) - model_forward_params = list( - inspect.signature(model.forward).parameters - ) - transformer_module.model_forward_params = set( - model_forward_params - ) | { - "input_ids", - "attention_mask", - "token_type_ids", - "inputs_embeds", - } - if max_seq_length is None: - pass - - # is this overkill? should we just force user to set it? - current_max_seq = max_seq_length - if current_max_seq is None: - if hasattr(model, "config") and hasattr( - model.config, "max_position_embeddings" - ): - current_max_seq = model.config.max_position_embeddings - elif hasattr(tokenizer, "model_max_length"): - current_max_seq = tokenizer.model_max_length - else: - current_max_seq = 512 - - transformer_module.max_seq_length = current_max_seq - transformer_module.config_keys = ["max_seq_length", "do_lower_case"] - transformer_module.save_in_root = True - if hasattr(model, "config"): - model.config.tokenizer_class = tokenizer.__class__.__name__ - modules[name] = transformer_module - - # load other modules else: + # load other modules (Pooling, Normalize, etc.) module_path = module_config["path"] if os.path.isdir(model_name): load_path = os.path.join(model_name, module_path) else: - # still looking try: - load_path = load_dir_path( - model_name, module_path, token = token - ) - except: - print( - f"Unsloth Warning: Could not download module {module_path} for {class_ref}. Skipping." - ) + load_path = load_dir_path(model_name, module_path, token=token) + except Exception as e: + print(f"Unsloth Warning: Could not download module {module_path}: {e}") continue module_class = import_from_string(class_ref) - # load module try: module = module_class.load(load_path) modules[name] = module except Exception as e: - print( - f"Unsloth Warning: Failed to load module {name} ({class_ref}) from {load_path}: {e}" - ) + print(f"Unsloth Warning: Failed to load module {name} ({class_ref}): {e}") + + return modules, False + + # fallback if no modules.json (non sentence-transformers models) + print("Unsloth: No modules.json found, falling back to [Transformer, Pooling, Normalize]") + + transformer_module = FastSentenceTransformer._create_transformer_module( + model_name, model, tokenizer, max_seq_length, trust_remote_code + ) + modules["0"] = transformer_module - else: - # fallback if no modules.json, is this necessary? - print( - "Unsloth: No modules.json found, falling back to [Transformer, Pooling, Normalize]" - ) - transformer_module = Transformer( - model_name, - max_seq_length = max_seq_length, - model_args = {"trust_remote_code": trust_remote_code}, - config_args = {"trust_remote_code": trust_remote_code}, - ) - transformer_module.auto_model = model - transformer_module.tokenizer = tokenizer + hidden_size = getattr(model.config, "hidden_size", 768) - # move tokenizer do_lower_case to transformer module - transformer_module.do_lower_case = getattr( - tokenizer, "do_lower_case", False - ) - model_forward_params = list(inspect.signature(model.forward).parameters) - transformer_module.model_forward_params = set(model_forward_params) | { - "input_ids", - "attention_mask", - "token_type_ids", - "inputs_embeds", - } + if pooling_mode == "mean": + pooling_mode = FastSentenceTransformer._read_pooling_mode(model_name, token) - if max_seq_length is None: - if hasattr(model, "config") and hasattr( - model.config, "max_position_embeddings" - ): - max_seq_length = model.config.max_position_embeddings - elif hasattr(tokenizer, "model_max_length"): - max_seq_length = tokenizer.model_max_length - else: - max_seq_length = 512 - transformer_module.max_seq_length = max_seq_length - transformer_module.config_keys = ["max_seq_length", "do_lower_case"] - transformer_module.save_in_root = True - # add tokenizer class to config for sentence-transformers - if hasattr(model, "config"): - model.config.tokenizer_class = tokenizer.__class__.__name__ - - modules["0"] = transformer_module - - hidden_size = ( - model.config.hidden_size - if hasattr(model.config, "hidden_size") - else 768 - ) - - if pooling_mode == "mean": - pooling_mode = FastSentenceTransformer._read_pooling_mode( - model_name, token - ) - - pooling_module = Pooling( - word_embedding_dimension = hidden_size, - pooling_mode = pooling_mode, - ) - # end of fallback - modules["1"] = pooling_module - modules["2"] = Normalize() - return modules + modules["1"] = Pooling(word_embedding_dimension=hidden_size, pooling_mode=pooling_mode) + modules["2"] = Normalize() + + return modules, True @staticmethod def from_pretrained( @@ -389,22 +543,38 @@ class FastSentenceTransformer(FastModel): qat_scheme = None, unsloth_tiled_mlp = False, pooling_mode = "mean", + for_inference = False, **kwargs, ): try: from sentence_transformers import SentenceTransformer from sentence_transformers.models import Transformer, Pooling, Normalize - from transformers import AutoModel except ImportError: raise ImportError( "Unsloth: To use `FastSentenceTransformer`, you must install `sentence-transformers`.\n" "Run `pip install sentence-transformers` to install it." ) + # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues + if for_inference: + st_model = SentenceTransformer(model_name, device=device_map, trust_remote_code=trust_remote_code) + return st_model + if "auto_model" not in kwargs: kwargs["auto_model"] = AutoModel - if "add_pooling_layer" not in kwargs: + transformers4 = Version(transformers.__version__).major < 5 + is_distilbert = "distilbert" in model_name.lower() + if not is_distilbert: + try: + # this becomes necessary after merging a trained model + config = AutoConfig.from_pretrained(model_name, token=token) + if getattr(config, "model_type", "") == "distilbert": + is_distilbert = True + except: + pass + + if "add_pooling_layer" not in kwargs and not is_distilbert: kwargs["add_pooling_layer"] = False # forces fp8 to be False since it's not supported @@ -417,8 +587,27 @@ class FastSentenceTransformer(FastModel): old_environ = os.environ.get("UNSLOTH_WARN_UNINITIALIZED", "1") os.environ["UNSLOTH_WARN_UNINITIALIZED"] = "0" - if "mpnet" in model_name.lower(): - FastSentenceTransformer._patch_mpnet() + if is_distilbert and transformers4: + FastSentenceTransformer._patch_distilbert_v4() + elif is_distilbert: + FastSentenceTransformer._patch_distilbert_v5() + elif (mpnet := "mpnet" in model_name.lower()) and transformers4: + FastSentenceTransformer._patch_mpnet_v4() + elif mpnet: + FastSentenceTransformer._patch_mpnet_v5() + + # check if modules.json exists - if not, force 16-bit training + # why? because i have to implement saving myself for these models, and i don't feel like adding dequantization + # to the save_pretrained_merged for a model that really should be trained in 16-bit anyway + has_modules_json = FastSentenceTransformer._module_path(model_name, token) is not None + + if not has_modules_json and load_in_4bit: + print( + "\033[1;33mUnsloth: No modules.json found. This is not a sentence-transformers model.\n" + "Forcing 16-bit loading to simplify merged model saving.\033[0m" + ) + load_in_4bit = False + load_in_16bit = True try: model, tokenizer = FastModel.from_pretrained( @@ -453,8 +642,7 @@ class FastSentenceTransformer(FastModel): # try to load modules, otherwise fallback to old hard-coded modules from sentence_transformers import SentenceTransformer - - modules = FastSentenceTransformer._load_modules( + modules, no_modules = FastSentenceTransformer._load_modules( model_name, token, model, @@ -465,6 +653,7 @@ class FastSentenceTransformer(FastModel): ) st_model = SentenceTransformer(modules = modules, device = device_map) + st_model.no_modules = no_modules def _save_pretrained_merged(self, save_directory, **kwargs): # sentence-transformers config and modules only get saved if we call save_pretrained @@ -477,11 +666,28 @@ class FastSentenceTransformer(FastModel): except: pass - # save merged weights tokenizer = kwargs.pop("tokenizer", self.tokenizer) - self[0].auto_model.save_pretrained_merged( - save_directory, tokenizer = tokenizer, **kwargs - ) + if self.no_modules: + # fallback for non-sentence-transformers models + print("Unsloth: No modules detected. Using standard merge_and_unload for saving...") + safe_kwargs = kwargs.copy() + # filter out Unsloth-specific args that are not in huggingface's save_pretrained + unsloth_args = [ + "save_method", + "temporary_location", + "maximum_memory_usage" + ] + for k in unsloth_args: + safe_kwargs.pop(k, None) + + merged_model = self[0].auto_model.merge_and_unload() + merged_model.save_pretrained(save_directory, **safe_kwargs) + if tokenizer is not None: + tokenizer.save_pretrained(save_directory) + else: + self[0].auto_model.save_pretrained_merged( + save_directory, tokenizer = tokenizer, **kwargs + ) st_model.save_pretrained_merged = types.MethodType( _save_pretrained_merged, st_model