fix: remove duplicate sequential code

This commit is contained in:
Christian Medina
2026-07-02 23:12:03 -04:00
parent 7040c84618
commit 9e2c7bec8d

View File

@@ -86,7 +86,7 @@ def streaming_quantize(model_path: str, output_path: str):
# Process in pairs (2 GPUs)
for i in range(0, len(shards), 2):
batch = shards[i:i+2]
if i == 0 and all((output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors").exists() for idx in range(i, i+len(batch))):
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
@@ -98,43 +98,6 @@ def streaming_quantize(model_path: str, output_path: str):
for future in concurrent.futures.as_completed(futures):
future.result()
state_dict = load_file(shard_file, device="cuda:0")
weight_keys = [
k for k, v in state_dict.items()
if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
]
print(f" Quantizing {len(weight_keys)} tensors...")
for key in weight_keys:
try:
weight = state_dict[key]
qweight, qstate = quantize_weight_nf4(weight)
state_dict[key] = qweight
del weight, qweight, qstate
gc.collect()
torch.cuda.empty_cache()
except Exception as e:
print(f" Warning: Failed on {key}: {e}")
gc.collect()
torch.cuda.empty_cache()
# Move to CPU and save
state_dict = {
k: v.cpu() if isinstance(v, torch.Tensor) else v
for k, v in state_dict.items()
}
shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
save_file(state_dict, output_path / shard_name)
print(f" Saved {shard_name}")
del state_dict
gc.collect()
torch.cuda.empty_cache()
config.save_pretrained(output_path)
print(f"\n✅ Done → {output_path}")