[cache] add a test to confirm we can use cache at train time

[gante]

What does this PR do?

In #35648, the idea of always disabling the cache at train-time was floated. However, @BenjaminBossan pointed to me that some fine-tuning scripts need cache at train time (e.g. prefix tuning in PEFT).

I couldn't find any test that ensured the survival of this property, so this PR adds a test that will fail if we decide to disable it :)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[BenjaminBossan]

Thanks a lot for adding this test to prevent possible regressions with PEFT methods like prefix tuning. I wonder if we can also try injecting the cache and checking if that works, basically like a mini prefix tuning:

        # simulate injecting virtual tokens like in prefix tuning
        num_virtual_tokens = 3
        past_key_values = [torch.randn(2, 1, 2, num_virtual_tokens, 8)] * 2
        past_key_values = DynamicCache.from_legacy_cache(past_key_values)
        model_inputs["attention_mask"] = torch.cat((
            model_inputs["attention_mask"],
            torch.ones(1, num_virtual_tokens).to(model_inputs["attention_mask"].device)
        ), dim=1)
        model_outputs = model(**model_inputs, past_key_values=past_key_values, use_cache=True)

Maybe this is overkill, LMK what you think.

[gante]

@BenjaminBossan added, thank you for the suggestion 👍 (the extra check is nearly free, and may save some pain down the road)

[BenjaminBossan]

Thanks so much for taking of ensuring that this works, LGTM.

[ArthurZucker]

Okay! I think we did break that once or twice, good addition! 🤗