feat: process 8 shards at once (4 per GPU) with better logging

quantize_streaming.py

@@ -35,20 +35,21 @@ def streaming_quantize(model_path: str, output_path: str):
     shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
     print(f"Found {len(shards)} shards\n")
 
-    # Process 2 shards at a time (one per GPU)
+    # Process 4 shards per GPU (8 total) for max parallelism
-    import concurrent.futures
+    shards_per_gpu = 4
+    num_gpus = 2
+    max_parallel = shards_per_gpu * num_gpus
     
-    def process_shard(idx, shard_file):
+    def process_shard(idx, shard_file, gpu_id):
         """Process a single shard (called per-GPU)."""
         shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
         if (output_path / shard_name).exists():
             return
         
-        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name}")
+        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})")
         
-        # Load to GPU based on index
+        # Load to assigned GPU
-        gpu = idx % 2
+        state_dict = load_file(shard_file, device=f"cuda:{gpu_id}")
-        state_dict = load_file(shard_file, device=f"cuda:{gpu}")
         
         weight_keys = [
             k for k, v in state_dict.items()
@@ -77,24 +78,25 @@ def streaming_quantize(model_path: str, output_path: str):
         }
         
         save_file(state_dict, output_path / shard_name)
-        print(f"   Saved {shard_name}")
+        print(f"   ✓ Saved {shard_name}")
         
         del state_dict
         gc.collect()
         torch.cuda.empty_cache()
     
-    # Process in pairs (2 GPUs)
+    # Process in batches (4 per GPU = 8 total)
-    for i in range(0, len(shards), 2):
+    for i in range(0, len(shards), max_parallel):
-        batch = shards[i:i+2]
+        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=2) as executor:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_parallel) as executor:
             futures = []
             for j, shard_file in enumerate(batch):
                 idx = i + j
-                futures.append(executor.submit(process_shard, idx, shard_file))
+                gpu_id = (j // shards_per_gpu) % num_gpus  # Assign to GPU
+                futures.append(executor.submit(process_shard, idx, shard_file, gpu_id))
             for future in concurrent.futures.as_completed(futures):
                 future.result()