fix: reduce to 3 shards per GPU (6 total) for memory

2026-07-03 00:23:21 -04:00
parent 23dcdd2c6f
commit 8a30ec92c6
1 changed files with 2 additions and 2 deletions
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -47,8 +47,8 @@ def streaming_quantize(model_path: str, output_path: str):
    
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)

-    # Process 4 shards per GPU (8 total) for max parallelism
-    shards_per_gpu = 4
+    # Process 3 shards per GPU (6 total) - adjust based on GPU memory
+    shards_per_gpu = 3
    num_gpus = 2
    max_parallel = shards_per_gpu * num_gpus