From 262b1181b77df1571f06eec6d2da81f0c3426d2d Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 23:39:24 -0400
Subject: [PATCH] fix: skip to first unsaved shard, fix GPU assignment

---
 quantize_streaming.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/quantize_streaming.py b/quantize_streaming.py
index b6ad47f..84beac1 100644
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -85,18 +85,26 @@ def streaming_quantize(model_path: str, output_path: str):
         gc.collect()
         torch.cuda.empty_cache()
     
-    # Process in batches (4 per GPU = 8 total)
-    for i in range(0, len(shards), max_parallel):
+    # Find first unsaved shard
+    start_idx = 0
+    for idx in range(len(shards)):
+        shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
+        if not (output_path / shard_name).exists():
+            start_idx = idx
+            break
+    
+    print(f"Starting from shard {start_idx+1}/16\n")
+    
+    # Process in batches from start_idx
+    for i in range(start_idx, len(shards), max_parallel):
         batch = shards[i:i+max_parallel]
-        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
         
         with concurrent.futures.ThreadPoolExecutor(max_workers=max_parallel) as executor:
             futures = []
             for j, shard_file in enumerate(batch):
                 idx = i + j
-                gpu_id = idx % num_gpus  # Alternate GPUs: 0,1,0,1...
+                # Assign 4 shards per GPU: 0-3→GPU0, 4-7→GPU1, 8-11→GPU0, etc.
+                gpu_id = (idx // shards_per_gpu) % num_gpus
                 futures.append(executor.submit(process_shard, idx, shard_file, gpu_id))
             for future in concurrent.futures.as_completed(futures):
                 future.result()