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

Merge branch 'main' into adding-support-for-mamba2

Browse Source
This commit is contained in:
Gökdeniz Gülmez
2024-12-10 14:32:44 +01:00
committed by GitHub
parent 38e5801edb 12083c4b7e
commit ddad2105ef
31 changed files with 1579 additions and 414 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -98,6 +98,7 @@ def linear_to_lora_layers(
"cohere",
"minicpm",
"deepseek",
"olmo2",
]:
keys = set(["self_attn.q_proj", "self_attn.v_proj"])
if model.model_type in ["mixtral", "phimoe"]:
@@ -150,6 +151,8 @@ def linear_to_lora_layers(
"mixer.out_proj",
]
)
elif model.model_type == "exaone":
keys = set(["attn.attention.q_proj", "attn.attention.v_proj"])
else:
raise ValueError(f"Lora does not support {model.model_type}")
@@ -256,12 +259,14 @@ def remove_lora_layers(model: nn.Module) -> nn.Module:
return model
def print_trainable_parameters(model):
def nparams(m):
if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
return m.weight.size * (32 // m.bits)
return sum(v.size for _, v in tree_flatten(m.parameters()))
def nparams(module):
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()))
def print_trainable_parameters(model):
leaf_modules = tree_flatten(
model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
)