Lora + DeepSpeed non-trainer integration does not work
Author: hrushikesh198
Pull request: no

### System Info

- `transformers` version: 4.37.2
- Platform: Linux-4.19.0-24-cloud-amd64-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.26.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.2 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes, trying to run deepspeed zero 3 on 8 A100 gpus

### Who can help?

cc: @pacman100

Tagging few folks who were discussing a similar issue before https://github.com/huggingface/transformers/issues/24445 
@1ytic, @don-tpanic, @olegsinavski 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I am trying to finetune Mistral-7B using LoRa, deepspeed zero 3, pytorch-lightning. As per the deepspeed non trainer integration, I have created the `dschf` object and kept it alive. 

```python
class Module(LightningModule):
    def configure_model(self) -> None:
        if self.model is not None:
            return 
        deepspeed_config = self.trainer.strategy.config
        self.dschf = HfDeepSpeedConfig(deepspeed_config)
        self.model = AutoModelForSequenceClassification.from_pretrained(...)
        self.model = get_peft_model(
            self.model,
            LoraConfig(
                task_type=TaskType.SEQ_CLS,
                inference_mode=False,
                target_modules=target_modules,
                r=256,
                lora_alpha=256,
                lora_dropout=0.5,
            ),
        )

def main():
    trainer = lightning.Trainer(
        ...
        strategy = DeepSpeedStrategy(stage=3)
    )
    model=Module()
    trainer.fit(model, datamodules)
```

The training script throws an error on the `get_peft_model` line

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/peft/peft_model.py", line 528, in __getattr__
    return super().__getattr__(name)  # defer to nn.Module's logic
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'PeftModelForSequenceClassification' object has no attribute '_ds_child_entered'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/peft/peft_model.py", line 528, in __getattr__
    return super().__getattr__(name)  # defer to nn.Module's logic
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'PeftModelForSequenceClassification' object has no attribute 'base_model'```
```
the seconds one continues until it reaches max recursion depth.


### Expected behavior

Lora model should initialize seamlessly and train as it works for deepspeed stage 2.

Comments:
cc @younesbelkada 

Hello @hrushikesh198 ,

PR https://github.com/huggingface/peft/pull/1450 fixes support for DeepSpeed Z3 with zero init and `modules_to_save` config of PEFT. In the PR, we show how you can run non trainer example with Accelerate and normal training loop with an official example. 

I don't have experience with Lightning, we could look into this if you can provide a minimal but end-to-end example that we can run. Above code is missing a lot of things to just run it as is.

Hi @pacman100 ,
Thank you for the attempt at fixing the DS3 issue. 
I tried installing the latest released transformer==4.38.1, peft based on latest github commit 1b3b7b5, and run the following script for loading phi-2 model using DeepSpeed3 on 2-A100 gpus. It still gives the same error.

Here are the training script and it's output:

<details>
<summary> training.py </summary>

```python
# Python version 3.10.10
# pip install torch==2.1.2 lightning==2.1.4 deepspeed==0.13.2 torchmetrics \
#     transformers==4.38.1
# pip install -U git+https://github.com/huggingface/peft@1b3b7b5 # latest on 2/23/24

import json
from functools import partial

import lightning as L
import peft
import torch
import torch.utils.data
from lightning_utilities.core.rank_zero import rank_zero_info
from transformers import AutoModelForCausalLM, PreTrainedTokenizer, AutoTokenizer
from transformers.integrations import HfDeepSpeedConfig

train_data = [
                 "Transformers is a media franchise produced by American toy company Hasbro",
                 "and Japanese toy company Takara Tomy.",
                 "It primarily follows the heroic Autobots",
                 "and the villainous Decepticons,",
                 "two alien robot factions at war that can transform into other forms,",
                 "such as vehicles and animals.",
             ] * 200

val_data = [
               "The franchise encompasses toys, animation,",
               "comic books, video games and films.",
               "As of 2011, it generated more than ¥2 trillion ($25 billion) in revenue"
           ] * 80


class TextDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index]


class LitModel(L.LightningModule):
    def __init__(self, pretrain_model_name, pad_token_id):
        super().__init__()
        self.pretrain_model_name = pretrain_model_name
        self.pad_token_id = pad_token_id
        self.hf_deepspeed_cfg = None
        self.model = None

    def configure_model(self):
        if self.model is not None:
            return

        ds_config = self.trainer.strategy.config
        rank_zero_info(f"DeepSpeed cfg={json.dumps(ds_config, indent=2)}")
        self.hf_deepspeed_cfg = HfDeepSpeedConfig(ds_config)  # keep this object alive

        self.model = AutoModelForCausalLM.from_pretrained(
            self.pretrain_model_name,
            pad_token_id=self.pad_token_id,
        )

        rank_zero_info(f"Pretrained Model >> {self.model}")

        self.model = peft.get_peft_model(self.model, peft.LoraConfig(
            task_type=peft.TaskType.CAUSAL_LM, r=128, lora_alpha=128, lora_dropout=0.1, target_modules=[
                'q_proj', 'k_proj', 'v_proj', 'dense', 'fc1', 'fc2', 'lm_head']))

        rank_zero_info(f"LoRA Model >> {self.model}")

    def configure_optimizers(self):
        return torch.optim.AdamW(self.trainer.model.parameters(), fused=True)

    def _step(self, batch, batch_idx, prefix):
        res = self.model.forward(**batch, labels=batch['input_ids'])
        self.log(f"{prefix}/loss", res[0], prog_bar=True)
        return res[0]

    def training_step(self, batch, batch_idx):
        return self._step(batch, batch_idx, 'train')

    def validation_step(self, batch, batch_idx):
        return self._step(batch, batch_idx, 'val')


def train():
    torch.set_float32_matmul_precision('high')

    strategy = "deepspeed_stage_3"
    trainer = L.Trainer(
        accelerator='gpu', devices=2, precision='32', strategy=strategy, max_epochs=2,
        num_sanity_val_steps=0)

    pretrain_model_name = "microsoft/phi-2"
    tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(pretrain_model_name)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
    tokenizer.truncation_side = 'right'
    tokenizer.padding_side = 'left'
    model = LitModel(pretrain_model_name, tokenizer.pad_token_id)

    collate = partial(tokenizer, return_tensors='pt', padding='longest', )
    train_dl = torch.utils.data.DataLoader(TextDataset(train_data), batch_size=2, collate_fn=collate, drop_last=True)
    val_dl = torch.utils.data.DataLoader(TextDataset(val_data), batch_size=2, collate_fn=collate, drop_last=True)

    trainer.fit(model, train_dl, val_dl)


if __name__ == '__main__':
    train()
```
</details>

