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

fix: free GPU memory after each layer quantization

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 22:43:27 -04:00
parent b9f6ddaf50
commit 8c13be4d1f
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
quantize_streaming.py
View File

@@ -51,12 +51,18 @@ def streaming_quantize(model_path: str, output_path: str):
for key in weight_keys: for key in weight_keys:
try: try:
qweight, qstate = quantize_weight_nf4(state_dict[key]) weight = state_dict[key]
qweight, qstate = quantize_weight_nf4(weight)
state_dict[key] = qweight state_dict[key] = qweight
if qstate is not None: if qstate is not None:
quant_states[f"{key}.quant_state"] = qstate quant_states[f"{key}.quant_state"] = qstate
del weight, qweight
gc.collect()
torch.cuda.empty_cache()
except Exception as e: except Exception as e:
print(f" Warning: Failed on {key}: {e}") print(f" Warning: Failed on {key}: {e}")
gc.collect()
torch.cuda.empty_cache()
# Move everything to CPU # Move everything to CPU
state_dict = { state_dict = {