fix: alternate GPUs for each shard (0,1,0,1)

2026-07-02 23:32:38 -04:00
parent 78817532ad
commit b742b5549e
1 changed files with 1 additions and 1 deletions
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -96,7 +96,7 @@ def streaming_quantize(model_path: str, output_path: str):
            futures = []
            for j, shard_file in enumerate(batch):
                idx = i + j
-                gpu_id = (j // shards_per_gpu) % num_gpus  # Assign to GPU
+                gpu_id = idx % num_gpus  # Alternate GPUs: 0,1,0,1...
                futures.append(executor.submit(process_shard, idx, shard_file, gpu_id))
            for future in concurrent.futures.as_completed(futures):
                future.result()