Cache size limit for generation

[Natooz]

Feature request

Add a cache_limit argument for generate, limiting the size of the cache (past_key_values).

Motivation

In some contexts one might want to generate long sequences. When doing so, the system can easily run out of memory. Keeping the cache to a maximum size would allow users to have more control and tweak others parameters such as batch size or number of beams, to generate faster and take the most out of their hardware.

Your contribution

I implemented it in GPT2 (PyTorch & TF, PR is ready), but I guess this could be implemented more broadly in generate so that every models could benefit it.
It might relate to #17574.
Waiting for your opinion on this, I can probably add it to generate.

[sgugger]

cc @gante but this might be a bitt too niche for us to incorporate in generate.

[gante]

Hey @Natooz 👋

Thank you for raising this issue! It may be helpful in some situations. However, since it is the first time I'm seeing a request for this, the answer depends on your implementation :) I'm on board if it consists of a few line changes (say, <100 in total for all models). Otherwise, the maintenance costs are too high.

Regardless of your answer and the decision on this issue, it's useful for us to have issues like this, to gauge demand and guide our future work 🤗

[Natooz]

Hi @gante, thanks for the feedback !

This is because the feature might be a very niche (eg story or music gen) that I asked if you would prefer it implemented per model (was thinking of prepare_inputs_for_generation) or more in a more DRY way, probably a method called in each decoding method at each step.

Here is how I implemented it for GPT2 in PyTorch (TF is pretty much the same) in prepare_inputs_for_generation, it stands in a few lines:

# past is retrieved from model_kwargs, and is of shape (L,2,N,NH,T,DH)
# (layer, keys/values, batch, attn_head, seq, dim_head)
# with two first dims as tuple, dim 2:-1 is a fixed length tensor
cache_limit = kwargs.get("cache_limit", None)

# check if using cache and a limit, and that the current cache does exceed it
# dim -2 of cache is always the same across all layers and kv, so checking past[0][0]
if past and cache_limit and past[0][0].shape[-2] > cache_limit:
    # reducing the time / seq dimension (-2)
    past = [[kv[..., -cache_limit:, :] for kv in layer] for layer in past]

    # we need to update the attention_mask as it is incremented in _update_model_kwargs_for_generation
    if attention_mask is not None:
        attention_mask = attention_mask[:, -cache_limit - 1:]
    # no need to update position_ids here as kwargs[attention_mask].shape[1]
    # is the nb of all past positions / tokens that have been processed so far

Now for implementing it in GenerationMixin, it would require that 1) all concerned models use the same cache shape (L,2,N,NH,T,DH) (is it the case?), and; 2) it does not mess with positional encoding. In GPT2, positions are created on the fly from the shape of the attention_mask, itself incremented at each decoding step.
For this direction, and if we don't want to touch any method overridden by models, I think the best solution would be to store
position_id in model_kwargs, that would be updated at each step.

I'll give a try to the GenerationMixin direction and come back ! 👨‍💻

[github-actions]

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 are likely to be ignored.