Idefics2 raises shape error in the backward pass with gradient_checkpointing

[chenzizhao]

I get a shape error if I PEFT the latest Idefics2 model at head with gradient checkpointing. The symptom is identical to this. Seems like the following block isn't compatible with gradient checkpointing.

transformers/src/transformers/models/idefics2/modeling_idefics2.py

Lines 1593 to 1596 in e4ea19b

	if use_cache:
	if not isinstance(past_key_values, Cache):
	past_key_values = DynamicCache.from_legacy_cache(past_key_values)
	past_seen_tokens = past_key_values.get_usable_length(seq_length)

In my local workspace a fix similar to 91f4b7e removed the shape error but a quick review of the logic from a maintainer would be lovely. Happy to submit PR after. Thanks!

[chenzizhao]

Patch: a317323

[amyeroberts]

Hi @chenzizhao, thanks for raising this and providing links to the patch and relevant commits!

Would you like to open a PR with these changes? This way you get the github contribution for the solution.

cc @younesbelkada @gante for reference

[younesbelkada]

Nice catch @chenzizhao !
This is not catched by our CI sadly, we need to work on that ! adding it on my TODOs !