add scripts/convert_8b.py

scripts/convert_8b.py

@@ -0,0 +1,108 @@
+"""P3b — convert nvidia/Nemotron-3-Embed-8B-BF16 to a quantized MLX checkpoint.
+
+The NVIDIA HF checkpoint's weight names (embed_tokens/layers.N/norm, no "model."
+prefix) and config keys match the proven bidirectional MLX implementation bundled
+with the mlx-community 1B port exactly (see runs/mlx_contract_1b_bf16.json: the
+implementation is 0.9999-exact vs Transformers). Conversion is therefore:
+mmap-load shards -> load into NemotronEmbedModel -> nn.quantize -> save.
+
+One recipe per process invocation (24GB unified memory: 16GB bf16 mmap + the
+quantized materialization must not coexist across recipes).
+
+Usage:
+  source ~/.venvs/nq-mlx/bin/activate
+  python mlx/convert_8b.py --bits 4 --group-size 32 \
+      --src ~/.cache/nq-models/nv-8b-bf16 --out ~/.cache/nq-models/mlx-8b-4bit-gs32
+  python mlx/convert_8b.py --bits 4 --group-size 32 --mixed \
+      --out ~/.cache/nq-models/mlx-8b-mixed-gs32   # embeddings kept 8-bit
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import shutil
+import sys
+from pathlib import Path
+
+MLX_1B_DIR = Path.home() / ".cache/nq-models/mlx-1b-4bit"  # source of the impl
+
+
+def load_impl():
+    spec = importlib.util.spec_from_file_location(
+        "nemotron3_embed_mlx", MLX_1B_DIR / "nemotron3_embed_mlx.py"
+    )
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", default=str(Path.home() / ".cache/nq-models/nv-8b-bf16"))
+    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)")
+    args = ap.parse_args()
+
+    import mlx.core as mx
+    import mlx.nn as nn
+    from mlx_lm.models.ministral3 import ModelArgs
+
+    src, out = Path(args.src).expanduser(), Path(args.out).expanduser()
+    out.mkdir(parents=True, exist_ok=True)
+    mod = load_impl()
+
+    cfg = json.loads((src / "config.json").read_text())
+    model = mod.NemotronEmbedModel(ModelArgs.from_dict(cfg))
+
+    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")
+    model.load_weights(list(weights.items()))
+
+    if args.mixed:
+        def class_predicate(path, module):
+            if not hasattr(module, "to_quantized"):
+                return False
+            if "embed_tokens" in path:
+                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}}}
+    else:
+        nn.quantize(model, group_size=args.group_size, bits=args.bits)
+        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"})
+
+    cfg_out = dict(cfg)
+    cfg_out["quantization"] = quant_cfg
+    cfg_out["quantization_config"] = quant_cfg
+    (out / "config.json").write_text(json.dumps(cfg_out, indent=2))
+    for f in ("tokenizer.json", "tokenizer_config.json"):
+        shutil.copy(src / f, out / f)
+    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)}))
+
+
+if __name__ == "__main__":
+    main()