fix bug when tie_word_embeddings is true

Author: yuhao-zh

src/parallax/server/shard_loader.py

@@ -223,7 +223,12 @@ def load(
                         is_needed = True
                         remapped_key = key.replace("model.", "", 1)
                         if model_shard.is_last_shard and config.get("tie_word_embeddings", False):
-                            shard_weights["lm_head.weight"] = mx.array(f.get_tensor(key))
+                            if key == "model.embed_tokens.weight":
+                                shard_weights["lm_head.weight"] = mx.array(f.get_tensor(key))
+                            elif key == "model.embed_tokens.scales":
+                                shard_weights["lm_head.scales"] = mx.array(f.get_tensor(key))
+                            elif key == "model.embed_tokens.biases":
+                                shard_weights["lm_head.biases"] = mx.array(f.get_tensor(key))
                     elif model_shard.is_last_shard:
                         if "model.norm" in key:
                             is_needed = True