<details>
<summary>training.log</summary>

```
[2024-02-23 21:09:29,721] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
initializing deepspeed distributed: GLOBAL_RANK: 0, MEMBER: 1/2
[2024-02-23 21:09:41,382] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
initializing deepspeed distributed: GLOBAL_RANK: 1, MEMBER: 2/2
DeepSpeed cfg={
  "zero_allow_untested_optimizer": true,
  "zero_optimization": {
    "stage": 3,
    "contiguous_gradients": true,
    "overlap_comm": true,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "allgather_bucket_size": 200000000,
    "reduce_bucket_size": 200000000,
    "sub_group_size": 1000000000000
  },
  "activation_checkpointing": {
    "partition_activations": false,
    "cpu_checkpointing": false,
    "contiguous_memory_optimization": false,
    "synchronize_checkpoint_boundary": false
  },
  "aio": {
    "block_size": 1048576,
    "queue_depth": 8,
    "single_submit": false,
    "overlap_events": true,
    "thread_count": 1
  },
  "gradient_accumulation_steps": 1,
  "train_micro_batch_size_per_gpu": 2,
  "gradient_clipping": 0.0
}
Loading checkpoint shards: 100%|   | 2/2 [00:03<00:00,  1.89s/it]
Loading checkpoint shards: 100%|   | 2/2 [00:03<00:00,  1.98s/it]
Pretrained Model >> PhiForCausalLM(
  (model): PhiModel(
    (embed_tokens): Embedding(51200, 2560, padding_idx=50256)
    (embed_dropout): Dropout(p=0.0, inplace=False)
    (layers): ModuleList(
      (0-31): 32 x PhiDecoderLayer(
        (self_attn): PhiAttention(
          (q_proj): Linear(in_features=2560, out_features=2560, bias=True)
          (k_proj): Linear(in_features=2560, out_features=2560, bias=True)
          (v_proj): Linear(in_features=2560, out_features=2560, bias=True)
          (dense): Linear(in_features=2560, out_features=2560, bias=True)
          (rotary_emb): PhiRotaryEmbedding()
        )
        (mlp): PhiMLP(
          (activation_fn): NewGELUActivation()
          (fc1): Linear(in_features=2560, out_features=10240, bias=True)
          (fc2): Linear(in_features=10240, out_features=2560, bias=True)
        )
        (input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
        (resid_dropout): Dropout(p=0.1, inplace=False)
      )
    )
    (final_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=2560, out_features=51200, bias=True)
)
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/peft/peft_model.py", line 529, in __getattr__
    return super().__getattr__(name)  # defer to nn.Module's logic
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'PeftModelForCausalLM' object has no attribute '_ds_child_entered'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/peft/peft_model.py", line 529, in __getattr__
    return super().__getattr__(name)  # defer to nn.Module's logic
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'PeftModelForCausalLM' object has no attribute 'base_model'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/peft/peft_model.py", line 529, in __getattr__
    return super().__getattr__(name)  # defer to nn.Module's logic
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'PeftModelForCausalLM' object has no attribute 'base_model'
....

Goes on until it hits max recursion depth
```
</details>

