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 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 bd5091d 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 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 :) )