VRAM optimizations during quant

conversion/adaptivegptq.py

@@ -325,7 +325,7 @@ class AdaptiveGPTQ:
             self.weights = self.weights[self.perm, :]
 
 
-    def quantize(self, keep_qweight = False, apply = False, drop = False):
+    def quantize(self, keep_qweight = False, apply = False):
 
         with torch.inference_mode():
 
@@ -401,19 +401,19 @@ class AdaptiveGPTQ:
             self.qscale_max = qscale_max.to(torch.float16)
             self.qgroups = torch.tensor(qgroups, dtype = torch.short)
 
+            # I love Python
+
+            weights = None
+            error = None
+            scale = None
+            qscale = None
+            qscale_max = None
+            qgroups = None
+            group_idx_list = None
+
             # Apply
 
             if apply:
-                if drop:
-                    weights = None
-                    error = None
-                    scale = None
-                    qscale = None
-                    qscale_max = None
-                    qgroups = None
-                    group_idx_list = None
-                    gc.collect()
-                    torch.cuda.empty_cache()
                 self.apply_quant()
 
 
@@ -431,10 +431,16 @@ class AdaptiveGPTQ:
 
     def apply_quant(self):
 
+        self.hessian = None
+
         qc = self.quant.cpu()
         invperm = self.invperm.cpu()
         q = qc[invperm, :].T
         q = q.reshape(self.quant.T.shape)
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
         q = q.to(self.quant.device)
         self.layer.weight.data = q
 

conversion/quantize.py

@@ -164,7 +164,7 @@ def quant_lm_head(job, module, hidden_states, quantizers, cache, attn_params):
     quantizers["lm_head"].prepare()
 
     qp = qparams_headoptions[job["head_bits"]]
-    quant_linear(job, module, quantizers["lm_head"], qp.get_dict())
+    quant_linear(job, module, quantizers["lm_head"], qp.get_dict(), drop = True)
 
 
 # def testc(module, states, target_states, norm, layers):
@@ -368,6 +368,9 @@ def quant(job, save_fn, model):
                 rfn_sum += (torch.linalg.norm(output - output_ref, 'fro') / torch.linalg.norm(output_ref, 'fro')).item()
                 rfn_count += 1
 
+                output_ref = None
+                output = None
+
             elif i < job["measurement_rows"]:
 
                 x = hidden_states[i].to("cuda:0")
@@ -383,6 +386,10 @@ def quant(job, save_fn, model):
                 logprob_sum += token_log_probs.sum().item()
                 logprob_count += target_ids.numel()
 
+                output = None
+                logits = None
+                token_log_probs = None
+
         if mode != "linear":
 
             err = rfn_sum / rfn_count
@@ -407,6 +414,7 @@ def quant(job, save_fn, model):
             # hidden_states = target_states
             # hidden_states = [(x + y) / 2 for x, y in zip(target_states, q_states)]
             hidden_states = q_states
+            q_states = None
 
         # Checkpoint