diff --git a/pyproject.toml b/pyproject.toml index 4f9c308b32..c7f67acfdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ triton = [ ] huggingface = [ - "unsloth_zoo>=2025.9.9", + "unsloth_zoo>=2025.9.10", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,<=4.55.4", @@ -453,7 +453,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3", ] colab-new = [ - "unsloth_zoo>=2025.9.9", + "unsloth_zoo>=2025.9.10", "packaging", "tyro", "transformers>=4.51.3,!=4.47.0,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,<=4.55.4", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 60abcea702..ef8fc0fba6 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2025.9.7" +__version__ = "2025.9.8" __all__ = [ "SUPPORTS_BFLOAT16", @@ -137,6 +137,7 @@ for temporary_patch in TEMPORARY_PATCHES: # ============================================= # Disable some warnings which can get annoying warnings.filterwarnings(action = "ignore", category = UserWarning, module = "torch") +warnings.filterwarnings(action = "ignore", category = FutureWarning, module = "torch") warnings.filterwarnings(action = "ignore", category = UserWarning, module = "huggingface_hub") warnings.filterwarnings(action = "ignore", category = FutureWarning, module = "huggingface_hub") warnings.filterwarnings(action = "ignore", category = UserWarning, module = "trl") diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6326f519f1..1b22542514 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2170,6 +2170,9 @@ class FastLlamaModel: m = m.model pass m.max_seq_length = max_seq_length + # Save to modules as well + for module in model.modules(): + module.max_seq_length = max_seq_length # We check the tokenizer first for errors if fix_tokenizer: @@ -2228,6 +2231,11 @@ class FastLlamaModel: # Add for_inference and for_training model.for_training = functools.partial(FastLlamaModel.for_training, model) model.for_inference = functools.partial(FastLlamaModel.for_inference, model) + m = model + while hasattr(m, "model"): + m.for_training = functools.partial(FastBaseModel.for_training, m) + m.for_inference = functools.partial(FastBaseModel.for_inference, m) + m = m.model # Patch generate is_classification = "Classification" in str(type(model)) @@ -2236,6 +2244,13 @@ class FastLlamaModel: unsloth_fast_generate.__doc__ = model._old_generate.__doc__ model.generate = types.MethodType(unsloth_fast_generate, model) pass + # Set weight[padding_idx] = 0 + with torch.no_grad(): + for name, module in model.named_modules(): + if type(module) is torch.nn.Embedding: + if getattr(module, "weight", None) is not None and getattr(module, "padding_idx", None) is not None: + if module.padding_idx < module.weight.shape[0]: + module.weight[module.padding_idx] = 0 return model, tokenizer pass @@ -2704,6 +2719,11 @@ class FastLlamaModel: # Add for_inference and for_training model.for_training = functools.partial(FastLlamaModel.for_training, model) model.for_inference = functools.partial(FastLlamaModel.for_inference, model) + m = model + while hasattr(m, "model"): + m.for_training = functools.partial(FastBaseModel.for_training, m) + m.for_inference = functools.partial(FastBaseModel.for_inference, m) + m = m.model return model pass @@ -2892,6 +2912,9 @@ class FastLlamaModel: internal_model = internal_model.model pass internal_model.max_seq_length = max_seq_length + # Save to modules as well + for module in model.modules(): + module.max_seq_length = max_seq_length # Patch tokenizer to pad to the right internal_model = model @@ -2916,6 +2939,11 @@ class FastLlamaModel: # Add for_inference and for_training model.for_training = functools.partial(FastLlamaModel.for_training, model) model.for_inference = functools.partial(FastLlamaModel.for_inference, model) + m = model + while hasattr(m, "model"): + m.for_training = functools.partial(FastBaseModel.for_training, m) + m.for_inference = functools.partial(FastBaseModel.for_inference, m) + m = m.model return model pass diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 3eb80fc0dd..98396bb754 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -82,20 +82,37 @@ from ._utils import ( ) global FORCE_FLOAT32 +# Forces float32 precision since float16 goes to infinity FORCE_FLOAT32 = [ - "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3,", # Add comma bc gemma3 will match gemma3n "gemma3n", "gpt_oss", ] +global DISABLE_COMPILE_MODEL_NAMES +# Must be alphabetically sorted for each entry +DISABLE_COMPILE_MODEL_NAMES = [ + "aya_vision", + "modernbert", + "granite,llava_next", # Granite-vision 3 +] + +global DISABLE_SDPA_MODEL_NAMES +# Disables some SDPA modules since it's wrong +DISABLE_SDPA_MODEL_NAMES = [ + "gemma3,", # Add comma bc gemma3 will match gemma3n +] + + class FastLanguageModel(FastLlamaModel): @staticmethod def from_pretrained( model_name = "unsloth/Llama-3.2-1B-Instruct", max_seq_length = 2048, dtype = None, - load_in_4bit = True, - load_in_8bit = False, + load_in_4bit = True, # 4bit QLoRA + load_in_8bit = False, # 8bit LoRA + load_in_16bit = False, # 16bit LoRA full_finetuning = False, token = None, device_map = "sequential", @@ -106,6 +123,7 @@ class FastLanguageModel(FastLlamaModel): resize_model_vocab = None, revision = None, use_exact_model_name = False, + offload_embedding = False, fast_inference = False, # uses vLLM gpu_memory_utilization = 0.5, @@ -131,6 +149,7 @@ class FastLanguageModel(FastLlamaModel): dtype = dtype, load_in_4bit = load_in_4bit, load_in_8bit = load_in_8bit, + load_in_16bit = load_in_16bit, full_finetuning = full_finetuning, token = token, device_map = device_map, @@ -143,6 +162,7 @@ class FastLanguageModel(FastLlamaModel): return_logits = False, # Return logits fullgraph = True, # No graph breaks use_exact_model_name = use_exact_model_name, + offload_embedding = offload_embedding, # Pass vLLM/inference parameters fast_inference = fast_inference, @@ -213,16 +233,27 @@ class FastLanguageModel(FastLlamaModel): peft_error = str(error) is_peft = False pass - model_types = get_transformers_model_type(peft_config or model_config) + + # Old transformers versions check + both_exist = (is_model and is_peft) and not SUPPORTS_LLAMA32 + + # Error out if both LoRA and normal model config exists. + if both_exist: + raise RuntimeError( + "Unsloth: Your repo has a LoRA adapter and a base model.\n"\ + "You have 2 files `config.json` and `adapter_config.json`.\n"\ + "We must only allow one config file.\n"\ + "Please separate the LoRA and base models to 2 repos." + ) + model_types = get_transformers_model_type( + peft_config if peft_config is not None else model_config + ) if len(model_types) == 1: model_type = model_types[0] else: # Leave as tuple if more than one arch model_type = model_types - # Old transformers versions check - both_exist = (is_model and is_peft) and not SUPPORTS_LLAMA32 - # New transformers need to check manually. if SUPPORTS_LLAMA32: # Check if folder exists locally @@ -240,17 +271,8 @@ class FastLanguageModel(FastLlamaModel): pass pass - # Error out if both LoRA and normal model config exists. - if both_exist: - raise RuntimeError( - "Unsloth: Your repo has a LoRA adapter and a base model.\n"\ - "You have 2 files `config.json` and `adapter_config.json`.\n"\ - "We must only allow one config file.\n"\ - "Please separate the LoRA and base models to 2 repos." - ) - - elif not is_model and not is_peft: - error = autoconfig_error or peft_error + if not is_model and not is_peft: + error = autoconfig_error if autoconfig_error is not None else peft_error # Old transformers version if "rope_scaling" in error.lower() and not SUPPORTS_LLAMA31: raise ImportError( @@ -368,6 +390,7 @@ class FastLanguageModel(FastLlamaModel): dtype = dtype, load_in_4bit = load_in_4bit, load_in_8bit = load_in_8bit, + load_in_16bit = load_in_16bit, full_finetuning = full_finetuning, token = token, device_map = device_map, @@ -380,6 +403,7 @@ class FastLanguageModel(FastLlamaModel): return_logits = False, # Return logits fullgraph = True, # No graph breaks use_exact_model_name = use_exact_model_name, + offload_embedding = offload_embedding, # Pass vLLM/inference parameters fast_inference = fast_inference, @@ -498,13 +522,6 @@ except: from transformers import AutoModelForVision2Seq pass -# Must be alphabetically sorted for each entry -DISABLE_COMPILE_MODEL_NAMES = [ - "aya_vision", - "modernbert", - "granite,llava_next", # Granite-vision 3 -] - class FastModel(FastBaseModel): @staticmethod @@ -512,8 +529,9 @@ class FastModel(FastBaseModel): model_name = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit", max_seq_length = 2048, dtype = None, - load_in_4bit = True, - load_in_8bit = False, + load_in_4bit = True, # 4bit QLoRA + load_in_8bit = False, # 8bit LoRA + load_in_16bit = False, # 16bit LoRA full_finetuning = False, token = None, device_map = "sequential", @@ -530,6 +548,7 @@ class FastModel(FastBaseModel): whisper_language = None, whisper_task = None, unsloth_force_compile = False, + offload_embedding = False, # Add the missing vLLM/inference parameters fast_inference = False, # uses vLLM @@ -565,15 +584,17 @@ class FastModel(FastBaseModel): if full_finetuning and (load_in_4bit or load_in_8bit): print("Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA.") - load_in_4bit = False - load_in_8bit = False + load_in_4bit = False + load_in_8bit = False + load_in_16bit = False pass - if load_in_4bit and load_in_8bit: + if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: raise RuntimeError( - "Unsloth: Can only load in 4bit or 8bit, not both!\n"\ + "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n"\ "Also, we by default set `load_in_4bit = True`.\n"\ - "If you want 8bit finetuning, set both `load_in_4bit = False` and `load_in_8bit = True`" + "If you want 8bit finetuning, set both `load_in_4bit = False` and `load_in_8bit = True`\n"\ + "If you want 16bit LoRA finetuning, set `load_in_16bit = True`" ) pass @@ -626,8 +647,20 @@ class FastModel(FastBaseModel): peft_error = str(error) is_peft = False pass - model_types = get_transformers_model_type(peft_config or model_config) - model_types_all = ",".join(model_types) + # Old transformers versions check + both_exist = (is_model and is_peft) and not SUPPORTS_LLAMA32 + # Error out if both LoRA and normal model config exists. + if both_exist: + raise RuntimeError( + "Unsloth: Your repo has a LoRA adapter and a base model.\n"\ + "You have 2 files `config.json` and `adapter_config.json`.\n"\ + "We must only allow one config file.\n"\ + "Please separate the LoRA and base models to 2 repos." + ) + model_types = get_transformers_model_type( + peft_config if peft_config is not None else model_config + ) + model_types_all = ",".join(model_types) + "," # Check versions lowered_model_name = model_name.lower() @@ -641,22 +674,24 @@ class FastModel(FastBaseModel): # Qwen 2.5 elif "qwen2_5" in model_types_all and transformers_version < Version("4.49.0"): raise RuntimeError("Unsloth: Qwen 2.5 only works on transformers >= 4.49.0." + LATEST) + # Gemma 3N must be before Gemma 3 + elif "gemma3n" in model_types_all: + if transformers_version < Version("4.53.0"): + raise RuntimeError("Unsloth: Gemma 3N only works on transformers >= 4.53.0" + LATEST) + os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" + os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \ + "float16;torch.float16;torch.float16;"\ + "if name.endswith('norm'): "\ + "module._pre_set_compute_dtype = torch.float32\n"\ + ";"\ + "from unsloth_zoo.temporary_patches.gemma3n import patch_Gemma3nConv_Embed_forwards; patch_Gemma3nConv_Embed_forwards()" + # Set norms to float32 since anyways they get upcasted to float32 + # common in both gemma-3 and gemma-3n + os.environ["UNSLOTH_HIGH_PRECISION_LAYERNORM"] = "1" # Gemma 3 elif "gemma3" in model_types_all: - if "gemma3n" in model_types_all: - if transformers_version < Version("4.53.0"): - raise RuntimeError("Unsloth: Gemma 3N only works on transformers >= 4.53.0" + LATEST) - os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" - os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \ - "float16;torch.float16;torch.float16;"\ - "if name.endswith('norm'): "\ - "module._pre_set_compute_dtype = torch.float32\n"\ - ";"\ - "from unsloth_zoo.temporary_patches.gemma3n import patch_Gemma3nConv_Embed_forwards; patch_Gemma3nConv_Embed_forwards()" - else: - if transformers_version < Version("4.50.0.dev0"): - raise RuntimeError("Unsloth: Gemma 3 only works on transformers >= 4.50.0." + NIGHTLY) - + if transformers_version < Version("4.50.0.dev0"): + raise RuntimeError("Unsloth: Gemma 3 only works on transformers >= 4.50.0." + NIGHTLY) # Set norms to float32 since anyways they get upcasted to float32 # common in both gemma-3 and gemma-3n os.environ["UNSLOTH_HIGH_PRECISION_LAYERNORM"] = "1" @@ -665,7 +700,7 @@ class FastModel(FastBaseModel): raise RuntimeError("Unsloth: Cohere's Command model only works on transformers >= 4.50.0." + NIGHTLY) # Sesame elif "csm" in model_types_all: - os.environ["UNSLOTH_COMPILE_DISABLE"] = "1" # Inference is too slow + os.environ["UNSLOTH_COMPILE_DISABLE"] = "partial" # Inference is too slow os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" # Sesame fails os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \ "all;torch.float32;torch.float16;"\ @@ -720,7 +755,7 @@ class FastModel(FastBaseModel): else: for check_model_name in DISABLE_COMPILE_MODEL_NAMES: if check_model_name in lowered_model_name: - os.environ["UNSLOTH_COMPILE_DISABLE"] = "1" + os.environ["UNSLOTH_COMPILE_DISABLE"] = "partial" os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" if transformers_version < Version("4.50.0.dev0"): raise RuntimeError(f"Unsloth: {check_model_name} only works on transformers >= 4.50.0." + NIGHTLY) @@ -732,9 +767,6 @@ class FastModel(FastBaseModel): os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" pass - # Old transformers versions check - both_exist = (is_model and is_peft) and not SUPPORTS_LLAMA32 - # New transformers need to check manually. if SUPPORTS_LLAMA32: # Check if folder exists locally @@ -751,17 +783,8 @@ class FastModel(FastBaseModel): pass pass - # Error out if both LoRA and normal model config exists. - if both_exist: - raise RuntimeError( - "Unsloth: Your repo has a LoRA adapter and a base model.\n"\ - "You have 2 files `config.json` and `adapter_config.json`.\n"\ - "We must only allow one config file.\n"\ - "Please separate the LoRA and base models to 2 repos." - ) - - elif not is_model and not is_peft: - error = autoconfig_error or peft_error + if not is_model and not is_peft: + error = autoconfig_error if autoconfig_error is not None else peft_error # Old transformers version if "rope_scaling" in error.lower() and not SUPPORTS_LLAMA31: raise ImportError( @@ -811,7 +834,7 @@ class FastModel(FastBaseModel): for disable_name in FORCE_FLOAT32: # add comma to model_types_all matching in case of exact match for end if (disable_name.lower() == model_type_arch.lower().replace("-", "").replace("_", "") or \ - disable_name.lower() in f'{model_types_all},') and \ + disable_name.lower() in model_types_all) and \ ((dtype == torch.float16) or not SUPPORTS_BFLOAT16): os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" dtype = torch.bfloat16 # Change to bfloat16 loading @@ -855,9 +878,10 @@ class FastModel(FastBaseModel): unsloth_force_compile = unsloth_force_compile, ) pass - # Fix SDPA - if "gemma3n" in model_types_all: - supports_sdpa = False + # Fix SDPA issues + for model_type in DISABLE_SDPA_MODEL_NAMES: + if model_type in model_types_all: + supports_sdpa = False pass # Check if this is local model since the tokenizer gets overwritten @@ -884,6 +908,7 @@ class FastModel(FastBaseModel): dtype = _get_dtype(dtype), load_in_4bit = load_in_4bit, load_in_8bit = load_in_8bit, + load_in_16bit = load_in_16bit, full_finetuning = full_finetuning, token = token, device_map = device_map, @@ -896,6 +921,8 @@ class FastModel(FastBaseModel): supports_sdpa = supports_sdpa, whisper_language = whisper_language, whisper_task = whisper_task, + auto_config = model_config, + offload_embedding = offload_embedding, # Pass vLLM/inference parameters fast_inference = fast_inference, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 3d5f6d084b..2b6293993d 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -116,6 +116,23 @@ from torch.nn import functional as F from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling from transformers.training_args import ParallelMode +# Wrap trainer with padding to right and enable training mode +import functools +from types import MethodType +def prepare_for_training_mode(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + # Enable training mode + if hasattr(self, 'model') and hasattr(self.model, "for_training"): + self.model.for_training() + output = f(self, *args, **kwargs) + # Return inference mode + if hasattr(self, 'model') and hasattr(self.model, "for_inference"): + self.model.for_inference() + return output + return wrapper +pass + torch_compile_options = {{ "epilogue_fusion" : True, "max_autotune" : False, @@ -174,7 +191,11 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: if getattr(args, "_n_gpu", 1) != 1: args._n_gpu = 1 + if "model" in locals() and hasattr(model, "for_training"): + model.for_training() super().__init__({RLTrainer_call_args}{RLTrainer_kwargs}) + if "model" in locals() and hasattr(model, "for_inference"): + model.for_inference() {RLTrainer_post} pass ''' @@ -460,7 +481,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Add accelerator scaler to model if "model" in call_args: - neftune_check = \ + accelerator_check = \ "if hasattr(self, 'accelerator'):\n"\ " scaler = self.accelerator.scaler\n"\ " current_model = model\n"\ @@ -469,7 +490,16 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): " current_model = current_model.model\n"\ " current_model.accelerator_scaler = scaler\n"\ "pass\n" - RLTrainer_post += neftune_check + RLTrainer_post += accelerator_check + pass + + # Add enabling and disabling training modes + if "model" in call_args: + training_check = \ + "if hasattr(self, 'train'):\n"\ + " self.train = MethodType(prepare_for_training_mode(self.__class__.train), self)\n"\ + "pass\n" + RLTrainer_post += training_check pass # Edit optional metrics @@ -933,6 +963,19 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import source = edit_function(function, source) pass + """ + import torch + X = torch.ones((2, 2048, 201088), dtype = torch.bfloat16, device = "cuda") + X[torch.randperm(2, dtype = torch.int64, device = X.device)] + + will error out in torch 2.8 AcceleratorError: CUDA error: invalid configuration argument + """ + source = re.sub( + r"(\n[\s]{4,})generation_batch = shuffle_sequence_dict\(generation_batch\)\n", + r"\n\1try: generation_batch = shuffle_sequence_dict(generation_batch)\n\1except: pass\n", + source, + ) + # llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model source = re.sub( r"(\n[\s]{4,}).+?model_executor\.driver_worker.+?\n", diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 3f5ae816ea..2cd95b0377 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -83,13 +83,18 @@ NUM_LOGITS_TO_KEEP = dict() VLLM_SUPPORTED_VLM = [ "qwen2_5_vl", "gemma3", + "mistral3", ] VLLM_NON_LORA_VLM = [ - "mllama" + "mllama", +] +PRE_COMPILE_INFERENCE = [ + "gpt_oss", ] from transformers import GenerationConfig, CompileConfig, HybridCache, AutoConfig, PretrainedConfig HAS_TORCH_DTYPE = "torch_dtype" in PretrainedConfig.__doc__ + from transformers import GenerationConfig, CompileConfig, HybridCache _compile_config = CompileConfig( @@ -217,8 +222,11 @@ def unsloth_base_fast_generate( if getattr(self, "_supports_static_cache", getattr(self, "_can_compile_fullgraph", True)): if os.environ.get("UNSLOTH_DISABLE_STATIC_GENERATION", "0") == "0": cache_implementation = "static" - else: + elif Version(transformers_version) < Version("4.56.0.dev0"): cache_implementation = None + else: + # Should work in latest transformers! + cache_implementation = "static" else: cache_implementation = None if cache_implementation is not None: @@ -242,10 +250,33 @@ def unsloth_base_fast_generate( kwargs["compile_config"] = _compile_config pass + # Delete cached Flex Attention masks to reset inference + for name, module in self.named_modules(): + if hasattr(module, "_flex_attention_cache"): + try: del module._flex_attention_cache + except: pass + # Solves AttributeError: 'SlidingWindowLayer' object has no attribute 'max_batch_size' + if hasattr(module, "_cache") and "cache_utils" in str(module._cache.__class__): + try: del module._cache + except: pass + pass + + # DO INFERENCE with torch.inference_mode(), autocaster: output = self._old_generate(*args, **kwargs) - FastBaseModel.for_training(self) + # Delete cached Flex Attention masks to reset inference + for name, module in self.named_modules(): + if hasattr(module, "_flex_attention_cache"): + try: del module._flex_attention_cache + except: pass + # Solves AttributeError: 'SlidingWindowLayer' object has no attribute 'max_batch_size' + if hasattr(module, "_cache") and "cache_utils" in str(module._cache.__class__): + try: del module._cache + except: pass + pass + + # FastBaseModel.for_training(self) return output pass @@ -258,6 +289,7 @@ class FastBaseModel: dtype = None, load_in_4bit = True, load_in_8bit = False, + load_in_16bit = False, full_finetuning = False, token = None, device_map = "sequential", @@ -269,7 +301,10 @@ class FastBaseModel: supports_sdpa = True, whisper_language = None, whisper_task = None, - fast_inference = False, + auto_config = None, + offload_embedding = False, + # vLLM parameters + fast_inference = False, gpu_memory_utilization = 0.5, float8_kv_cache = False, random_state = 3407, @@ -421,19 +456,21 @@ class FastBaseModel: if not ("attn_implementation" in kwargs): kwargs["attn_implementation"] = "sdpa" if not supports_sdpa: - print(f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager.") + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") == "0": + print(f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager.") del kwargs["attn_implementation"] pass bnb_config = None if full_finetuning and (load_in_4bit or load_in_8bit): print("Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA.") - load_in_4bit = False - load_in_8bit = False + load_in_4bit = False + load_in_8bit = False + load_in_16bit = False pass - if load_in_4bit and load_in_8bit: - raise RuntimeError("Unsloth: Can only load in 4bit or 8bit, not both!") + if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: + raise RuntimeError("Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!") if load_in_4bit: bnb_config = BitsAndBytesConfig( load_in_4bit = True, @@ -447,6 +484,8 @@ class FastBaseModel: load_in_8bit = True, llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy(), ) + elif load_in_16bit: + bnb_config = None elif not load_in_4bit and not load_in_8bit and not full_finetuning: print("Unsloth: QLoRA and full finetuning all not selected. Switching to 16bit LoRA.") pass @@ -468,10 +507,28 @@ class FastBaseModel: # Cannot be None, since HF now checks for the config if load_in_4bit: # Ignore load_in_4bit / load_in_8bit for MXFP4 - best to get config file - if "gpt-oss" in model_name.lower(): + if "gpt-oss-20b" in model_name.lower() or "gpt-oss-120b" in model_name.lower(): pass else: kwargs["quantization_config"] = bnb_config + else: + if auto_config is None: + auto_config = AutoConfig.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + ) + if hasattr(auto_config, "quantization_config"): + from transformers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING + quantization_config = auto_config.quantization_config + quantizer = AUTO_QUANTIZATION_CONFIG_MAPPING[quantization_config["quant_method"]] + quantizer_kwargs = {} + # We cannot dequantize since gpt-oss-20b MXFP4 will now be gpt-oss-20b-BF16 + # if "dequantize" in inspect.signature(quantizer).parameters: + # quantizer_kwargs["dequantize"] = True + quantization_config = quantizer.from_dict(quantization_config, **quantizer_kwargs) + kwargs["quantization_config"] = quantization_config + pass pass # Check if using forced float32 - we load it in bfloat16, then cast to float16! @@ -495,6 +552,26 @@ class FastBaseModel: if hasattr(model, 'generate'): model.fast_generate = model.generate model.fast_generate_batches = error_out_no_vllm + if offload_embedding: + embed_tokens = model.get_input_embeddings() + nbytes = embed_tokens.weight.numel() * embed_tokens.weight.itemsize + ngb = round(nbytes / 1024 / 1024 / 1024, 2) + print(f"Unsloth: Offloading embeddings to RAM to save {ngb} GB.") + embed_tokens.to("cpu") + + # Add hooks to move inputs to CPU and back to CUDA + # [TODO] Doesn't seem to work! + # def pre_hook(module, args): + # args[0]._old_device = args[0].device + # return (args[0].to("cpu", non_blocking = True)) + # def post_hook(module, args, output): + # old_device = getattr(args[0], "_old_device", "cuda") + # return output.to(old_device, non_blocking = True) + # embed_tokens.register_forward_pre_hook(pre_hook, prepend = True) + # embed_tokens.register_forward_hook (post_hook, prepend = True) + # Must free GPU memory otherwise will not free! + torch.cuda.empty_cache() + gc.collect() else: from unsloth_zoo.vllm_utils import ( load_vllm, @@ -573,7 +650,7 @@ class FastBaseModel: if (whisper_language and whisper_task) or auto_model.__name__.endswith("ForConditionalGeneration"): tokenizer = auto_processor.from_pretrained( tokenizer_name, - padding_side = "right", + padding_side = "left", token = token, language = whisper_language, task = whisper_task, @@ -582,19 +659,19 @@ class FastBaseModel: try: tokenizer = auto_processor.from_pretrained( tokenizer_name, - padding_side = "right", + padding_side = "left", token = token, ) except: tokenizer = get_auto_processor( tokenizer_name, - padding_side = "right", + padding_side = "left", token = token, ) if hasattr(tokenizer, "tokenizer"): __tokenizer = tokenizer.tokenizer # Add padding side as well - __tokenizer.padding_side = "right" + __tokenizer.padding_side = "left" # Check bos, eos, pad tokens if hasattr(__tokenizer, "bos_token"): tokenizer.bos_token = __tokenizer.bos_token @@ -642,6 +719,9 @@ class FastBaseModel: m = m.model pass m.max_seq_length = max_seq_length + # Save to modules as well + for module in model.modules(): + module.max_seq_length = max_seq_length m._saved_temp_tokenizer = tokenizer # Also set is_loaded_in_8bit to disable incorrect DDP m.is_loaded_in_8bit = True if not full_finetuning else False @@ -659,6 +739,8 @@ class FastBaseModel: model, use_gradient_checkpointing = use_gradient_checkpointing, trust_remote_code = trust_remote_code, + model_type = model_type_arch, + tokenizer = tokenizer, ) # Clear deleted GPU items for _ in range(3): @@ -671,6 +753,51 @@ class FastBaseModel: return model, tokenizer pass + @staticmethod + def pre_compile_for_inference(model_type, model, tokenizer): + """ + We need to invoke torch.compile to save VRAM usage and make it faster downstream. + Sometimes torch.compile can use 3GB weirdly on large batches, then it goes down to <1GB. + So we invoke torch.compile on short batches to reduce VRAM usage. + """ + if model_type is None or model is None or tokenizer is None: return + if str(model_type).lower() not in PRE_COMPILE_INFERENCE: return + if getattr(tokenizer, "chat_template", None) is None: return + # Check if already compiled and exit + for module in model.modules(): + if hasattr(module, "_pre_compiled_for_inference"): return + pass + print(f"🦥 Unsloth: Pre compiling {model_type.title()} model for faster inference - this might take 3 minutes or so!") + print("========= Pre compiling model for faster inference. Please be patient thank you! =========") + # Do single inference + messages = [ + [ + {"role": "user", "content": f"What is 1+1 equal to?"}, + ], + ]*1 + inputs = tokenizer.apply_chat_template( + messages, + add_generation_prompt = True, + return_tensors = "pt", + return_dict = True, + ).to(model.device) + _ = model.generate(**inputs, max_new_tokens = 1) + # Do batched inference + messages = [ + [ + {"role": "user", "content": f"1+1"}, + ], + ]*4 + inputs = tokenizer.apply_chat_template( + messages, + add_generation_prompt = True, + return_tensors = "pt", + return_dict = True, + ).to(model.device) + _ = model.generate(**inputs, max_new_tokens = 2) + # Set we already pre compiled + model._pre_compiled_for_inference = True + pass @staticmethod def get_peft_model( @@ -777,6 +904,9 @@ class FastBaseModel: trust_remote_code = getattr(model, "_unsloth_trust_remote_code", False) model = FastBaseModel.post_patch_model(model, use_gradient_checkpointing, trust_remote_code = trust_remote_code) model.max_seq_length = max_seq_length + # Save to modules as well + for module in model.modules(): + module.max_seq_length = max_seq_length # Clear deleted GPU items for _ in range(3): gc.collect() @@ -791,6 +921,11 @@ class FastBaseModel: # Add for_inference and for_training model.for_training = functools.partial(FastBaseModel.for_training, model) model.for_inference = functools.partial(FastBaseModel.for_inference, model) + m = model + while hasattr(m, "model"): + m.for_training = functools.partial(FastBaseModel.for_training, m) + m.for_inference = functools.partial(FastBaseModel.for_inference, m) + m = m.model return model pass @@ -800,6 +935,8 @@ class FastBaseModel: model, use_gradient_checkpointing = True, trust_remote_code = False, + model_type = None, + tokenizer = None, ): full_finetuning = os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1" @@ -826,12 +963,12 @@ class FastBaseModel: pass patch_saving_functions(model, vision = True) - # Patch tokenizer to pad to the right + # Patch tokenizer to pad to the left m = model while hasattr(m, "model"): if hasattr(m, "_saved_temp_tokenizer"): if hasattr(m._saved_temp_tokenizer, "tokenizer"): - m._saved_temp_tokenizer.tokenizer.padding_side = "right" + m._saved_temp_tokenizer.tokenizer.padding_side = "left" pass # Also set is_loaded_in_8bit to disable incorrect DDP m.is_loaded_in_8bit = True if not full_finetuning else False @@ -839,7 +976,7 @@ class FastBaseModel: pass if hasattr(m, "_saved_temp_tokenizer"): if hasattr(m._saved_temp_tokenizer, "tokenizer"): - m._saved_temp_tokenizer.tokenizer.padding_side = "right" + m._saved_temp_tokenizer.tokenizer.padding_side = "left" pass # Also set is_loaded_in_8bit to disable incorrect DDP m.is_loaded_in_8bit = True if not full_finetuning else False @@ -855,6 +992,20 @@ class FastBaseModel: # Add for_inference and for_training model.for_training = functools.partial(FastBaseModel.for_training, model) model.for_inference = functools.partial(FastBaseModel.for_inference, model) + m = model + while hasattr(m, "model"): + m.for_training = functools.partial(FastBaseModel.for_training, m) + m.for_inference = functools.partial(FastBaseModel.for_inference, m) + m = m.model + # Set weight[padding_idx] = 0 + with torch.no_grad(): + for name, module in model.named_modules(): + if type(module) is torch.nn.Embedding: + if getattr(module, "weight", None) is not None and getattr(module, "padding_idx", None) is not None: + if module.padding_idx < module.weight.shape[0]: + module.weight[module.padding_idx] = 0 + # Patch for torch.compiled inference + # FastBaseModel.pre_compile_for_inference(model_type, model, tokenizer) return model pass @@ -922,7 +1073,12 @@ class FastBaseModel: # Pad tokenizer to the left if hasattr(m, "_saved_temp_tokenizer"): m._saved_temp_tokenizer.padding_side = "right" # Set a flag for generation! - if hasattr(m, "_flag_for_generation"): del m._flag_for_generation + if hasattr(m, "_flag_for_generation"): + try: + # Weirdly sometimes cannot succeed so do a try except + del m._flag_for_generation + except: + pass pass m = model while hasattr(m, "model"):