diff --git a/quantize_streaming.py b/quantize_streaming.py index a0f1f4b..5aecf2d 100644 --- a/quantize_streaming.py +++ b/quantize_streaming.py @@ -30,11 +30,22 @@ def streaming_quantize(model_path: str, output_path: str): output_path = Path(output_path) output_path.mkdir(parents=True, exist_ok=True) - config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) - import glob shards = sorted(glob.glob(f"{model_path}/*.safetensors")) + + # Rename existing 0-indexed files to 1-indexed + for existing in output_path.glob("*.safetensors"): + parts = existing.name.split("-") + if len(parts) >= 3 and existing.name.startswith("model-00000-"): + num = int(parts[1]) + new_name = f"model-{num+1:05d}-of-{len(shards):05d}.safetensors" + new_path = output_path / new_name + existing.rename(new_path) + print(f"Renamed: {existing.name} -> {new_name}") + print(f"Found {len(shards)} shards\n") + + config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) # Process 4 shards per GPU (8 total) for max parallelism shards_per_gpu = 4