When using HF trainer + PEFT + DeepSpeed ZeRO3, there's only hacky way to save the base model
Author: yundai424
Pull request: no

### System Info

- `transformers` version: 4.34.1
- Platform: Linux-3.10.0-1160.102.1.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.10.12
- Huggingface_hub version: 0.17.1
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0+cu121 (False)
- 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

### Who can help?

deepspeed: HF Trainer/Accelerate: @pacman100

### 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

train any peft model with DeepSpeed ZeRO3. For simplicity this is a sample script that uses `trl.SFTTrainer` to fine tune on alpaca dataset:

```python
import transformers
import datasets
import trl
from dataclasses import dataclass, field
import peft
import torch
import time
import callbacks


@dataclass
class CustomArguments():
    model_path: str
    data_path: str = field(default="alpaca_data.json")
    max_seq_length: int = field(default=512)
    lora: bool = field(default=False)
    lora_r: int = field(default=8)
    lora_alpha: int = field(default=32)
    lora_dropout: float = field(default=0.1)


def formatting_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        output_texts.append(f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{example["instruction"][i]}

### Input:
{example["input"][i]}

### Response:
{example["output"][i]}""")
    return output_texts



def main():
    parser = transformers.HfArgumentParser((transformers.TrainingArguments, CustomArguments))
    training_args, custom_args = parser.parse_args_into_dataclasses()

    dataset = datasets.load_dataset("json", data_files=custom_args.data_path, split=[
        'train'])[0].train_test_split(test_size=0.2, shuffle=True)
    dataset_train, dataset_eval = dataset['train'], dataset['test']

    if torch.distributed.get_rank() == 0:
        print(custom_args, training_args)

    model = transformers.AutoModelForCausalLM.from_pretrained(custom_args.model_path,
            trust_remote_code=True,
            use_cache=False)

    peft_config = peft.LoraConfig(task_type=peft.TaskType.CAUSAL_LM,
            inference_mode=False,
            r=custom_args.lora_r,
            lora_alpha=custom_args.lora_alpha,
            lora_dropout=custom_args.lora_dropout) if custom_args.lora else None

    tokenizer = transformers.AutoTokenizer.from_pretrained(custom_args.model_path)
    tokenizer.pad_token = tokenizer.eos_token

    trainer = trl.SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset_train,
        eval_dataset=dataset_eval,
        formatting_func=formatting_func,
        max_seq_length=custom_args.max_seq_length,
        peft_config=peft_config,
        args=training_args,
    )

    trainer.train()
    tokenizer.save_pretrained(training_args.output_dir)
    trainer.save_model()
    model.base_model.save_pretrained(training_args.output_dir)

if __name__ == "__main__":
    main()
```

If we run it with DS ZeRO3, this line `model.base_model.save_pretrained(training_args.output_dir)` will only save out a tiny model bin (several megabytes). I believe that's because when initializing with ZeRO3, the original model's `state_dict` will include mostly just [empty torch tensor](https://github.com/microsoft/DeepSpeed/blob/9dfb06de36bb29293b1e94dc1e48d6f2adf54d2c/deepspeed/runtime/zero/partition_parameters.py#L534) except for layer norm stats, while the actual weights are placed in deepspeed internal variables. 

There's certainly a hack to work around:
```python
    state_dict = trainer.accelerator.get_state_dict(trainer.deepspeed)
    import collections
    if torch.distributed.get_rank() == 0:
        renamed_state_dict = collections.OrderedDict()
        for key, value in state_dict.items():
            renamed_state_dict[key.lstrip("base_model.model.")] = value
        model.base_model.save_pretrained(training_args.output_dir", state_dict=renamed_state_dict)
```
to get the original full state dict of the lora model by `trainer.accelerator.get_state_dict(trainer.deepspeed)`, and then rename the keys in the state dict to trim the `base_model.model.` path so that the base model can be accessed. 


(We understand that the base model is not necessary -- as long as the PEFT model is saved, once it's reloaded, the base model specified in the `"base_model_name_or_path"` in `adapter_config.json` will be loaded into memory as well. However for our production cases the base model path is not always there and may be cleaned up from time to time so a copy of base model is needed along with the peft model)

### Expected behavior

There may be a public API to get the full state dict of the base model that will work with ZeRO3, might be implemented in a way similar to the hack we're having now 

Comments:
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.

Gently pinging @pacman100 and @younesbelkada 

i will let @pacman100 answer on this as he is more familiar than I am on DS, @yundai424 can you try on latest transformers and peft ? `pip install -U transformers peft`

@younesbelkada that gives the same result. Though I found https://huggingface.co/docs/accelerate/usage_guides/deepspeed#saving-and-loading gives a more official solution (i.e. using accelerate `unwrap_model` so if this is the ultimate solution i can close this issue :D

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.