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

fix mamba models conversion (#1065)

Browse Source
This commit is contained in:
Awni Hannun 2024-10-22 15:44:08 -07:00 committed by GitHub
parent d1d480867b
commit 9000e280ae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
llms/mlx_lm/models/mamba.py
View File

@ -205,7 +205,7 @@ class Model(nn.Module):
def sanitize(self, weights): def sanitize(self, weights):
for k, v in weights.items(): for k, v in weights.items():
if "conv1d.weight" in k and v.ndim == 3: if "conv1d.weight" in k and v.shape[-1] != 1:
weights[k] = v.moveaxis(2, 1) weights[k] = v.moveaxis(2, 1)
return weights return weights

2
llms/mlx_lm/models/recurrent_gemma.py
View File

@ -440,7 +440,7 @@ class Model(nn.Module):
def sanitize(self, weights): def sanitize(self, weights):
for k, v in weights.items(): for k, v in weights.items():
if "conv_1d.weight" in k and v.ndim == 3: if "conv_1d.weight" in k and v.shape[-1] != 1:
weights[k] = v.moveaxis(2, 1) weights[k] = v.moveaxis(2, 1)
if "lm_head.weight" not in weights: if "lm_head.weight" not in weights:
self.pop("lm_head") self.pop("lm_head")