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

fix: free GPU memory after each quantization

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 22:41:52 -04:00
parent 1751d33f69
commit b9f6ddaf50
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
test_quantize_single_shard.py
View File

@@ -38,16 +38,21 @@ def test_single_shard(model_path, output_dir):
quantized = 0 quantized = 0
failed = 0 failed = 0
for key in weight_keys[:5]: # Test first 5 layers for i, key in enumerate(weight_keys[:5]):
try: try:
weight = state_dict[key].to("cuda:0") weight = state_dict[key].to("cuda:0")
qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True) qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
state_dict[key] = qweight.cpu() state_dict[key] = qweight.cpu()
del weight, qweight
gc.collect()
torch.cuda.empty_cache()
quantized += 1 quantized += 1
print(f"✓ {key}") print(f"✓ {key}")
except Exception as e: except Exception as e:
failed += 1 failed += 1
print(f"✗ {key}: {e}") print(f"✗ {key}: {e}")
gc.collect()
torch.cuda.empty_cache()
print(f"\nResults: {quantized} quantized, {failed} failed") print(f"\nResults: {quantized} quantized, {failed} failed")