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

adding olmoe to training

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2025-03-04 21:08:55 +01:00
parent ef8ec7a27a
commit bbde6ea4bc
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -98,6 +98,7 @@ def linear_to_lora_layers(
"minicpm", "minicpm",
"deepseek", "deepseek",
"olmo2", "olmo2",
"olmoe",
"internlm3", "internlm3",
]: ]:
keys = set(["self_attn.q_proj", "self_attn.v_proj"]) keys = set(["self_attn.q_proj", "self_attn.v_proj"])
@ -106,6 +107,8 @@ def linear_to_lora_layers(
if model.model_type == "qwen2_moe": if model.model_type == "qwen2_moe":
keys.add("mlp.gate") keys.add("mlp.gate")
keys.add("mlp.shared_expert_gate") keys.add("mlp.shared_expert_gate")
if model.model_type == "olmoe":
keys.add("mlp.gate")
elif model.model_type == "gpt_bigcode": elif model.model_type == "gpt_bigcode":
keys = set(["attn.c_attn"]) keys = set(["attn.c_attn"])