fix embed_tokens for last layer in qwen models

9c1bea97e86c7a1eb277a2410ac9052cf998c7bb · 2025-01-28 23:09:45 +0000 · Alex Cheema

Files touched

M exo/inference/mlx/models/qwen2.py

Diff

commit 9c1bea97e86c7a1eb277a2410ac9052cf998c7bb
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Tue Jan 28 23:09:45 2025 +0000

    fix embed_tokens for last layer in qwen models
---
 exo/inference/mlx/models/qwen2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/inference/mlx/models/qwen2.py b/exo/inference/mlx/models/qwen2.py
index f2ffc7fb..492fbb27 100644
--- a/exo/inference/mlx/models/qwen2.py
+++ b/exo/inference/mlx/models/qwen2.py
@@ -31,7 +31,7 @@ class Qwen2Model(nn.Module):
     self.num_hidden_layers = args.num_hidden_layers
     assert self.vocab_size > 0
 
-    if self.args.shard.is_first_layer():
+    if self.args.shard.is_first_layer() or (self.args.shard.is_last_layer() and args.tie_word_embeddings):
       self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
 
     self.layers = []

← af171f06 propagate prompts to other nodes so they can display them, c · back to Exo · Fix AMD device capabilities fields 9cf6818f →