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

---
 quantize_streaming.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

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