Update baichuan_m1.py (#127)

* Update baichuan_m1.py The sanitize method was failing with quantized models. I've simplified it to be consistent with implementations in other models. It works now. * Reimplement pre-normalization for non-quantized model * Update baichuan_m1.py --------- Co-authored-by: Awni Hannun <[email protected]>
2025-04-26 06:30:59 -07:00
parent d6a9e61572
commit 36faf2edfd
1 changed files with 6 additions and 5 deletions
@@ -205,12 +205,13 @@ class Model(nn.Module):
        return caches

    def sanitize(self, weights: dict) -> dict:
-        if self.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-        else:
-            # Pre-normalize the lm_head
+        is_quantized = "lm_head.scales" in weights
+        if not is_quantized and "lm_head.weight" in weights:
            w = weights["lm_head.weight"]
-            w = w / (mx.linalg.norm(w, axis=-1, keepdims=True) + 1e-7)
+            dtype = w.dtype
+            w = w.astype(mx.float32)
+            norm = mx.linalg.norm(w, axis=-1, keepdims=True)
+            w = (w / (norm + 1e-7)).astype(dtype)
            weights["lm_head.weight"] = w
        return weights