zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 09:56:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

count bias too

Browse Source
This commit is contained in:
Alex Barron 2024-12-08 14:05:12 -08:00
parent 46109e4141
commit 4345c7c8bb
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
llms/mlx_lm/tuner/utils.py
View File

@ -250,8 +250,9 @@ def remove_lora_layers(model: nn.Module) -> nn.Module:
def nparams(module): def nparams(module):
if isinstance(module, (nn.QuantizedLinear, nn.QuantizedEmbedding)): if hasattr(module, "bits"):
return module.weight.size * 32 // 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())) return sum(v.size for _, v in tree_flatten(module.parameters()))