CI sync from nemotron-embed-quant @ 8ce22703

nemotron3_embed_mlx.py

@@ -8,19 +8,24 @@ Changes:  converted the causal-decoder mlx-lm ministral3 implementation into a
 
 Dependencies: mlx, mlx-lm, transformers, numpy
 """
+
 import json
 from pathlib import Path
 
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
-from mlx_lm.models.ministral3 import (ModelArgs, TransformerBlock,
-                                      _get_llama_4_attn_scale)
+from mlx_lm.models.ministral3 import (
+    ModelArgs,
+    TransformerBlock,
+    _get_llama_4_attn_scale,
+)
 from transformers import AutoTokenizer
 
 NEG_INF = -1e9
 PREFIXES = {"query": "query: ", "passage": "passage: "}
 
+
 class NemotronEmbedModel(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -32,7 +37,8 @@ class NemotronEmbedModel(nn.Module):
     def __call__(self, input_ids: mx.array, attention_mask: mx.array) -> mx.array:
         h = self.embed_tokens(input_ids)
         attn_scale = _get_llama_4_attn_scale(
-            input_ids.shape[1], 0,
+            input_ids.shape[1],
+            0,
             self.args.rope_parameters["llama_4_scaling_beta"],
             self.args.rope_parameters["original_max_position_embeddings"],
         ).astype(h.dtype)
@@ -44,17 +50,23 @@ class NemotronEmbedModel(nn.Module):
         emb = (h * m).sum(axis=1) / m.sum(axis=1)
         return emb / mx.linalg.norm(emb, axis=-1, keepdims=True)
 
+
 def load(path: str):
     p = Path(path)
     if not p.is_dir():
         from huggingface_hub import snapshot_download
+
         p = Path(snapshot_download(path))
     cfg = json.loads((p / "config.json").read_text())
     model = NemotronEmbedModel(ModelArgs.from_dict(cfg))
     q = cfg.get("quantization")
     if q:
-        nn.quantize(model, group_size=q["group_size"], bits=q["bits"],
-                    mode=q.get("mode", "affine"))
+        nn.quantize(
+            model,
+            group_size=q["group_size"],
+            bits=q["bits"],
+            mode=q.get("mode", "affine"),
+        )
     model.load_weights(str(p / "model.safetensors"))
     model.eval()
     mx.eval(model.parameters())
@@ -62,14 +74,19 @@ def load(path: str):
     tok.padding_side = "right"
     return model, tok
 
-def encode(model, tokenizer, texts, input_type=None, batch_size=8,
-           max_length=4096):
+
+def encode(model, tokenizer, texts, input_type=None, batch_size=8, max_length=4096):
     if input_type is not None:
         texts = [PREFIXES[input_type] + t for t in texts]
     out = []
     for i in range(0, len(texts), batch_size):
-        b = tokenizer(texts[i:i + batch_size], padding=True, truncation=True,
-                      max_length=max_length, return_tensors="np")
+        b = tokenizer(
+            texts[i : i + batch_size],
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors="np",
+        )
         e = model(mx.array(b["input_ids"]), mx.array(b["attention_mask"]))
         out.append(np.asarray(e.astype(mx.float32)))
         mx.clear_cache()

scripts/convert_8b.py

@@ -44,8 +44,11 @@ def main():
     ap.add_argument("--out", required=True)
     ap.add_argument("--bits", type=int, required=True)
     ap.add_argument("--group-size", type=int, required=True)
-    ap.add_argument("--mixed", action="store_true",
-                    help="keep embed_tokens at 8-bit (body at --bits)")
+    ap.add_argument(
+        "--mixed",
+        action="store_true",
+        help="keep embed_tokens at 8-bit (body at --bits)",
+    )
     args = ap.parse_args()
 
     import mlx.core as mx
@@ -62,10 +65,13 @@ def main():
     weights = {}
     for shard in sorted(src.glob("model-*.safetensors")):
         weights.update(mx.load(str(shard)))  # mmap — not materialized yet
-    print(f"loaded {len(weights)} tensors from {len(list(src.glob('model-*.safetensors')))} shards")
+    print(
+        f"loaded {len(weights)} tensors from {len(list(src.glob('model-*.safetensors')))} shards"
+    )
     model.load_weights(list(weights.items()))
 
     if args.mixed:
+
         def class_predicate(path, module):
             if not hasattr(module, "to_quantized"):
                 return False
@@ -73,23 +79,30 @@ def main():
                 return {"group_size": args.group_size, "bits": 8}
             return True
 
-        nn.quantize(model, group_size=args.group_size, bits=args.bits,
-                    class_predicate=class_predicate)
-        quant_cfg = {"group_size": args.group_size, "bits": args.bits,
-                     "mode": "affine",
-                     "overrides": {"embed_tokens": {"bits": 8}}}
+        nn.quantize(
+            model,
+            group_size=args.group_size,
+            bits=args.bits,
+            class_predicate=class_predicate,
+        )
+        quant_cfg = {
+            "group_size": args.group_size,
+            "bits": args.bits,
+            "mode": "affine",
+            "overrides": {"embed_tokens": {"bits": 8}},
+        }
     else:
         nn.quantize(model, group_size=args.group_size, bits=args.bits)
-        quant_cfg = {"group_size": args.group_size, "bits": args.bits,
-                     "mode": "affine"}
+        quant_cfg = {"group_size": args.group_size, "bits": args.bits, "mode": "affine"}
 
     mx.eval(model.parameters())
 
     from mlx.utils import tree_flatten
 
     flat = dict(tree_flatten(model.parameters()))
-    mx.save_safetensors(str(out / "model.safetensors"), flat,
-                        metadata={"format": "mlx"})
+    mx.save_safetensors(
+        str(out / "model.safetensors"), flat, metadata={"format": "mlx"}
+    )
 
     cfg_out = dict(cfg)
     cfg_out["quantization"] = quant_cfg
@@ -100,8 +113,16 @@ def main():
     shutil.copy(MLX_1B_DIR / "nemotron3_embed_mlx.py", out / "nemotron3_embed_mlx.py")
 
     size_gb = (out / "model.safetensors").stat().st_size / 1e9
-    print(json.dumps({"out": str(out), "quantization": quant_cfg,
-                      "n_tensors": len(flat), "size_gb": round(size_gb, 2)}))
+    print(
+        json.dumps(
+            {
+                "out": str(out),
+                "quantization": quant_cfg,
+                "n_tensors": len(flat),
+                "size_gb": round(size_gb, 2),
+            }
+        )
+    )
 
 
 if __name__ == "__main__":

scripts/eval_8b.py

@@ -41,8 +41,12 @@ def main():
     ap.add_argument("--make-ref", action="store_true")
     ap.add_argument("--mlx-dir", default=None)
     ap.add_argument("--out", default=None)
-    ap.add_argument("--min-cos", type=float, default=0.0,
-                    help="optional gate on min fixture cosine (0 = report only)")
+    ap.add_argument(
+        "--min-cos",
+        type=float,
+        default=0.0,
+        help="optional gate on min fixture cosine (0 = report only)",
+    )
     args = ap.parse_args()
 
     import numpy as np
@@ -74,24 +78,31 @@ def main():
 
     ref = np.load(REF_NPZ)["embeddings"]
     m = embed_fn(texts)
-    cos = [float(np.dot(a / np.linalg.norm(a), b / np.linalg.norm(b)))
-           for a, b in zip(m, ref)]
+    cos = [
+        float(np.dot(a / np.linalg.norm(a), b / np.linalg.norm(b)))
+        for a, b in zip(m, ref)
+    ]
 
     suite = pooling_fixtures.run(embed_fn, f"mlx-8b:{mlx_dir.name}")
     report = {
         "mlx_dir": str(mlx_dir),
-        "quantization": json.loads((mlx_dir / "config.json").read_text()).get("quantization"),
+        "quantization": json.loads((mlx_dir / "config.json").read_text()).get(
+            "quantization"
+        ),
         "fixture_parity_vs_transformers_bf16": [round(c, 6) for c in cos],
         "parity_mean": round(float(np.mean(cos)), 6),
         "parity_min": round(float(np.min(cos)), 6),
         "pooling_suite": suite,
-        "arch_ok": bool(suite["results"]["padding_invariance"]["pass"]
-                        and suite["results"]["suffix_sensitivity"]["pass"]),
+        "arch_ok": bool(
+            suite["results"]["padding_invariance"]["pass"]
+            and suite["results"]["suffix_sensitivity"]["pass"]
+        ),
     }
     out = Path(args.out) if args.out else ROOT / "runs" / f"mlx_8b_{mlx_dir.name}.json"
     out.write_text(json.dumps(report, indent=2))
-    print(json.dumps({k: v for k, v in report.items() if k != "pooling_suite"},
-                     indent=2))
+    print(
+        json.dumps({k: v for k, v in report.items() if k != "pooling_suite"}, indent=2)
+    )
     if args.min_cos and min(cos) < args.min_cos:
         sys.exit(1)