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

fix: save quantization states with torch.save

Browse Source
This commit is contained in:
Christian Medina
2026-07-03 01:51:23 -04:00
parent aa7961605f
commit d8393969fc
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
quantize_streaming.py
View File

@@ -69,11 +69,14 @@ def streaming_quantize(model_path: str, output_path: str):
print(f" Quantizing {len(weight_keys)} tensors...") print(f" Quantizing {len(weight_keys)} tensors...")
quant_states = {}
for key in weight_keys: for key in weight_keys:
try: try:
weight = state_dict[key] weight = state_dict[key]
qweight, qstate = quantize_weight_nf4(weight) qweight, qstate = quantize_weight_nf4(weight)
state_dict[key] = qweight state_dict[key] = qweight
if qstate is not None:
quant_states[f"{key}.quant_state"] = qstate
del weight, qweight, qstate del weight, qweight, qstate
gc.collect() gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()
@@ -88,8 +91,8 @@ def streaming_quantize(model_path: str, output_path: str):
for k, v in state_dict.items() for k, v in state_dict.items()
} }
save_file(state_dict, output_path / shard_name) torch.save({**state_dict, **quant_states}, output_path / shard_name)
print(f" ✓ Saved {shard_name}") print(f" ✓ Saved {shard_name} ({len(quant_states)} quant states)")
del state_dict del state_dict
gc.collect() gc.collect()