From a23ecc49f0ef9423a0b24769b380d5248af3400d Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 20:27:06 -0400
Subject: [PATCH] feat: use transformers built-in BnB loading with
 device_map=cpu

---
 quantize_to_bnb.py | 81 ++++++++++++----------------------------------
 1 file changed, 21 insertions(+), 60 deletions(-)

diff --git a/quantize_to_bnb.py b/quantize_to_bnb.py
index 2bafc26..2bd11cd 100644
--- a/quantize_to_bnb.py
+++ b/quantize_to_bnb.py
@@ -1,84 +1,45 @@
 #!/usr/bin/env python3
-"""Properly quantize model to BnB 4-bit using BnB API."""
+"""Quantize bf16 model to BnB 4-bit using transformers' built-in mechanism."""
 
 import argparse
 import gc
 import torch
 from pathlib import Path
-from transformers import AutoModelForCausalLM, AutoConfig
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 
 
 def quantize_model(model_path, output_path):
-    """Load bf16 model, properly quantize to BnB 4-bit, save."""
+    """Load bf16 model, quantize to BnB 4-bit, save."""
     
     print(f"Loading model from: {model_path}")
+    
+    # Load with BnB quantization config and device_map="cpu"
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+    )
+    
+    print("Loading with BnB 4-bit quantization to CPU...")
     model = AutoModelForCausalLM.from_pretrained(
         model_path,
+        quantization_config=bnb_config,
         device_map="cpu",
-        torch_dtype=torch.bfloat16,
+        torch_dtype=torch.float16,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print(f"✓ Model loaded to CPU (~70GB bf16)")
+    print("✓ Model loaded with BnB 4-bit to CPU")
     
-    # Count parameters
-    total_params = sum(p.numel() for p in model.parameters())
-    print(f"   Total parameters: {total_params / 1e9:.2f}B")
-    
-    # Quantize using BnB's actual API
-    print("\nQuantizing with BnB 4-bit...")
-    from bitsandbytes.nn import Linear4bit
-    from torch import nn
-    
-    quantized_count = 0
-    for name, module in list(model.named_modules()):
-        if isinstance(module, nn.Linear) and 'lm_head' not in name:
-            # Create new Linear4bit
-            new_module = Linear4bit(
-                module.in_features,
-                module.out_features,
-                bias=module.bias is not None,
-                compute_dtype=torch.float16,
-                quant_type='nf4',
-            )
-            
-            # Copy weights - Linear4bit will quantize on first forward
-            with torch.no_grad():
-                new_module.weight = nn.Parameter(module.weight.data.clone())
-                if module.bias is not None:
-                    new_module.bias = nn.Parameter(module.bias.data.clone())
-            
-            # Replace in model
-            layers = name.split('.')
-            parent = model
-            for layer in layers[:-1]:
-                parent = getattr(parent, layer)
-            setattr(parent, layers[-1], new_module)
-            quantized_count += 1
-    
-    print(f"✓ Replaced {quantized_count} Linear layers with Linear4bit")
-    
-    # Force quantization by running a dummy forward pass
-    print("   Forcing quantization with dummy input...")
-    dummy_input = torch.randn(1, 32, model.config.hidden_size)
-    with torch.no_grad():
-        try:
-            model(dummy_input)
-        except Exception as e:
-            print(f"   (Dummy forward may fail, but weights should be quantized)")
-    
-    # Count quantized parameters
-    bnb_params = sum(
-        1 for p in model.parameters() 
-        if hasattr(p, 'quant_state') and p.quant_state is not None
+    # Check if model is actually quantized
+    bnb_modules = sum(
+        1 for m in model.modules() 
+        if hasattr(m, 'weight') and hasattr(m.weight, 'quant_state')
     )
-    print(f"   Quantized modules: {bnb_params}")
+    print(f"   BnB quantized modules: {bnb_modules}")
     
-    # Save model config
+    # Save model
     print(f"\nSaving to: {output_path}")
-    model.config.save_pretrained(output_path)
-    
-    # Save weights
     model.save_pretrained(output_path)
     print("✓ Model saved")