Many bug fixes

pyproject.toml

@@ -24,7 +24,8 @@ classifiers = [
 dependencies = [
     "torch",
     "triton ; platform_system == 'Linux'",
-    "packaging",
+    "triton_windows ; platform_system == 'Windows'",
+    "packaging>=24.1",
     "tyro",
     "transformers>=4.46.1,!=4.47.0",
     "datasets>=2.16.0",

unsloth_zoo/__init__.py

@@ -14,7 +14,7 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-__version__ = "2025.3.12"
+__version__ = "2025.3.13"
 
 from importlib.util import find_spec
 if find_spec("unsloth") is None:

unsloth_zoo/compiler.py

@@ -163,6 +163,9 @@ def get_transformers_model_type(
     config = str(config.to_dict())
     model_types = re.findall(r"'model_type': '([^\s\']{1,})'", config)
     model_types = [x.replace("-", "_").lower() for x in model_types]
+    # Add splitted modules for eg gemma3_text -> gemma3
+    model_types += [x.split("_")[0] for x in model_types]
+    model_types = list(dict().fromkeys(model_types))
 
     from transformers import models
     models = dir(models)
@@ -1213,6 +1216,34 @@ def lora_forward(result, lora_A, lora_B, dropout, x, scaling):
 
 """
 
+COMPILED_LORA_FORWARD_forced_float32 = """
+torch_addmm = torch.addmm
+torch_add   = torch.add
+@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options)
+def lora_forward(result, lora_A, lora_B, dropout, x, scaling):
+    xA = dropout(x.to(torch.float16)) @ lora_A.weight.to(torch.float16).t()
+    # output = result + scaling * xA @ lora_B.weight.t()
+    shape = result.shape
+    output = torch_addmm(
+        result.view(-1, shape[-1]),
+        xA.view(-1, xA.shape[-1]),
+        lora_B.weight.to(torch.float16).t(),
+        alpha = scaling,
+        beta = 1,
+    ).view(shape)
+
+    bias = lora_B.bias
+    if bias is not None:
+        output = torch_add(
+        output,
+        bias.to(torch.float16),
+        alpha = scaling,
+    )
+    return output
+pass
+
+"""
+
 def patch_lora_forwards(torch_compile_options):
     # All Unsloth Zoo code licensed under LGPLv3
     Linear_LoRA_Layers = get_lora_layer_modules()
@@ -1254,26 +1285,37 @@ def patch_lora_forwards(torch_compile_options):
         )
 
         # Check failed upcasting
-        if "torch.is_autocast_enabled()" not in source:
-            source = source.replace(
-                "x = x.to(lora_A.weight.dtype)",
-                "if not torch.is_autocast_enabled(): "\
-                "result, x = "\
-                    "result.to(lora_A.weight.dtype), "\
-                    "x.to(lora_A.weight.dtype)"
-            )
+        replacements = [
+            "x = x.to(lora_A.weight.dtype)",
+            "x = self._cast_input_dtype(x, lora_A.weight.dtype)",
+        ]
+        if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "0":
+            if "torch.is_autocast_enabled()" not in source:
+                new = "if not torch.is_autocast_enabled(): "\
+                    "result, x = "\
+                        "result.to(lora_A.weight.dtype), "\
+                        "x.to(lora_A.weight.dtype)"
+                for replace in replacements:
+                    source = source.replace(replace, new)
+        else:
+            for replace in replacements:
+                source = source.replace(replace, "")
         pass
-
         source = source.replace(
             "self._check_forward_args(x, *args, **kwargs)",
             "",
         )
 
         if hash(source) != old_hash:
             success += 1
+            compiled_lora_forward = \
+                COMPILED_LORA_FORWARD \
+                if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "0" \
+                else COMPILED_LORA_FORWARD_forced_float32
+
             forward = create_new_function(
                 f"{child}_peft_forward",
-                COMPILED_LORA_FORWARD + source,
+                compiled_lora_forward + source,
                 parent,
                 dir(eval(parent)),
                 prepend = \

unsloth_zoo/dataset_utils.py

@@ -492,6 +492,7 @@ def sft_prepare_dataset(
     if max_seq_length == 0: max_seq_length = getattr(args, "max_seq_length", 0)
     if max_seq_length == 0: max_seq_length = getattr(self, "max_seq_length", 0)
     if max_seq_length == 0: max_seq_length = getattr(self, "max_seq", 0)
+    if max_seq_length == 0: raise RuntimeError("Unsloth: max_seq_length is 0! Please specify one!")
     dataset_text_field = getattr(args, "dataset_text_field", "text")
     do_truncation = max_seq_length != 0
     do_formatting_func = False

unsloth_zoo/loss_utils.py

@@ -168,7 +168,7 @@ def fused_linear_cross_entropy(
     reduction = "sum" if num_items_in_batch is not None else "mean"
     if logit_softcapping == 0: logit_softcapping = None
     loss = linear_cross_entropy(
-        hidden_states,
+        hidden_states.to(lm_weight.dtype),
         lm_weight,
         targets      = labels,
         ignore_index = ignore_index,

unsloth_zoo/patching_utils.py

@@ -183,7 +183,13 @@ def patch_torch_compile(debug = False, O3 = False, ignore_errors = True):
 pass
 
 
-def patch_model_and_tokenizer(model, tokenizer, downcast_rope = True, fix_embeddings = True):
+def patch_model_and_tokenizer(
+    model,
+    tokenizer,
+    downcast_rope = True,
+    fix_embeddings = True,
+    do_forced_float32 = False,
+):
     # All Unsloth Zoo code licensed under LGPLv3
     assert(type(downcast_rope) is bool)
     import gc
@@ -221,7 +227,44 @@ def patch_model_and_tokenizer(model, tokenizer, downcast_rope = True, fix_embedd
         correct_dtype = _get_dtype(model.config.torch_dtype)
     except:
         correct_dtype = model.get_input_embeddings().weight.dtype
+    # If we force float32, we first use bfloat16, then downcast to float16
+    if do_forced_float32:
+        correct_dtype = torch.float16
+        for name, module in model.named_modules():
+            if "down_proj" in name or "up_proj" in name or "gate_proj" in name:
+                exec(f"module.to(torch.float16)")
+            if "q_proj" in name or "k_proj" in name or "v_proj" in name or "o_proj" in name:
+                exec(f"module.to(torch.float16)")
+            if "lm_head" in name or "embed_tokens" in name:
+                exec(f"module.to(torch.float16)")
+            if "norm" in name:
+                exec(f"module.to(torch.float32)")
+                assert(module.weight.dtype == torch.float32)
+            torch.cuda.empty_cache()
+        pass
 
+        # Correct torch_dtype
+        def __fix_dtype(config):
+            if not hasattr(config, "to_dict"): return
+            dicts = config.to_dict()
+            for key, value in dicts.items():
+                if key == "torch_dtype":
+                    setattr(config, "torch_dtype", torch.float16)
+                else:
+                    __fix_dtype(getattr(config, key))
+        m = model
+        while hasattr(m, "model"):
+            if hasattr(m, "dtype"):
+                try: setattr(m, "dtype", torch.float16)
+                except: pass
+            if hasattr(m, "config"): __fix_dtype(m.config)
+            m = m.model
+        pass
+        if hasattr(m, "config"): __fix_dtype(m.config)
+        if hasattr(m, "dtype"):
+            try: setattr(m, "dtype", torch.float16)
+            except: pass
+    pass
     # Check all params and patch!
     for name, module in model.named_modules():
         if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)):

unsloth_zoo/peft_utils.py

@@ -208,7 +208,8 @@ def requires_grad_pre_hook(module, input):
                 raise RuntimeError("Unsloth: Failed to make input require gradients!")
                 # print(f"  WARNING: Empty list input to {module.__class__.__name__}!") # 
                 # return
-            input[0].requires_grad_(True)
+            if torch.is_floating_point(input[0]):
+                input[0].requires_grad_(True)
         else:
             raise RuntimeError("Unsloth: Failed to make input require gradients!")
     pass
@@ -247,6 +248,10 @@ def requires_grad_pre_hook(module, input):
             if f"in self.{module_list}:" in forward:
                 final_where = j
                 break
+            elif re.search(r"for [^\s]{3,} in self\." + module_list, forward) is not None:
+                # Might have failed finding self.layers: like self.layers[...]:
+                final_where = j
+                break
             pass
         pass
     pass

unsloth_zoo/rl_replacements.py

@@ -78,9 +78,9 @@ def grpo_compute_loss(old_logits, new_logits, input_ids, mask, beta, advantages)
     n_mask_per_reward = mask.sum(1)
 
     # See https://github.com/huggingface/trl/pull/2881
-    # loss_per_reward = (loss_i * mask).sum(1) / n_mask_per_reward
-    # loss = loss_per_reward.mean()
-    loss = (loss_i * mask).sum() / mask.sum()
+    loss_per_reward = (loss_i * mask).sum(1) / n_mask_per_reward
+    loss = loss_per_reward.mean()
+    # loss = (loss_i * mask).sum() / mask.sum()
 
     # Get metrics as well which are folded
     with torch.inference_mode():

unsloth_zoo/saving_utils.py

@@ -61,6 +61,7 @@
 pass
 from transformers.modeling_utils import PushToHubMixin
 import json
+import os
 from pathlib import Path
 import tempfile
 from peft import PeftModelForCausalLM
@@ -540,7 +541,13 @@ def merge_and_overwrite_lora(
         model_name = model.config._name_or_path
 
     # Find repository's max shard size and total size of everything
-    file_list = HfFileSystem(token = token).ls(model_name, detail = True)
+    try:
+        file_list = HfFileSystem(token = token).ls(model_name, detail = True)
+    except:
+        original_model_id = get_original_model_id(model_name)
+        model_name = original_model_id
+        file_list = HfFileSystem(token = token).ls(model_name, detail = True)
+
     safetensors_list = []
     max_size_in_bytes = 0
     total_size_in_bytes = 0
@@ -909,6 +916,30 @@ def merge_lora_weights(state_dict, name):
     pass
 pass
 
+def get_original_model_id(local_path: str):
+    import json
+    import os
+
+    config_path = os.path.join(local_path, "config.json")
+    if os.path.exists(config_path):
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        # Check for _name_or_path that's not a local path
+        # When we load using AutoConfig, the _name_or_path changed into the local path instead
+        if "_name_or_path" in config:
+            return config["_name_or_path"]
+
+    config_path = os.path.join(local_path, "adapter_config.json")
+    if os.path.exists(config_path):
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        if "base_model_name_or_path" in config:
+            return config["base_model_name_or_path"]
+
+    return None
+
 # Unsloth Zoo - Utilities for Unsloth
 # Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
 #