Repository: huggingface/peft

Issue:
GPT2 LoRA _find_and_replace raise tensor shape 0 error with DeepSpeed CPU offloading
my accelerate yaml is 

```yaml
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 0.7
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: True
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
mixed_precision: 'no'
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
use_cpu: false
```

reproduce file 

```python
import os
os.environ['HF_DATASETS_CACHE']="/nfs-nlp-ali/ray_train/huggingface_cache"
os.environ['TRANSFORMERS_CACHE']='/nfs-nlp-ali/ray_train/huggingface_cache/hub'
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import get_peft_model, LoraConfig, TaskType
from accelerate import Accelerator, DistributedType
from accelerate.logging import get_logger
from accelerate.utils import set_seed
peft_config = LoraConfig(
  task_type=TaskType.CAUSAL_LM,
  inference_mode=False,
  r=16,
  lora_alpha=32,
  lora_dropout=0.1,
)
accelerator = Accelerator(gradient_accumulation_steps=1)
tokenizer = AutoTokenizer.from_pretrained("gpt2-xl")
model = AutoModelForCausalLM.from_pretrained("gpt2-xl")
model = get_peft_model(model, peft_config)
accelerator.print(model.print_trainable_parameters())
```


run command 

```bash
accelerate launch --config_file ds_zero3_cpu.yaml test.py
```
error output
```bash
Traceback (most recent call last):
  File "test.py", line 19, in <module>
    model = get_peft_model(model, peft_config)
  File "/opt/miniconda/lib/python3.8/site-packages/peft/mapping.py", line 145, in get_peft_model
    return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config)
  File "/opt/miniconda/lib/python3.8/site-packages/peft/peft_model.py", line 512, in __init__
    super().__init__(model, peft_config)
  File "/opt/miniconda/lib/python3.8/site-packages/peft/peft_model.py", line 80, in __init__
    self.base_model = LoraModel(peft_config, model)
  File "/opt/miniconda/lib/python3.8/site-packages/peft/tuners/lora.py", line 118, in __init__
    self._find_and_replace()
  File "/opt/miniconda/lib/python3.8/site-packages/peft/tuners/lora.py", line 163, in _find_and_replace
    in_features, out_features = target.weight.shape
```




Discussion:
i have try to modify the https://github.com/huggingface/peft/blob/ce4e6f3dd91e759ccaf91c76f73b2e03de1c71c0/src/peft/tuners/lora.py#L163  to 

```python
in_features, out_features = target.weight.ds_shape if hasattr(target.weight,"ds_shape") else \
                            target.weight.shape
```

but the error was changed to another one like issue #122 

And i also find the current peft not supporting using fp16  with deepspeed #138 which seems that the new add parameters is not following fp16 data type. 

Resolving patch:
diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -195,7 +195,9 @@ def _setup_prompt_encoder(self):
                     self.transformer_backbone_name = name
 
         if self.peft_config.num_transformer_submodules is None:
-            self.peft_config.num_transformer_submodules = 2 if self.peft_config.task_type == TaskType.SEQ_2_SEQ_LM else 1
+            self.peft_config.num_transformer_submodules = (
+                2 if self.peft_config.task_type == TaskType.SEQ_2_SEQ_LM else 1
+            )
 
         for named_param, value in list(transformer_backbone.named_parameters()):
             if value.shape[0] == self.base_model.config.vocab_size:
@@ -733,8 +735,9 @@ def forward(
                 decoder_inputs_embeds = torch.cat(
                     (prompts[:, self.peft_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1
                 )
-                return self.base_model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs)
-
+                return self.base_model(
+                    inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs
+                )
 
     def generate(self, **kwargs):
         if not isinstance(self.peft_config, PromptLearningConfig):
diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -145,7 +145,7 @@ def _find_and_replace(self):
                     is_target_modules_in_base_model = True
                 parent, target, target_name = self._get_submodules(key)
                 bias = target.bias is not None
-                if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt) and self.peft_config.enable_lora is None:
+                if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt):
                     kwargs.update(
                         {
                             "has_fp16_weights": target.state.has_fp16_weights,
@@ -154,13 +154,19 @@ def _find_and_replace(self):
                             "index": target.index,
                         }
                     )
