diff --git a/quantize_streaming.py b/quantize_streaming.py index 867cee1..0890b66 100644 --- a/quantize_streaming.py +++ b/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) - import concurrent.futures + # Process 4 shards per GPU (8 total) for max parallelism + 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 - gpu = idx % 2 - state_dict = load_file(shard_file, device=f"cuda:{gpu}") + # Load to assigned GPU + state_dict = load_file(shard_file, device=f"cuda:{gpu_id}") 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) - for i in range(0, len(shards), 2): - batch = shards[i:i+2] + # Process in batches (4 per GPU = 8 total) + for i in range(0, len(shards), max_parallel): + 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()