fixed bug in peft installation for gptqmodel (#81)

* fixed bug in peft installation for gptqmodel Signed-off-by: 1000850000 user <[email protected]> * changed peft installation on parent qlinear Signed-off-by: 1000850000 user <[email protected]> * additional checks on peft_config target modules Signed-off-by: 1000850000 user <[email protected]> --------- Signed-off-by: 1000850000 user <[email protected]>
foundation-model-stack · Sep 5, 2024 · 2e2736b · 2e2736b
1 parent 6028250
commit 2e2736b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 7 deletions.
diff --git a/plugins/accelerated-peft/src/fms_acceleration_peft/autogptq_utils.py b/plugins/accelerated-peft/src/fms_acceleration_peft/autogptq_utils.py
@@ -32,6 +32,18 @@
 # these parameters are to be patched for triton v2
 # consider making a map if patching more kernels
 PATCH_FOR_FSDP_TRITON_V2 = ["qweight", "qzeros"]
+PEFT_ALL_LINEAR = "all-linear"
+
+
+def requires_installation_on_all_linears(peft_config):
+    tm = peft_config.target_modules
+    assert isinstance(tm, (list, str)), "target modules can only be list or string"
+    if isinstance(tm, list):
+        if PEFT_ALL_LINEAR not in tm:
+            return False
+        assert len(tm) == 1, f"`{PEFT_ALL_LINEAR}` must exist alone in target modules"
+        return True
+    return tm == PEFT_ALL_LINEAR
 
 
 def build_patch_to_view_tensor_to_parameter_for_fsdp_gptq(

diff --git a/plugins/accelerated-peft/src/fms_acceleration_peft/framework_plugin_autogptq.py b/plugins/accelerated-peft/src/fms_acceleration_peft/framework_plugin_autogptq.py
@@ -33,7 +33,10 @@
 import torch.distributed
 
 # Local
-from .autogptq_utils import register_tensors_as_parameters_patch_rule
+from .autogptq_utils import (
+    register_tensors_as_parameters_patch_rule,
+    requires_installation_on_all_linears,
+)
 
 
 class AutoGPTQAccelerationPlugin(AccelerationPlugin):
@@ -318,7 +321,7 @@ def augmentation(
         model = get_gptq_peft_model(
             model,
             peft_config=peft_config,
-            auto_find_all_linears=peft_config.target_modules is None,
+            auto_find_all_linears=requires_installation_on_all_linears(peft_config),
             train_mode=True,  # install adapaters for training
         )
         modifiable_args = (None,)  # return a None for peft_config

diff --git a/plugins/accelerated-peft/src/fms_acceleration_peft/gptqmodel/utils/peft.py b/plugins/accelerated-peft/src/fms_acceleration_peft/gptqmodel/utils/peft.py
@@ -35,7 +35,7 @@
 
 # Local
 from ..models.base import BaseGPTQModel
-from ..nn_modules.qlinear.qlinear_tritonv2 import QuantLinear as QuantLinearTriton
+from ..nn_modules.qlinear import BaseQuantLinear
 
 
 class GPTQLoraConfig(LoraConfig):
@@ -61,7 +61,7 @@ def _create_new_module(
         lora_config: LoraConfig,
         adapter_name: str,
         target: torch.nn.Module,
-        target_cls: torch.nn.Module = QuantLinearTriton,
+        target_cls: torch.nn.Module = BaseQuantLinear,
         **kwargs,
     ):
         # if the base layer module matches a supported class, dispatch the lora linear
@@ -92,12 +92,12 @@ def find_all_linear_names(
 ):
     if not ignore:
         ignore = []
-    lm_head_name = model.lm_head_name
+    lm_head_name = model.lm_head
     if ignore_lm_head and lm_head_name not in ignore:
         ignore.append(lm_head_name)
     results = set()
     for n, m in model.named_modules():
-        if isinstance(m, torch.nn.Linear):
+        if isinstance(m, BaseQuantLinear):
             res = n.split(".")[-1]
             if res not in ignore:
                 results.add(res)
@@ -127,6 +127,7 @@ def get_gptq_peft_model(
     adapter_name: str = "default",
     auto_find_all_linears: bool = True,
     train_mode: bool = False,
+    ignore_lm_head=True,
 ):
     if train_mode and not peft_config:
         raise ValueError("peft_config not specified when in train mode.")
@@ -142,7 +143,7 @@ def get_gptq_peft_model(
         if peft_type in [PeftType.LORA.value]:
             if auto_find_all_linears:
                 peft_config.target_modules = find_all_linear_names(
-                    model, ignore_lm_head=True
+                    model, ignore_lm_head=ignore_lm_head
                 )
             if peft_type == PeftType.LORA.value and not isinstance(
                 peft_config, GPTQLoraConfig