electron-rare/mlx-lm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Apertus tie_word_embeddings fix (#1143)

Browse Source
This commit is contained in:
Andrei Panferov
2026-04-20 23:47:40 +02:00
committed by GitHub
parent a401730941
commit e92138cb01
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx_lm/models/apertus.py
+9 -2
View File
@@ -167,7 +167,8 @@ class Model(nn.Module):
self.args = args
self.model_type = args.model_type
self.model = ApertusModel(args)
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
if not args.tie_word_embeddings:
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
def __call__(
self,
@@ -175,12 +176,18 @@ class Model(nn.Module):
cache: Optional[Any] = None,
) -> mx.array:
out = self.model(inputs, cache)
return self.lm_head(out)
if self.args.tie_word_embeddings:
out = self.model.embed_tokens.as_linear(out)
else:
out = self.lm_head(out)
return out
def sanitize(self, weights):
for k, v in weights.items():
if k.endswith("alpha_p") or k.endswith("alpha_n"):
weights[k] = v.squeeze()
if self.args.tie_word_embeddings:
weights.pop("lm_head.weight", None)
return weights
@property