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

feat: skip already processed shards (resume support)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 23:04:52 -04:00
parent d1f0c81ac4
commit d698a32f0f
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
quantize_streaming.py
View File

@@ -36,6 +36,12 @@ def streaming_quantize(model_path: str, output_path: str):
print(f"Found {len(shards)} shards\n") print(f"Found {len(shards)} shards\n")
for idx, shard_file in enumerate(shards): for idx, shard_file in enumerate(shards):
# Skip already processed shards
shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
if (output_path / shard_name).exists():
print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (already saved, skipping)")
continue
print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name}") print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name}")
state_dict = load_file(shard_file, device="cuda:0") state_dict = load_file(shard_file, device="cuda:0")