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

fix: use 1-indexing for shard names (model-00001-of-00016)

Browse Source
This commit is contained in:
Christian Medina
2026-07-03 00:15:39 -04:00
parent 50623cdcd0
commit b4538c013b
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
quantize_streaming.py
View File

@@ -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