From 9e2c7bec8d6fb70ef368a6b88a1ddf0a25f850b9 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 23:12:03 -0400
Subject: [PATCH] fix: remove duplicate sequential code

---
 quantize_streaming.py | 39 +--------------------------------------
 1 file changed, 1 insertion(+), 38 deletions(-)

diff --git a/quantize_streaming.py b/quantize_streaming.py
index 7874486..867cee1 100644
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -86,7 +86,7 @@ def streaming_quantize(model_path: str, output_path: str):
     # Process in pairs (2 GPUs)
     for i in range(0, len(shards), 2):
         batch = shards[i:i+2]
-        if i == 0 and all((output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors").exists() for idx in range(i, i+len(batch))):
+        if all((output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors").exists() for idx in range(i, i+len(batch))):
             print(f"Shards {i+1}-{i+len(batch)} already saved, skipping")
             continue
         
@@ -98,43 +98,6 @@ def streaming_quantize(model_path: str, output_path: str):
             for future in concurrent.futures.as_completed(futures):
                 future.result()
 
-        state_dict = load_file(shard_file, device="cuda:0")
-
-        weight_keys = [
-            k for k, v in state_dict.items()
-            if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
-        ]
-
-        print(f"   Quantizing {len(weight_keys)} tensors...")
-
-        for key in weight_keys:
-            try:
-                weight = state_dict[key]
-                qweight, qstate = quantize_weight_nf4(weight)
-                state_dict[key] = qweight
-                del weight, qweight, qstate
-                gc.collect()
-                torch.cuda.empty_cache()
-            except Exception as e:
-                print(f"   Warning: Failed on {key}: {e}")
-                gc.collect()
-                torch.cuda.empty_cache()
-
-        # Move to CPU and save
-        state_dict = {
-            k: v.cpu() if isinstance(v, torch.Tensor) else v
-            for k, v in state_dict.items()
-        }
-
-        shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
-        save_file(state_dict, output_path / shard_name)
-
-        print(f"   Saved {shard_name}")
-
-        del state_dict
-        gc.collect()
-        torch.cuda.empty_cache()
-
     config.save_pretrained(output_path)
     print(f"\n✅ Done → {output_path}")