Workaround Low-Mem-Mode Patch for GPTQ-LoRA

plugins/accelerated-peft/src/fms_acceleration_peft/autogptq_utils.py

@@ -16,14 +16,76 @@
 # https://spdx.dev/learn/handling-license-info/
 
 # Standard
-from typing import Callable, List
+from typing import Any, Callable, List
+import importlib
 
 # Third Party
 from peft import LoraConfig
 from peft.tuners.lora.gptq import QuantLinear as LoraLinearGPTQ
 import torch
 
 
+# This function will be replaced after merging
+# https://github.com/foundation-model-stack/fms-acceleration/pull/25
+def _patch_target_module(
+    to_patch: str,
+    replace_with: Any,
+    target_module: str = None,
+):
+    to_patch = to_patch.split(".")
+    assert len(to_patch) > 1, "must have an object to patch"
+
+    to_patch, obj_name_to_patch = to_patch[:-1], to_patch[-1]
+    to_patch = ".".join(to_patch)
+    source = importlib.import_module(to_patch)
+    original_obj = getattr(source, obj_name_to_patch)
+    setattr(source, obj_name_to_patch, replace_with)
+
+    if target_module is not None:
+        # reload and this should get the patched object
+        target_module = importlib.import_module(target_module)
+        importlib.reload(target_module)
+
+        # replace it
+        setattr(source, obj_name_to_patch, original_obj)
+
+
+def make_sure_no_tensor_in_meta_device(
+    model,
+    use_triton: bool,
+    desc_act: bool,
+    group_size: int,
+    bits: int,
+    disable_exllama: bool,
+    disable_exllamav2: bool,
+    use_marlin: bool = False,
+    use_tritonv2: bool = False,
+):
+    # Third Party
+    from auto_gptq.utils.import_utils import (  # pylint: disable=import-outside-toplevel,import-error
+        dynamically_import_QuantLinear,
+    )
+
+    QuantLinear = dynamically_import_QuantLinear(
+        use_triton,
+        desc_act,
+        group_size,
+        bits=bits,
+        disable_exllama=disable_exllama,
+        disable_exllamav2=disable_exllamav2,
+        use_marlin=use_marlin,
+        use_tritonv2=use_tritonv2,
+    )
+    for _, m in model.named_modules():
+        bias = getattr(m, "bias", None)
+        if bias:
+            if isinstance(m, QuantLinear) and bias.device == torch.device("meta"):
+                m.register_buffer(
+                    "bias",
+                    torch.zeros((m.outfeatures), dtype=torch.float16, device="cpu"),
+                )
+
+
 def replace_module_peft(self, parent_module, child_name, new_module, old_module):
 
     # replace the lora linear

plugins/accelerated-peft/src/fms_acceleration_peft/framework_plugin_autogptq.py

@@ -27,6 +27,7 @@
 from peft import LoraConfig, prepare_model_for_kbit_training
 from peft.tuners.lora.model import LoraModel
 from transformers import AutoModelForCausalLM, TrainingArguments
+from transformers.modeling_utils import is_fsdp_enabled
 import torch
 import torch.distributed
 
@@ -48,14 +49,20 @@ def __init__(self, configurations: Dict[str, Dict]):
         )
 
     def model_loader(self, model_name: str, **kwargs):
-
         # guarded imports
         # Third Party
-        from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig #pylint: disable=import-outside-toplevel,import-error
-        from auto_gptq.nn_modules.qlinear.qlinear_tritonv2 import QuantLinear #pylint: disable=import-outside-toplevel,import-error
+        from auto_gptq import (  # pylint: disable=import-outside-toplevel,import-error
+            AutoGPTQForCausalLM,
+            BaseQuantizeConfig,
+        )
+        from auto_gptq.nn_modules.qlinear.qlinear_tritonv2 import (  # pylint: disable=import-outside-toplevel,import-error
+            QuantLinear,
+        )
 
         # Local
-        from .autogptq_utils import patch_forward_to_view_attributes_before_call #pylint: disable=import-outside-toplevel
+        from .autogptq_utils import (  # pylint: disable=import-outside-toplevel
+            patch_forward_to_view_attributes_before_call,
+        )
 
         # Currently we allow only a quantized checkpoint to be loaded, we do not
         # implement the quantization process here.
@@ -84,20 +91,6 @@ def model_loader(self, model_name: str, **kwargs):
         low_cpu_mem_usage = kwargs.get("low_cpu_mem_usage")
         attn_implementation = kwargs.get("attn_implementation")
 
