Skip saving frozen parameters if using peft model with deepspeed
Author: VeryLazyBoy
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

Currently, when using `PeftModel`, `transformers` only save the weights of the adapter and support resuming training from these saved weights. However, if `deepspeed` is used on top of `PeftModel`, the entire model weights are saved. This behavior differs from that of `PeftModel`.

This PR integrates a newly added parameter `exclude_frozen_weights` from deepspeed to skip saving frozen weights if using peft.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@pacman100

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


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.

Gentle ping @pacman100 

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.

I believe this is resolved in #27825 

@amyeroberts Hi Amy, you are right. We implemented a similar solution to save only a portion of the weights. However, I'm afraid that issue could arise when attempting to load the weights again, unless the `load_module_strict` parameter is set to False, as demonstrated in this pull request.

```diff
# @@ def deepspeed_load_checkpoint(deepspeed_engine, checkpoint_path):
if len(deepspeed_checkpoint_dirs) > 0:
        logger.info(f"Attempting to resume from {checkpoint_path}")
+
+       load_module_strict = True
+        if version.parse(deepspeed_version) > version.parse("0.10.0"):
+           if is_peft_available() and isinstance(deepspeed_engine.module, PeftModel):
+              load_module_strict = False
        # this magically updates self.optimizer and self.lr_scheduler
        load_path, _ = deepspeed_engine.load_checkpoint(
-        checkpoint_path, load_optimizer_states=True, load_lr_scheduler_states=True
+            checkpoint_path,
+            load_optimizer_states=True,
+            load_lr_scheduler_states=True,
+            load_module_strict=load_module_strict,
        )

```

@VeryLazyBoy Thanks for pointing out! @muellerzr @pacman100 could you give this a first review? 

I think I am running into this issue when resuming,
```
 File "/data/chirag/llm-finetune/train.py", line 883, in _train
    trainer.train(resume_from_checkpoint=last_checkpoint_dir)
  File "/data/v/ft/lib/python3.10/site-packages/transformers/trainer.py", line 1543, in train
    return inner_training_loop(
  File "/data/v/ft/lib/python3.10/site-packages/transformers/trainer.py", line 1699, in _inner_training_loop
    deepspeed_load_checkpoint(self.model_wrapped, resume_from_checkpoint)
  File "/data/v/ft/lib/python3.10/site-packages/transformers/integrations/deepspeed.py", line 402, in deepspeed_load_checkpoint
    load_path, _ = deepspeed_engine.load_checkpoint(
  File "/data/v/ft/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2725, in load_checkpoint
    load_path, client_states = self._load_checkpoint(load_dir,
  File "/data/v/ft/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2795, in _load_checkpoint
    self.load_module_state_dict(checkpoint=checkpoint,
  File "/data/v/ft/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2588, in load_module_state_dict
    self.module.load_state_dict(
  File "/data/v/ft/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2152, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for PeftModelForCausalLM:
        Missing key(s) in state_dict: "base_model.model.model.embed_tokens.weight", "base_model.model.model.layers.0.self_attn.q_proj.weight", ...
```

The `load_module_strict=False` is necessary here to load from peft and resume. I tested it ad-hoc on top of main

> @amyeroberts Hi Amy, you are right. We implemented a similar solution to save only a portion of the weights. However, I'm afraid that issue could arise when attempting to load the weights again, unless the `load_module_strict` parameter is set to False, as demonstrated in this pull request.
> 
> ```diff
> # @@ def deepspeed_load_checkpoint(deepspeed_engine, checkpoint_path):
> if len(deepspeed_checkpoint_dirs) > 0:
>         logger.info(f"Attempting to resume from {checkpoint_path}")
> +
> +       load_module_strict = True
> +        if version.parse(deepspeed_version) > version.parse("0.10.0"):
> +           if is_peft_available() and isinstance(deepspeed_engine.module, PeftModel):
> +              load_module_strict = False
>         # this magically updates self.optimizer and self.lr_scheduler
>         load_path, _ = deepspeed_engine.load_checkpoint(
> -        checkpoint_path, load_optimizer_states=True, load_lr_scheduler_states=True
> +            checkpoint_path,
> +            load_optimizer_states=True,
> +            load_lr_scheduler_states=True,
> +            load_module_strict=load_module_strict,
>         )
> ```

I encountered the same issue previously. Moreover, when resuming from the checkpoint in Deepspeed (with `DeepSpeedEngine.load_checkpoint` you mentioned), some objects like the lr_scheduler object are instantiated from scratch which means the lr_scheduler cycle in fact does not resume and starts again, as Deepspeed does not save the lr_scheduler at all. This is because Deepspeed [does not support](https://github.com/microsoft/DeepSpeed/blob/c37fe9cbfb8bc10c8dd6ccd8cac9b34ded218990/deepspeed/runtime/lr_schedules.py#L23) default lr_scheduler type used in transformers i.e. `LambdaLR`.

Currently, Transformers delegates checkpointing to Deepspeed if it is enabled, and in that scenario lr_scheduler is not saved and loaded in checkpoints. I tweaked the Transformers a bit to handle this scenario with minimal changes in this [commit](https://github.com/huggingface/transformers/commit/4acce7a43be3aa1d29d32c353ca2d15dd27bddf0).

> Deepspeed does not save the lr_scheduler at all. This is because Deepspeed [does not support](https://github.com/microsoft/DeepSpeed/blob/c37fe9cbfb8bc10c8dd6ccd8cac9b34ded218990/deepspeed/runtime/lr_schedules.py#L23) default lr_scheduler type used in transformers i.e. LambdaLR

Hi @kazemf78 , have you checked this https://github.com/huggingface/transformers/pull/25863? By the way, the issue you mentioned is not related to this PR, if there indeed exist a bug, you can create a separate issue or pr :)

The `load_module_strict` is a required change for resuming to work correctly,  @amyeroberts can we please get that merged?
Is the blocker that partial work of this PR is already merged? Should we create another PR quickly?


cc @pacman100 for first review 

Hello, Thank you for @VeryLazyBoy for your PR. I'm extremely sorry and apologize that this PR got missed. However, this issue is fixed in https://github.com/huggingface/transformers/pull/28746. 

Thank you for providing this information. I am pleased to hear that the issue has been resolved. @pacman100