From 7040c846187df9e5d696325a1cd264d455dc6c53 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 23:05:57 -0400
Subject: [PATCH] feat: process 2 shards at a time (one per GPU)

---
 quantize_streaming.py | 62 ++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 58 insertions(+), 4 deletions(-)

diff --git a/quantize_streaming.py b/quantize_streaming.py
index c7ebf62..7874486 100644
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -35,14 +35,68 @@ def streaming_quantize(model_path: str, output_path: str):
     shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
     print(f"Found {len(shards)} shards\n")
 
-    for idx, shard_file in enumerate(shards):
-        # Skip already processed shards
+    # Process 2 shards at a time (one per GPU)
+    import concurrent.futures
+    
+    def process_shard(idx, shard_file):
+        """Process a single shard (called per-GPU)."""
         shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
         if (output_path / shard_name).exists():
-            print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (already saved, skipping)")
-            continue
+            return
         
         print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name}")
+        
+        # Load to GPU based on index
+        gpu = idx % 2
+        state_dict = load_file(shard_file, device=f"cuda:{gpu}")
+        
+        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()
+        }
+        
+        save_file(state_dict, output_path / shard_name)
+        print(f"   Saved {shard_name}")
+        
+        del state_dict
+        gc.collect()
+        torch.cuda.empty_cache()
+    
+    # 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))):
+            print(f"Shards {i+1}-{i+len(batch)} already saved, skipping")
+            continue
+        
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+            futures = []
+            for j, shard_file in enumerate(batch):
+                idx = i + j
+                futures.append(executor.submit(process_shard, idx, shard_file))
+            for future in concurrent.futures.as_completed(futures):
+                future.result()
 
         state_dict = load_file(shard_file, device="cuda:0")