-        if low_cpu_mem_usage:
-            # Note that low_cpu_mem_usage is typically set to
-            # transformers.modeling_utils.is_fsdp_enabled.
-            # e.g.,
-            # https://github.com/huggingface/transformers/blob/a98c41798cf6ed99e1ff17e3792d6e06a2ff2ff3/src/transformers/modeling_utils.py#L2989-L2990
-            # but not doing that now as AutoGPTQ will call make_sure_no_tensor_in_meta_device
-            # https://github.com/AutoGPTQ/AutoGPTQ/blob/ea829c7bbe83561c2b1de26795b6592992373ef7/auto_gptq/modeling/_base.py#L982C17-L982C51
-            # which does not properly check if a QuantLayer has a bias set or not,
-            # https://github.com/AutoGPTQ/AutoGPTQ/blob/ea829c7bbe83561c2b1de26795b6592992373ef7/auto_gptq/modeling/_utils.py#L514
-            raise ValueError(
-                "low_cpu_mem_usage set to True. This may raise error if model has no bias, "
-                "due to AutoGPTQ bug. Not supporting at the moment."
-            )
-
         # there are some kwargs that we wont be passed to AutoModel, so we need
         # to patch them in
         _old_from_config = AutoModelForCausalLM.from_config
@@ -107,14 +100,38 @@ def model_loader(self, model_name: str, **kwargs):
         )
         AutoModelForCausalLM.from_config = _from_config  # patch
 
-        # NOTE: need to set the device map as below as we want to
-        # use AutoGPTQ for training.
+        if is_fsdp_enabled():
+            # Local
+            from .autogptq_utils import (  # pylint: disable=import-outside-toplevel
+                _patch_target_module,
+                make_sure_no_tensor_in_meta_device,
+            )
+
+            # We patch `make_sure_no_tensor_in_meta_device`
+            # from autogptq to avoid errors on models without bias
+            _patch_target_module(
+                to_patch="auto_gptq.modeling._utils.make_sure_no_tensor_in_meta_device",
+                replace_with=make_sure_no_tensor_in_meta_device,
+                target_module="auto_gptq.modeling._base",
+            )
+            low_cpu_mem_usage = True
+
+        # NOTE: need to set the device map as below as we want to use AutoGPTQ for training.
         # device_map is for inference only
-        # ref: https://huggingface.co/docs/accelerate/en/concept_guides/big_model_inference
-        # Thus we set it as below to effectively disable it.
-        device_map = (
-            {"": torch.cuda.current_device()} if torch.cuda.is_available() else None
-        )
+        # https://huggingface.co/docs/accelerate/en/concept_guides/big_model_inference
+        # For low_cpu_mem_usage = True, we have to set the device map to load checkpoints to "cpu"
+        # to avoid gpu consumption before train
+        # This approach will divert consumption to cpu memory,
+        # a better approach would be to load the checkpoints to meta device
+        # QLoRA is currently implemented by the former approach and will encounter the same issue.
+        # see https://github.com/huggingface/transformers/pull/25107#issuecomment-2134833262
+        device_map = {
+            "": (
+                (torch.cuda.current_device() if not low_cpu_mem_usage else "cpu")
+                if torch.cuda.is_available()
+                else None
+            )
+        }
 
         # currently only enable triton_v2, because the triton kernels are the only ones
         # that have backwards
@@ -202,11 +219,19 @@ def augmentation(
     ):
         # guarded imports
         # Third Party
-        from auto_gptq.nn_modules.qlinear.qlinear_tritonv2 import QuantLinear #pylint: disable=import-outside-toplevel,import-error
-        from auto_gptq.utils.peft_utils import GPTQLoraModel, get_gptq_peft_model #pylint: disable=import-outside-toplevel,import-error
+        from auto_gptq.nn_modules.qlinear.qlinear_tritonv2 import (  # pylint: disable=import-outside-toplevel,import-error
+            QuantLinear,
+        )
+        from auto_gptq.utils.peft_utils import (  # pylint: disable=import-outside-toplevel,import-error
+            GPTQLoraModel,
+            get_gptq_peft_model,
+        )
 
         # Local
-        from .autogptq_utils import create_new_module_peft, replace_module_peft #pylint: disable=import-outside-toplevel
+        from .autogptq_utils import (  # pylint: disable=import-outside-toplevel
+            create_new_module_peft,
+            replace_module_peft,
+        )
 
         (peft_config,) = modifiable_args  # unpack modifiable args