fix quantized model parameter count method (#2855)

* fix quantized model parameter count method

* function cleanup

* parameter space cleanup
This commit is contained in:
Roland Tannous 2025-07-02 09:36:59 +03:00 committed by GitHub
commit 3691534111
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -206,33 +206,18 @@ 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):
def extract_quant_model_param_count(model):
"""
Extract approximate parameter count from model config's name_or_path
Returns int (param count) or None if not found.
Calculate quant model param count based on difference in param class. Returns int for param count.
"""
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)
count: int = 0
for name, p in model.named_parameters():
if p.__class__.__name__ == "Params4bit":
count += 2 * p.numel()
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
count += p.numel()
return count
pass
def get_model_param_count(model, trainable_only = False):
"""
@ -248,7 +233,7 @@ def get_model_param_count(model, trainable_only = False):
if (not trainable_only) and \
hasattr(model, "config") and \
hasattr(model.config, "quantization_config"):
approx = extract_approx_params_from_config(model.config)
approx = extract_quant_model_param_count(model)
if approx is not None:
s = approx
return s
@ -370,7 +355,7 @@ if is_openai_available():
def _is_openai_available(): return False
transformers.utils.is_openai_available = _is_openai_available
pass
pass
pass
# =============================================
# Get Flash Attention v2 if Ampere (RTX 30xx, A100)
@ -1085,7 +1070,7 @@ pass
def patch_gradient_accumulation_fix(Trainer):
# Fixes gradient accumulation
# Fixes gradient accumulation
import inspect
if hasattr(Trainer, "get_batch_samples"):
if Trainer.get_batch_samples.__name__ == "_unsloth_get_batch_samples": return
@ -1159,10 +1144,10 @@ def patch_gradient_accumulation_fix(Trainer):
"\2if num_items_in_batch is None:\n"\
"\3loss = loss / self.args.gradient_accumulation_steps\n"\
"\1self.accelerator.backward(loss, **kwargs)",
function,
)
exec(function, globals())
Trainer.training_step = _unsloth_training_step
pass
@ -1356,7 +1341,7 @@ def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, m
)
loftq_config = LoftQConfig(loftq_bits = 4, loftq_iter = 1)
pass
if hasattr(model.config, "quantization_config"):
raise ValueError(
"Unsloth: You are using `loftq` init, yet `load_in_4bit = True` was set.\n"\
@ -1365,4 +1350,4 @@ def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, m
pass
pass
return loftq_config
return loftq_config