Here is a workaround: use a CustomDeepSpeedStrategy class that does not call `deepspeed.zero.Init()`.

Complete working example:
<details>
<summary>trainer.py</summary>

```python
# Python version 3.10.10
# pip install torch==2.1.2 \
#     lightning==2.1.4 deepspeed==0.13.2 torchmetrics
# pip install -U transformers==4.38.1
# pip install -U git+https://github.com/huggingface/peft@1b3b7b5 # latest on 2/23/24

import json
from contextlib import contextmanager
from functools import partial
from typing import Generator

import lightning as L
import peft
import torch
import torch.utils.data
from lightning.pytorch.strategies import DeepSpeedStrategy
from lightning_utilities.core.rank_zero import rank_zero_info
from transformers import AutoModelForCausalLM, PreTrainedTokenizer, AutoTokenizer
from transformers.integrations import HfDeepSpeedConfig


class CustomDeepSpeedStrategy(DeepSpeedStrategy):
    @contextmanager
    def model_sharded_context(self) -> Generator[None, None, None]:
        # lightning DeepSpeedStrategy used to call deepspeed.zero.Init() here
        yield



train_data = [
                 "Transformers is a media franchise produced by American toy company Hasbro",
                 "and Japanese toy company Takara Tomy.",
                 "It primarily follows the heroic Autobots",
                 "and the villainous Decepticons,",
                 "two alien robot factions at war that can transform into other forms,",
                 "such as vehicles and animals.",
             ] * 200

val_data = [
               "The franchise encompasses toys, animation,",
               "comic books, video games and films.",
               "As of 2011, it generated more than ¥2 trillion ($25 billion) in revenue"
           ] * 80


class TextDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index]


class LitModel(L.LightningModule):
    def __init__(self, pretrain_model_name, pad_token_id):
        super().__init__()
        self.pretrain_model_name = pretrain_model_name
        self.pad_token_id = pad_token_id
        self.hf_deepspeed_cfg = None
        self.model = None

    def configure_model(self):
        if self.model is not None:
            return
        #
        ds_config = self.trainer.strategy.config
        rank_zero_info(f"DeepSpeed cfg={json.dumps(ds_config, indent=2)}")
        self.hf_deepspeed_cfg = HfDeepSpeedConfig(ds_config)  # keep this object alive

        self.model = AutoModelForCausalLM.from_pretrained(
            self.pretrain_model_name,
            pad_token_id=self.pad_token_id,
        )

        rank_zero_info(f"Pretrained Model >> {self.model}")

        self.model = peft.get_peft_model(self.model, peft.LoraConfig(
            task_type=peft.TaskType.CAUSAL_LM, r=128, lora_alpha=128, lora_dropout=0.1, target_modules=[
                'q_proj', 'k_proj', 'v_proj', 'dense', 'fc1', 'fc2', 'lm_head']))

        rank_zero_info(f"LoRA Model >> {self.model}")

    def configure_optimizers(self):
        return torch.optim.AdamW(self.trainer.model.parameters(), fused=True)

    def _step(self, batch, batch_idx, prefix):
        res = self.model.forward(**batch, labels=batch['input_ids'])
        self.log(f"{prefix}/loss", res[0], prog_bar=True)
        return res[0]

    def training_step(self, batch, batch_idx):
        return self._step(batch, batch_idx, 'train')

    def validation_step(self, batch, batch_idx):
        return self._step(batch, batch_idx, 'val')


def train():
    torch.set_float32_matmul_precision('high')
    batch_size = 2
    strategy = CustomDeepSpeedStrategy(stage=3, logging_batch_size_per_gpu=batch_size)
    trainer = L.Trainer(
        accelerator='gpu', devices=4, precision='32', strategy=strategy, max_epochs=2,
        num_sanity_val_steps=0)

    pretrain_model_name = "microsoft/phi-2"
    tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(pretrain_model_name)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
    tokenizer.truncation_side = 'right'
    tokenizer.padding_side = 'left'
    model = LitModel(pretrain_model_name, tokenizer.pad_token_id)
    collate = partial(tokenizer, return_tensors='pt', padding='longest', )
    train_dl = torch.utils.data.DataLoader(
        TextDataset(train_data), batch_size=batch_size, collate_fn=collate, drop_last=True)
    val_dl = torch.utils.data.DataLoader(
        TextDataset(val_data), batch_size=batch_size, collate_fn=collate, drop_last=True)

    trainer.fit(model, train_dl, val_dl)


if __name__ == '__main__':
    train()
```
</details>

Looks like the ds3 issues are coming from multiple zero init calls. I am not sure which library needs to solve this Huggingface or Lightning or Deepspeed, it is causing lot of trouble for the users :( 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

It actually works now with PyTorch lightning