feat: continuous processing - add to GPU with fewer tasks (max 4 per GPU)

quantize_streaming.py

@@ -47,10 +47,9 @@ def streaming_quantize(model_path: str, output_path: str):
     
     config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
 
-    # Process 3 shards per GPU (6 total) - adjust based on GPU memory
-    shards_per_gpu = 3
+    # Process continuously: max 4 shards per GPU, add to whichever GPU has room
+    max_per_gpu = 4
     num_gpus = 2
-    max_parallel = shards_per_gpu * num_gpus
     
     def process_shard(idx, shard_file, gpu_id):
         """Process a single shard (called per-GPU)."""
@@ -106,19 +105,61 @@ def streaming_quantize(model_path: str, output_path: str):
     
     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]
+    # Process continuously: submit to GPU with fewer active tasks
+    import threading
+    gpu_locks = [threading.Lock() for _ in range(num_gpus)]
+    gpu_counts = [0] * num_gpus
     
-        with concurrent.futures.ThreadPoolExecutor(max_workers=max_parallel) as executor:
-            futures = []
-            for j, shard_file in enumerate(batch):
-                idx = i + j
-                # 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()
+    def get_next_gpu():
+        """Get the GPU with fewest active tasks (max 4 per GPU)."""
+        with gpu_locks[0]:
+            with gpu_locks[1]:
+                if gpu_counts[0] <= gpu_counts[1] and gpu_counts[0] < max_per_gpu:
+                    return 0
+                elif gpu_counts[1] < max_per_gpu:
+                    return 1
+                else:
+                    return None  # Both GPUs at max
+    
+    # Create a queue of unprocessed shards
+    from collections import deque
+    remaining = deque(range(start_idx, len(shards)))
+    
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_per_gpu * num_gpus) as executor:
+        futures = []
+        
+        def submit_next():
+            """Submit next shard to available GPU."""
+            gpu_id = get_next_gpu()
+            if gpu_id is None or not remaining:
+                return
+            
+            idx = remaining.popleft()
+            shard_file = shards[idx]
+            
+            with gpu_locks[gpu_id]:
+                gpu_counts[gpu_id] += 1
+            
+            future = executor.submit(process_shard, idx, shard_file, gpu_id)
+            futures.append(future)
+            
+            # When this future completes, release GPU slot and submit next
+            def on_complete(f):
+                with gpu_locks[gpu_id]:
+                    gpu_counts[gpu_id] -= 1
+                submit_next()  # Try to submit next shard
+            
+            future.add_done_callback(on_complete)
+        
+        # Start submitting
+        while remaining:
+            gpu_id = get_next_gpu()
+            if gpu_id is None:
+                break
+            submit_next()
+        
+        # Wait for all to complete
+        concurrent.futures.wait(futures)
 
     config.save_pretrained(output_path)
     print(f"\n✅ Done → {output_path}")