From b9f6ddaf508a2e35a9cb68620f38a29fc340459e Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 22:41:52 -0400
Subject: [PATCH] fix: free GPU memory after each quantization

---
 test_quantize_single_shard.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

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