From 0878b11e9801e24e8b0abf758d0cc45b82b523b1 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:41:51 -0400
Subject: [PATCH] fix: restore streaming with device_map=auto and max_memory

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

diff --git a/quantize_streaming.py b/quantize_streaming.py
index 3afd486..0f5ef02 100644
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -39,9 +39,10 @@ def streaming_quantize(model_path: str, output_path: str):
         bnb_4bit_use_double_quant=True,
     )
 
-    print("Loading model with BnB 4-bit (shards streamed to GPUs)...")
-    print("   This will distribute across both GPUs\n")
+    print("Loading model with BnB 4-bit (shards streamed via accelerate)...")
+    print("   This will stream shards and quantize them\n")
 
+    # This streams shards and applies BnB correctly
     model = AutoModelForCausalLM.from_pretrained(
         model_path,
         quantization_config=bnb_config,