diff --git a/quantize_streaming.py b/quantize_streaming.py index eac58e8..a0f1f4b 100644 --- a/quantize_streaming.py +++ b/quantize_streaming.py @@ -43,11 +43,11 @@ def streaming_quantize(model_path: str, output_path: str): 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" + shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors" if (output_path / shard_name).exists(): return - print(f"[{idx}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})") + print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})") # Load to assigned GPU state_dict = load_file(shard_file, device=f"cuda:{gpu_id}") @@ -88,7 +88,7 @@ def streaming_quantize(model_path: str, output_path: str): # Find first unsaved shard start_idx = 0 for idx in range(len(shards)): - shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors" + shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors" if not (output_path / shard_name).exists(): start_idx = idx break