-                    new_module = Linear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs)
+                    if self.peft_config.enable_lora is None:
+                        new_module = Linear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs)
+                    else:
+                        kwargs.update({"enable_lora": self.peft_config.enable_lora})
+                        new_module = MergedLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs)
                 elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None:
                     new_module = Linear(target.in_features, target.out_features, bias=bias, **kwargs)
                 elif self.peft_config.enable_lora is not None:
                     kwargs.update({"enable_lora": self.peft_config.enable_lora})
                     if isinstance(target, Conv1D):
-                        in_features, out_features = target.weight.shape
+                        in_features, out_features = (
+                            target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape
+                        )
                     else:
                         in_features, out_features = target.in_features, target.out_features
                         if kwargs["fan_in_fan_out"]:
@@ -509,3 +515,82 @@ def forward(self, x: torch.Tensor):
                     output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling
                     result += output
             return result
+
+    class MergedLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer):
+        # Lora implemented in a dense layer
+        def __init__(
+            self,
+            in_features: int,
+            out_features: int,
+            r: int = 0,
+            lora_alpha: int = 1,
+            lora_dropout: float = 0.0,
+            enable_lora: List[bool] = [False],
+            **kwargs,
+        ):
+            bnb.nn.Linear8bitLt.__init__(
+                self,
+                in_features,
+                out_features,
+                bias=kwargs.get("bias", True),
+                has_fp16_weights=kwargs.get("has_fp16_weights", True),
+                memory_efficient_backward=kwargs.get("memory_efficient_backward", False),
+                threshold=kwargs.get("threshold", 0.0),
+                index=kwargs.get("index", None),
+            )
+            LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False)
+            if out_features % len(enable_lora) != 0:
+                raise ValueError("The length of enable_lora must divide out_features")
+            self.enable_lora = enable_lora
+            # Actual trainable parameters
+            if r > 0 and any(enable_lora):
+                self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False)
+                self.lora_B = nn.Conv1d(
+                    r * sum(enable_lora),
+                    out_features // len(enable_lora) * sum(enable_lora),
+                    kernel_size=1,
+                    groups=2,
+                    bias=False,
+                )
+                self.scaling = self.lora_alpha / self.r
+                # Freezing the pre-trained weight matrix
+                self.weight.requires_grad = False
+                # Compute the indices
+                self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1)
+                self.lora_ind[enable_lora, :] = True
+                self.lora_ind = self.lora_ind.view(-1)
+            self.reset_parameters()
+
+        def reset_parameters(self):
+            if hasattr(self, "lora_A"):
+                # initialize A the same way as the default for nn.Linear and B to zero
+                nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
+                nn.init.zeros_(self.lora_B.weight)
+
+        def zero_pad(self, x):
+            result = x.new_zeros((*x.shape[:-1], self.out_features))
+            result = result.view(-1, self.out_features)
+            result[:, self.lora_ind] = x.reshape(
+                -1, self.out_features // len(self.enable_lora) * sum(self.enable_lora)
+            )
+            return result.view((*x.shape[:-1], self.out_features))
+
+        def forward(self, x: torch.Tensor):
+            result = super().forward(x)
+            if self.disable_adapters:
+                return result
+            elif self.r > 0:
+                if not torch.is_autocast_enabled():
+                    expected_dtype = result.dtype
+                    if x.dtype != torch.float32:
+                        x = x.float()
+                    after_A = self.lora_A(self.lora_dropout(x))
+                    after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1)
+                    output = self.zero_pad(after_B).to(expected_dtype) * self.scaling
+                    result += output
+                else:
+                    after_A = self.lora_A(self.lora_dropout(x))
+                    after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1)
+                    output = self.zero_pad(after_B) * self.scaling
+                    result += output
+            return result
diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py
--- a/src/peft/utils/config.py
+++ b/src/peft/utils/config.py
@@ -160,6 +160,8 @@ class PromptLearningConfig(PeftConfig):
     token_dim: int = field(
         default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"}
     )
-    num_transformer_submodules: Optional[int] = field(default=None, metadata={"help": "Number of transformer submodules"})
+    num_transformer_submodules: Optional[int] = field(
+        default=None, metadata={"help": "Number of transformer submodules"}
+    )
     num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"})
     num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})


Test patch: