Merge branch 'main' into nightly

.github/ISSUE_TEMPLATE/bug---issue.md

@@ -15,5 +15,5 @@ assignees: ''
 6. Which trainer? `SFTTrainer`, `GRPOTrainer` etc
 7. **Minimal code to reproduce error Remove Hugging Face token!**
 
-For quick replies, got to https://discord.com/invite/unsloth.
-Have you tried https://docs.unsloth.ai/basics/errors-troubleshooting
+You can also join our Discord: https://discord.com/invite/unsloth
+Have you tried visiting our Docs? https://docs.unsloth.ai/basics/errors-troubleshooting

unsloth/models/_utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2025.5.9"
+__version__ = "2025.5.10"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",
@@ -201,6 +201,35 @@ except:
 
 # Patch get_model_param_count to record correct 4bit / 8bit
 from transformers.trainer_pt_utils import is_deepspeed_zero3_enabled
+
+def extract_approx_params_from_config(config):
+    """
+    Extract approximate parameter count from model config's name_or_path
+    Returns int (param count) or None if not found.
+    """
+    lowercase_b_families = ["gemma"] # gemma uses small 'b' : google/gemma-3-1b-it
+    model_name = getattr(config, "name_or_path", "")
+    import re
+    cleaned = re.sub(r"[-_]?bnb[-_]?4bit|[-_]?4bit|[-_]?8bit|[-_]?bnb", "", model_name, flags=re.IGNORECASE) # replace bnb and xbit
+    match_B = re.search(r"([0-9]+(?:\.[0-9]+)?)\s*B", cleaned) # first prefer searching 'B'
+    if match_B:
+        # most model names would come in this flow
+        billions = float(match_B.group(1))
+        return int(1_000_000_000 * billions)
+    else:
+        if any(fam in cleaned.lower() for fam in lowercase_b_families):
+            match_b = re.search(r"([0-9]+(?:\.[0-9]+)?)\s*b", cleaned)
+            if match_b:
+                billions = float(match_b.group(1))
+                return int(1_000_000_000 * billions)
+        else:
+            match_any = re.search(r"([0-9]+(?:\.[0-9]+)?)\s*[bB]", cleaned)
+            if match_any:
+                billions = float(match_any.group(1))
+                return int(1_000_000_000 * billions)
+    return None
+
+
 def get_model_param_count(model, trainable_only = False):
     """
     Calculate model's total param count. If trainable_only is True then count only those requiring grads
@@ -215,12 +244,9 @@ def get_model_param_count(model, trainable_only = False):
     if (not trainable_only) and \
         hasattr(model, "config") and \
         hasattr(model.config, "quantization_config"):
-
-        billions = re.findall(r"([0-9]{1,})(?:b|B)", model.config.name_or_path)
-        if len(billions) != 0:
-            billions = int(billions[0])
-            s = 1_000_000_000 * billions
-    pass
+        approx = extract_approx_params_from_config(model.config)
+        if approx is not None:
+            s = approx
     return s
 pass
 import transformers.trainer_pt_utils

unsloth/models/llama.py

@@ -99,7 +99,7 @@ torch_nn_functional_softmax = torch.nn.functional.softmax
 SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__
 
 # Fix new HF's inference code
-def _fast_prepare_inputs_for_generation(self, input_ids, **kwargs,):
+def _fast_prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs,):
     past_key_values = kwargs.get("past_key_values", None)
     if past_key_values is not None:
         # Check for uninitialized DynamicCache
@@ -107,11 +107,38 @@ def _fast_prepare_inputs_for_generation(self, input_ids, **kwargs,):
             past_key_values = None
             kwargs["past_key_values"] = None
         else:
+            bs, cache_length = input_ids.shape
             input_ids = input_ids[:,[-1]]
-            kwargs["attention_mask"] = kwargs["attention_mask"][:,[-1]]
+            
+            # Get to the base model
+            base_model = self
+            if hasattr(base_model, 'base_model_prefix'):
+                base_model = getattr(base_model, base_model.base_model_prefix)
+                
+            if hasattr(base_model, "_prepare_4d_causal_attention_mask_with_cache_position"):
+                attention_mask = base_model._prepare_4d_causal_attention_mask_with_cache_position(
+                    attention_mask,
+                    sequence_length=1,
+                    target_length=cache_length,
+                    dtype=self.dtype,
+                    device=input_ids.device,
+                    cache_position=torch.arange(cache_length, cache_length+1, device=input_ids.device),
+                    batch_size=bs,
+                    config=self.config,
+                    past_key_values=past_key_values,
+                )
+            else:
+                attention_mask = attention_mask[:,[-1]]
+                logger.warning_once(
+                    f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method "
+                    "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're "
+                    "writing code, see Llama for an example implementation. If you're a user, please report this "
+                    "issue on GitHub."
+                )
+
     if "cache_position" in kwargs:
         kwargs["position_ids"] = kwargs["cache_position"]
-    return { "input_ids" : input_ids, **kwargs, }
+    return { "input_ids" : input_ids, "attention_mask": attention_mask, **kwargs, }
 pass
 
 

unsloth/models/mapper.py

@@ -864,6 +864,11 @@ __INT_TO_FLOAT_MAPPER = \
         "mistralai/Devstral-Small-2505",
         "unsloth/Devstral-Small-2505-bnb-4bit",
     ),
+    "unsloth/DeepSeek-R1-0528-Qwen3-8B-unsloth-bnb-4bit" : (
+        "unsloth/DeepSeek-R1-0528-Qwen3-8B",
+        "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
+        "unsloth/DeepSeek-R1-0528-Qwen3-8B-bnb-4bit",
+    ),
 }
 
 INT_TO_FLOAT_MAPPER  = {}