cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 23:32:38 -04:00
parent 78817532ad
commit b742b5549e
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
quantize_streaming.py
View File

@@ -96,7 +96,7 @@ def streaming_quantize(model_path: str, output_path: str):
futures = [] futures = []
for j, shard_file in enumerate(batch): for j, shard_file in enumerate(batch):
idx = i + j 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)) futures.append(executor.submit(process_shard, idx, shard_file, gpu_id))
for future in concurrent.futures.as_completed(futures): for future in concurrent.futures.as_completed(futures):
future.result() future.result()