From 9a5d455dacc093ef53d46618e10765ffb70ed66f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 14 Mar 2025 02:00:03 -0700 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 075e3b769c..af423d560f 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -72,6 +72,7 @@ from platform import system as platform_system platform_system = platform_system() import numpy as np import contextlib +import re import warnings, subprocess, re, inspect, psutil, os, math from unsloth_zoo.utils import Version @@ -194,20 +195,15 @@ def get_model_param_count(model, trainable_only=False): def numel(p): return p.numel() s = sum(numel(p) for p in model.parameters() if not trainable_only or p.requires_grad) - if hasattr(model, "config") and hasattr(model.config, "quantization_config"): - quantization_config = model.config.quantization_config - if "load_in_4bit" in quantization_config: - load_in_4bit = quantization_config["load_in_4bit"] - else: - load_in_4bit = False - if "load_in_8bit" in quantization_config: - load_in_8bit = quantization_config["load_in_8bit"] - else: - load_in_8bit = False - if load_in_4bit: - s *= 4.5 - elif load_in_8bit: - s *= 2.0 + 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 = max(int(x) for x in billions) + s = 1_000_000_000 * billions + pass return s pass import transformers.trainer_pt_utils