fix: remove duplicate sequential code
This commit is contained in:
@@ -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}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user