fixed bug in peft installation for gptqmodel

Signed-off-by: 1000850000 user <[email protected]>
foundation-model-stack · Sep 4, 2024 · 486d0ac · 486d0ac
1 parent b159600
commit 486d0ac
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
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
@@ -298,11 +298,11 @@ def augmentation(
                 replace_module_peft, GPTQLoraModel
             )
 
-        # Install GPTQ adapters using the AutoGPTQ package (with the above patches)
+        # Install GPTQ adapters using the AutoGPTQ package (with the above patches)        
         model = get_gptq_peft_model(
             model,
             peft_config=peft_config,
-            auto_find_all_linears=peft_config.target_modules is None,
+            auto_find_all_linears=peft_config.target_modules == 'all-linear',
             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
@@ -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, QuantLinearTriton):
             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