diff --git a/quantize_streaming.py b/quantize_streaming.py index b6ad47f..84beac1 100644 --- a/quantize_streaming.py +++ b/quantize_streaming.py @@ -85,18 +85,26 @@ def streaming_quantize(model_path: str, output_path: str): gc.collect() torch.cuda.empty_cache() - # Process in batches (4 per GPU = 8 total) - for i in range(0, len(shards), max_parallel): + # Find first unsaved shard + start_idx = 0 + for idx in range(len(shards)): + shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors" + if not (output_path / shard_name).exists(): + start_idx = idx + break + + 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] - 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=max_parallel) as executor: futures = [] for j, shard_file in enumerate(batch): idx = i + j - gpu_id = idx % num_gpus # Alternate GPUs: 0,1,0,1... + # 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()