fix gpu cache

gptq_quantization.py

@@ -138,38 +138,34 @@ class GPTQLayerWrapper:
 
         if is_transformer_conv1d(self.layer):
             Q = Q.t()
-        self.layer.weight = nn.Parameter(Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype), requires_grad=False)
-
+        shape = self.layer.weight.shape
+        dtype = self.layer.weight.data.dtype
+        del self.layer.weight
+        setattr(self.layer, "weight", nn.Parameter(Q.reshape(shape).to(dtype), requires_grad=False))
         del self.H
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    def release_gpu_memory(self):
-        if hasattr(self, "H"):
-            del self.H
 
 
 class GPTQBlockWrapper:
-    def __init__(self, module_name: str, module: nn.Module, weight_bit_width=8):
+    def __init__(self, block_name: str, block: nn.Module, weight_bit_width=8):
         self.layer_wrappers = {}
         self.hook_handles = []
-        # module order in the whole network
+        # block order in the whole network
         self.order = 0
-        self.module_name = module_name
+        self.block_name = block_name
 
         def get_hook(layer_name):
             def record_hook(_, x):
                 self.layer_wrappers[layer_name].record_h(x[0])
             return record_hook
 
-        for layer_name, layer in module.named_modules():
+        for layer_name, layer in block.named_modules():
             if isinstance(layer, tuple(QUANT_LAYERS)):
-                full_layer_name = f"{module_name}.{layer_name}" if layer_name else f"{module_name}"
+                full_layer_name = f"{block_name}.{layer_name}" if layer_name else f"{block_name}"
                 self.layer_wrappers[full_layer_name] = GPTQLayerWrapper(full_layer_name, layer, weight_bit_width)
                 handle = layer.register_forward_pre_hook(get_hook(full_layer_name))
                 self.hook_handles.append(handle)
 
-    def quant_module(self):
+    def quant_block(self):
         for _, wrapper in self.layer_wrappers.items():
             wrapper.quant_weight()
 
@@ -190,10 +186,6 @@ class GPTQBlockWrapper:
         for n, l in self.layer_wrappers.items():
             l.is_record = False
 
-    def release_gpu_memory(self):
-        for _, wrapper in self.layer_wrappers.items():
-            wrapper.release_gpu_memory()
-
 
 class GPTQuantizer:
     def __init__(self, block_type: Optional[List[type]] = None):
@@ -207,19 +199,13 @@ class GPTQuantizer:
                 child_prefix = f"{prefix}.{name}" if prefix else name
                 if isinstance(child, tuple(self.block_type)):
                     self.gptq_block_wrappers[name] = GPTQBlockWrapper(child_prefix, child, weight_bit_width)
-                    LOGGER.debug(f"Calibrate module {child_prefix} as a whole block in GPTQ")
+                    LOGGER.debug(f"Calibrate block {child_prefix} as a whole block in GPTQ")
                 else:
                     wrap_block(child, child_prefix)
 
         wrap_block(model)
         return model
 
-    def quantize(self, model: nn.Module):
-        for _, module_wrapper in self.gptq_block_wrappers.items():
-            module_wrapper.quant_module()
-
-        return model
-
     @property
     def calibration_iters(self):
         return len(self.gptq_block_wrappers)
@@ -230,56 +216,59 @@ class GPTQuantizer:
         record_handles = []
         orders = {}
         try:
-            def get_record_order_hook(module_name):
+            def get_record_order_hook(block_name):
                 def record_hook(*args, **kwargs):
                     nonlocal counter
-                    if module_name not in orders:
-                        orders[module_name] = counter
+                    if block_name not in orders:
+                        orders[block_name] = counter
                         counter += 1
                 return record_hook
 
-            for module_name, module_wrapper in self.gptq_block_wrappers.items():
+            for block_name, block_wrapper in self.gptq_block_wrappers.items():
                 # disable the record
-                for _, layer_wrapper in module_wrapper.layer_wrappers.items():
+                for _, layer_wrapper in block_wrapper.layer_wrappers.items():
                     layer_wrapper.is_record = False
 
-                one_layer_wrapper_in_module = list(module_wrapper.layer_wrappers.values())[0]
-                handles = one_layer_wrapper_in_module.layer.register_forward_pre_hook(get_record_order_hook(module_name))
+                one_layer_wrapper_in_block = list(block_wrapper.layer_wrappers.values())[0]
+                handles = one_layer_wrapper_in_block.layer.register_forward_pre_hook(get_record_order_hook(block_name))
                 record_handles.append(handles)
             yield
         except Exception as e:
             logging.warning(e)
         finally:
-            for module_name, order in orders.items():
-                self.gptq_block_wrappers[module_name].set_order(order)
+            for block_name, order in orders.items():
+                self.gptq_block_wrappers[block_name].set_order(order)
 
             for h in record_handles:
                 h.remove()
 
-            for module_name, module_wrapper in self.gptq_block_wrappers.items():
+            for _, block_wrapper in self.gptq_block_wrappers.items():
                 # disable the record
-                for _, layer_wrapper in module_wrapper.layer_wrappers.items():
+                for _, layer_wrapper in block_wrapper.layer_wrappers.items():
                     layer_wrapper.is_record = True
 
 
     @contextlib.contextmanager
     def start_calib_iter(self, i):
         assert i < len(self.gptq_block_wrappers)
-        target_module_wrapper = None
+        target_block_wrapper = None
         try:
-            for _, module_wrapper in self.gptq_block_wrappers.items():
-                if module_wrapper.get_order() == i:
-                    module_wrapper.enable()
-                    target_module_wrapper = module_wrapper
+            for _, block_wrapper in self.gptq_block_wrappers.items():
+                if block_wrapper.get_order() == i:
+                    block_wrapper.enable()
+                    target_block_wrapper = block_wrapper
                 else:
-                    module_wrapper.disable()
+                    block_wrapper.disable()
             yield
         finally:
-            target_module_wrapper.quant_module()
+            target_block_wrapper.quant_block()
 
-    def release_gpu_memory(self):
-        for block_name, block_wrapper in self.gptq_block_wrappers.items():
-            block_wrapper.release_gpu_memory()
+    def release_reference(self):
+        # delete reference so that `torch.cuda.empty_cache()` can
+        # release all the gpu memory cache used during calibration
+        for _, block_wrapper in self.gptq_block_wrappers.items():
+            for _, layer_wrapper in block_wrapper.layer_wrappers.items():
+                del layer_wrapper.layer
 
         torch.cuda.empty_cache()
 
@@ -301,10 +290,12 @@ def gptq_quantize(model, tokenizer, weight_bit_width, calib_data):
     calib_model = quantizer.wrap_model(model, weight_bit_width)
     with quantizer.record_order():
         calib_model.chat(tokenizer, calib_data[0], history=[])
+
     logging.info("Start doing calibration using GPTQ ")
     for i in range(quantizer.calibration_iters):
         logging.info(f"Process: {i + 1}/{quantizer.calibration_iters}")
         # todo: should add early return to speed up the calibration
+        # todo: add cpu offload to reduce the gpu memory requirements.
         with quantizer.start_calib_iter(i):
             for prompt in calib_data:
                 model.chat(tokenizer, prompt, history=[])
@@ -328,5 +319,6 @@ def gptq_quantize(model, tokenizer, weight_bit_width, calib_data):
             )
             parent.add_module(name_in_parent, quantized_layer)
 
-    torch.cuda.empty_cache()
+    # release the memory caache during calibration
+    quantizer.release_reference()
     return