From b624e90b028bdf28fe12f66e224a8b6fdb88dc95 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 20:22:39 -0400
Subject: [PATCH] fix: use BnB quantize_batch to actually quantize weights

---
 quantize_to_bnb.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/quantize_to_bnb.py b/quantize_to_bnb.py
index ab971f6..5638d69 100644
--- a/quantize_to_bnb.py
+++ b/quantize_to_bnb.py
@@ -29,11 +29,20 @@ def quantize_model(model_path, output_path):
     print("\nQuantizing with BnB 4-bit...")
     from bitsandbytes.nn import Linear4bit
     from torch import nn
+    from bitsandbytes import quantize_batch
     
     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 with proper quantization
+            # Quantize weights using BnB
+            weight_2d = module.weight.data.view(-1, module.weight.data.size(-1))
+            quantized_weight, quant_state = quantize_batch(
+                weight_2d.to(torch.float16),
+                blocksize=64,
+                quant_type='nf4',
+            )
+            
+            # Create new Linear4bit with quantized weights
             new_module = Linear4bit(
                 module.in_features,
                 module.out_features,
@@ -42,11 +51,11 @@ def quantize_model(model_path, output_path):
                 quant_type='nf4',
             )
             
-            # Copy weights (BnB will quantize during forward)
-            with torch.no_grad():
-                new_module.weight.data = module.weight.data.clone()
-                if module.bias is not None:
-                    new_module.bias.data = module.bias.data.clone()
+            # Set quantized weights
+            new_module.weight = nn.Parameter(quantized_weight.view_as(module.weight.data))
+            new_module.quant_state = quant_state
+            if module.bias is not None:
+                new_module.bias = nn.Parameter(module.bias.data.clone())
             
             # Replace in model
             layers = name.split('.')