From 4345c7c8bb66495865b11f52db0f34c1fdd496bb Mon Sep 17 00:00:00 2001 From: Alex Barron Date: Sun, 8 Dec 2024 14:05:12 -0800 Subject: [PATCH] count bias too --- llms/mlx_lm/tuner/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/llms/mlx_lm/tuner/utils.py b/llms/mlx_lm/tuner/utils.py index 87f31f68..109d7eee 100644 --- a/llms/mlx_lm/tuner/utils.py +++ b/llms/mlx_lm/tuner/utils.py @@ -250,8 +250,9 @@ def remove_lora_layers(model: nn.Module) -> nn.Module: def nparams(module): - if isinstance(module, (nn.QuantizedLinear, nn.QuantizedEmbedding)): - return module.weight.size * 32 // module.bits + if hasattr(module, "bits"): + n = 0 if not hasattr(module, "bias") else module.bias.size + return n + module.weight.size * 32 // module.bits return sum(v.size for _, v in tree_flatten(module.parameters()))