diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index a3fc12f6d0..075e3b769c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -181,6 +181,37 @@ try: except: pass +# Patch get_model_param_count to record correct 4bit / 8bit +from transformers.trainer_pt_utils import is_deepspeed_zero3_enabled +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 + """ + if is_deepspeed_zero3_enabled(): + def numel(p): + return p.ds_numel if hasattr(p, "ds_numel") else p.numel() + else: + 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 + return s +pass +import transformers.trainer_pt_utils +transformers.trainer_pt_utils.get_model_param_count = get_model_param_count # ============================================= # =============================================