fix prompt tunning + deepspeed zero3 + checkpoint_saving hang issue
Author: sywangyi
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.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      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.

<!-- 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:
Hi @younesbelkada. When I use deepspeed zero3 and prompt tuning to finetune large model. find training hang after saving checkpoint.
 the prompt-tuning has prompt_encoder forward in model save_pretrained. only rank0 call the prompt_encoder forward in deepspeed zero3 is not enough. all the rank should call prompt_encoder forward. 

cc @pacman100 for deepspeed! 

yes, I think should_save should be move to _save. instead of controlling if _save is called like _save_tpu, WDYT?

any thought about my proposal? I think should_save should be move to _save. instead of controlling if _save is called like _save_tpu. @pacman100 

@pacman100 , @ArthurZucker what's your suggestion to move on for this issue? any comments on Yi's proposal?

how about this? @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.

cc @muellerzr if you can take this review over! 🤗 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29980). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> the prompt-tuning has prompt_encoder forward in model save_pretrained. only rank0 call the prompt_encoder forward in deepspeed zero3 is not enough. all the rank should call prompt_encoder forward.

thanks for the review!  the root cause of this issue is that the prompt-tuning has prompt_encoder forward in model **save_pretrained**. only rank0 call the prompt_encoder forward in deepspeed zero3 is not enough. all the rank should call prompt_encoder forward. or else it will hang...    so **save_pretrained** in **_save** should be called in all the rank. however. self.args.should_save is only true when local_process_index=0. so I move it to **_save**.

Again though: this is deepspeed specific based on model sharding/splitting, so we should only modify the code for deepspeed specifically, unless there is a reason not to. (Aka, deepspeed and fsdp specific I imagine).

Am I incorrect in this understanding?

> Again though: this is deepspeed specific based on model sharding/splitting, so we should only modify the code for deepspeed specifically, unless there is a reason not to. (Aka, deepspeed and fsdp specific I imagine).
> 
> Am I incorrect in this understanding?

yes, deepspeed zero3 and FSDP specific

Great, let's adjust the logic for that case specifically then, and leave the already existing base logic in-tact then

@muellerzr  I have updated the PR, could you revisit it?

> Thanks for adding this, but having this fix indicates that there's likely something wrong in how we control our saving logic more generally.
> 
> Having to have lots of `should_save` checks within the `_save` method is odd: the `_save` method should only be called if `should_save` is true OR, at worst, the `_save` method should be directly exited at the start of the call, we shouldn't have to have the whole method's logic conditioned on this flag

the main objective is that the save_pretrained in `_save` should be called in every rank in FSDP, DEEPSPEED+prompt tuning case. so, '_save' method could not be directly exited at the start of the call or controlled by `should_save`.
@amyeroberts I just fix the comment. 

> @sywangyi Thanks for the explanation. I understand the intended logic. My previous comment still stands: we shouldn't need to condition so much of the `_save` logic on the `should_save` flag, it's OK to pass in `is_main_process`, but it doesn't really make sense for saving e.g. the tokenizer. Instead, I would look at:
> 
> * How we define the main process - we might need a different flag e.g. `self._is_main_process`
> * Update `should_save` to be true for all ranks in FSDP, DEEPSPEED+prompt tuning case
> 
> cc @muellerzr Does this seem reasonable?

_save_tpu has the similar logic. use should_save flag internal.

also. in `save_pretrained` (https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L2503 or https://github.com/huggingface/peft/blob/main/src/peft/peft_model.py#L294) is_main_process is equal to `should_save`,`should_save` is controlled by https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py#L2322. not equal to `process_index` == 0. If you see `_save_tpu` logic, it's what I would like to imitate. calling logic is https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L3420. 
function logic is like https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L3460. `should_save` is put in the save function internally.

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.

@amyeroberts , do you have any further comments on this? Thx.

@yao-matrix Beyond [my comment here](https://github.com/huggingface/transformers/pull/29980#pullrequestreview-2161152984) and @muellerzr's [comment here](https://github.com/huggingface/transformers/pull/29980#discussion_r1699045322), I don't have any else to add, but happy to clarify my point in any way or help iterate on a solution