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

fix moe conversion (#802)

Browse Source
This commit is contained in:
Awni Hannun
2024-05-31 12:36:05 -07:00
committed by GitHub
parent f49c5f2829
commit 09aaeac72c
3 changed files with 17 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
llms/mlx_lm/models/phixtral.py
View File

@@ -182,11 +182,11 @@ class Model(nn.Module):
prefix = f"transformer.h.{l}"
for n in ["fc1", "fc2"]:
for k in ["weight", "scales", "biases", "bias"]:
to_join = [
weights.pop(f"{prefix}.moe.mlp.{e}.{n}.{k}")
for e in range(self.args.num_local_experts)
]
if to_join:
if f"{prefix}.moe.mlp.0.{n}.{k}" in weights:
to_join = [
weights.pop(f"{prefix}.moe.mlp.{e}.{n}.{k}")
for e in range(self.args.num_local_experts)
]
weights[f"{prefix}.moe.switch_mlp.{n}.{k}"] = mx.stack(to_join)
return weights