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)
|
# Process in pairs (2 GPUs)
|
||||||
for i in range(0, len(shards), 2):
|
for i in range(0, len(shards), 2):
|
||||||
batch = shards[i:i+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")
|
print(f"Shards {i+1}-{i+len(batch)} already saved, skipping")
|
||||||
continue
|
continue
|
||||||
|
|
||||||
@@ -98,43 +98,6 @@ def streaming_quantize(model_path: str, output_path: str):
|
|||||||
for future in concurrent.futures.as_completed(futures):
|
for future in concurrent.futures.as_completed(futures):
|
||||||
future.result()
|
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)
|
config.save_pretrained(output_path)
|
||||||
print(f"\n✅ Done → {output_path}")
|
print(f"\n✅ Done → {output_path}")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user