From d8393969fc31f6c9c8edc0078ba98f4a81b3ebd6 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:51:23 -0400
Subject: [PATCH] fix: save quantization states with torch.save

---
 quantize_streaming.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/quantize_streaming.py b/quantize_streaming.py
index b85546f..bf580df 100644
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -69,11 +69,14 @@ def streaming_quantize(model_path: str, output_path: str):
         
         print(f"   Quantizing {len(weight_keys)} tensors...")
         
+        quant_states = {}
         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, qstate
                 gc.collect()
                 torch.cuda.empty_cache()
@@ -88,8 +91,8 @@ def streaming_quantize(model_path: str, output_path: str):
             for k, v in state_dict.items()
         }
         
-        save_file(state_dict, output_path / shard_name)
-        print(f"   ✓ Saved {shard_name}")
+        torch.save({**state_dict, **quant_states}, output_path / shard_name)
+        print(f"   ✓ Saved {shard_name} ({len(quant_states)} quant states)")
         
         del state_dict
         gc.collect()