From fa98cee8f460b25627fdc797d13db5c071edf5a5 Mon Sep 17 00:00:00 2001 From: datta0 Date: Sat, 31 May 2025 18:52:46 +0000 Subject: [PATCH 1/2] Fix quant model param fetch regex --- unsloth/models/_utils.py | 39 +++++++++++++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 6 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 0ad258889e..f41f95a85f 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -201,6 +201,36 @@ 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: + for fam in lowercase_b_families: + if fam in cleaned.lower(): + 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 +245,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 From e6b1a3703d7fe19b581074277084c1fa610d4bcc Mon Sep 17 00:00:00 2001 From: datta0 Date: Sun, 1 Jun 2025 05:57:43 +0000 Subject: [PATCH 2/2] Make replacement logic conscise --- unsloth/models/_utils.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index f41f95a85f..0230f84565 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -217,12 +217,11 @@ def extract_approx_params_from_config(config): billions = float(match_B.group(1)) return int(1_000_000_000 * billions) else: - for fam in lowercase_b_families: - if fam in cleaned.lower(): - 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) + 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: