From 36faf2edfd6e1edd5ae32bca11ed68b4d29092ce Mon Sep 17 00:00:00 2001 From: Chris McMaster Date: Sat, 26 Apr 2025 23:30:59 +1000 Subject: [PATCH] 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 --- mlx_lm/models/baichuan_m1.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/mlx_lm/models/baichuan_m1.py b/mlx_lm/models/baichuan_m1.py index be1a4ab..ad1c404 100644 --- a/mlx_lm/models/baichuan_m1.py +++ b/mlx_lm/models/baichuan_m1.py @@ -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