From 18ef0f73efbd7bf73085935ecd98c028a52940c9 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 22:48:09 -0400
Subject: [PATCH] fix: don't save quant_state objects (BnB rebuilds on load)

---
 quantize_streaming.py | 15 +++------------
 1 file changed, 3 insertions(+), 12 deletions(-)

diff --git a/quantize_streaming.py b/quantize_streaming.py
index 1f95a2c..67c1a05 100644
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -47,16 +47,12 @@ def streaming_quantize(model_path: str, output_path: str):
 
         print(f"   Quantizing {len(weight_keys)} tensors...")
 
-        quant_states = {}  # We need to save these too for loading later
-
         for key in weight_keys:
             try:
                 weight = state_dict[key]
                 qweight, qstate = quantize_weight_nf4(weight)
                 state_dict[key] = qweight
-                if qstate is not None:
-                    quant_states[f"{key}.quant_state"] = qstate
-                del weight, qweight
+                del weight, qweight, qstate
                 gc.collect()
                 torch.cuda.empty_cache()
             except Exception as e:
@@ -64,19 +60,14 @@ def streaming_quantize(model_path: str, output_path: str):
                 gc.collect()
                 torch.cuda.empty_cache()
 
-        # Move everything to CPU
+        # Move to CPU and save
         state_dict = {
             k: v.cpu() if isinstance(v, torch.Tensor) else v
             for k, v in state_dict.items()
         }
-        quant_states = {
-            k: v.cpu() if isinstance(v, torch.Tensor) else v
-            for k, v in quant_states.items()
-        }
 
-        # Save shard + quant states
         shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
-        save_file({**state_dict, **quant_states}, output_path / shard_name)
+        save_file(state_dict, output_path / shard_name)
 
         print(f"   Saved {shard_name}")