From bee66833961d88e45735fe6577648548976bd4d0 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 22:32:38 -0400
Subject: [PATCH] feat: simple NF4 quantization with device_map=auto (proven
 method)

---
 quantize_model.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/quantize_model.py b/quantize_model.py
index 78219fd..e657d79 100644
--- a/quantize_model.py
+++ b/quantize_model.py
@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
-"""Quantize Ornith-1.0-35B to 4-bit NF4 using bitsandbytes (recommended way)."""
+"""Simple NF4 quantization using BnB with device_map auto-distribution."""
 
 import torch
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer
 
+
 def quantize_model(model_path, output_path):
     print(f"Quantizing model from: {model_path}")
-    print("This will use bitsandbytes NF4 quantization with double quantization.\n")
 
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
@@ -15,27 +15,35 @@ def quantize_model(model_path, output_path):
         bnb_4bit_use_double_quant=True,
     )
 
-    print("Loading model with 4-bit quantization (this may take a while)...")
+    print("Loading model with 4-bit quantization...")
+    print("   This will distribute across both GPUs automatically\n")
+
     model = AutoModelForCausalLM.from_pretrained(
         model_path,
         quantization_config=bnb_config,
         device_map="auto",
+        max_memory={0: "28GiB", 1: "28GiB"},  # Leave room for training
         low_cpu_mem_usage=True,
         trust_remote_code=True,
     )
 
-    print("\nSaving quantized model...")
+    print("\n✓ Model loaded and quantized")
+    print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    print(f"   GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
+
+    print(f"\nSaving quantized model to: {output_path}")
     model.save_pretrained(output_path)
-    
-    # Also save tokenizer if it exists
+
+    # Save tokenizer
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_path)
         tokenizer.save_pretrained(output_path)
+        print("✓ Tokenizer saved")
     except:
-        pass
+        print("⚠ No tokenizer found")
 
     print(f"\n✅ Quantized model saved to: {output_path}")
-    print("You can now load it with: AutoModelForCausalLM.from_pretrained(..., load_in_4bit=True)")
+
 
 if __name__ == "__main__":
     import argparse
@@ -43,5 +51,5 @@ if __name__ == "__main__":
     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)