Bugfix for generation with an early-stopping process
Author: ojh31
Pull request: yes

# What does this PR do?
Fixes a [bug](https://github.com/huggingface/transformers/issues/32603) where we get a shape mismatch from the forward call to a model which uses key-value caching if one of the processes finishes early during generation.

Fixes https://github.com/huggingface/transformers/issues/32603

## 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#create-a-pull-request),
      Pull Request section?
- [x] 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
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc

Documentation: @stevhliu

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:
Failing tests seem to be related to versioning issues with CircleCI runners so marking this as ready for review anyway

@ojh31 Could you try rebasing on main? This should resolve the failing tests. 

cc @gante for review

Hi @ojh31 👋 Thank you for opening #32603 and this PR!

I see, multi-gpu + new cache changes likely brought problems. I may request changes to this PR, but let me first sync with our `transformers` multi-device master!
___________________________

@SunMarc are the current shape-related problems of FSDP + `generate` on your radar? Is the snippet in the [original issue](https://github.com/huggingface/transformers/issues/32603) the intended usage for FSDP + `generate`? 

This PR fixes the issue by removing the `continue` when `synced_gpus=True`, which makes me believe the original cause is a mismatch between the cache size (which gets updated in `forward`) and its indexing tensor `cache_positions` (which does NOT get updated because of the `continue`)

Hey @ojh31, usually for generation, we need to first unwrap the model. This is what is done in TRL for example : https://github.com/huggingface/trl/blob/54f806b6ffdfa49f584340aec18d079a58a3a342/trl/trainer/online_dpo_trainer.py#L536. Can you try that and see if you still get the error ? Also did it worked on past version of transformers ? 

@gante , we never had issues with FSDP + generate in the past. I'll investigate a bit ! If we have this issue in FSDP, we should also have for deepspeed. 

> Hey @ojh31, usually for generation, we need to first unwrap the model. This is what is done in TRL for example : https://github.com/huggingface/trl/blob/54f806b6ffdfa49f584340aec18d079a58a3a342/trl/trainer/online_dpo_trainer.py#L536. Can you try that and see if you still get the error ? Also did it worked on past version of transformers ?

Thanks for taking a look at this @SunMarc! I get the same error modifying the last block as follows:
```
@contextmanager
def unwrap_model_for_generation(
    model: PreTrainedModel, accelerator: "Accelerator"
):
    yield accelerator.unwrap_model(model)


for batch_input_ids, batch_attention_mask in dataloader:
    with torch.no_grad():
        model.forward(input_ids=batch_input_ids, attention_mask=batch_attention_mask)
    with FSDP.summon_full_params(model, recurse=False):
        with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
            outputs = unwrapped_model.generate(
                input_ids=batch_input_ids,
                attention_mask=batch_attention_mask, 
                tokenizer=tokenizer,
                synced_gpus=True,
            )
```
I get the same error on transformers 4.42.4 and 4.44.0. However, I just randomly tried 4.35 and that did not hit the same error!

git bisect finds that bd5091df8db7cea1a9f94f797fc11487f840ade1 is the breaking commit

Wow thanks for testing and finding the commit that created this issue @ojh31 ! Do you have any insights regarding this commit and the error that we are facing @gante ? Is this linked to your intuition of why this is not working ? 

I am hitting the same issue with deepspeed zero3 on 4.44.0. For me, the mismatch is between the `attention_mask` (which is being updated after the `continue`) and the `key_states`/`value_states` (updated during forward). 

(see [this](https://github.com/huggingface/transformers/issues/32603#issuecomment-2324563712) comment in the original issue to avoid parallel discussions 🤗 )

with #34095, this PR should not be needed 🤗 

(please reopen and ping me if the original issue is not fixed :) )