From 49b99889ef5c10751d3b458c3973a7b83fceed03 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:39:27 -0400
Subject: [PATCH] feat: add proper BnB quantization script (loads to CPU,
 saves)

---
 quantize_proper.py | 59 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)
 create mode 100644 quantize_proper.py

diff --git a/quantize_proper.py b/quantize_proper.py
new file mode 100644
index 0000000..9cf5418
--- /dev/null
+++ b/quantize_proper.py
@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""Proper NF4 quantization using BnB during loading."""
+
+import torch
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
+
+
+def quantize_model(model_path, output_path):
+    print(f"Quantizing model from: {model_path}")
+    print(f"Output: {output_path}\n")
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    print("Loading model with BnB 4-bit to CPU (~70GB bf16 -> ~17GB 4-bit)...")
+    print("   This will use CPU RAM for loading\n")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        quantization_config=bnb_config,
+        device_map="cpu",
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
+
+    print("✓ Model loaded and quantized to CPU")
+    
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    bnb_params = sum(p.numel() for p in model.parameters() if hasattr(p, 'quant_state') and p.quant_state is not None)
+    print(f"   Total: {total_params / 1e9:.2f}B parameters")
+    print(f"   Quantized: {bnb_params / 1e9:.2f}B parameters ({bnb_params/total_params*100:.1f}%)\n")
+
+    print(f"Saving quantized model to: {output_path}")
+    model.save_pretrained(output_path)
+    
+    # Save tokenizer
+    try:
+        tokenizer = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+        # Tokenizer saving would go here if needed
+    except:
+        pass
+
+    print(f"\n✅ Quantized model saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B")
+    parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
+    args = parser.parse_args()
+    
+    quantize_model(args.model_path, args.output_path)