Sink cache: crop if exceeds max cache len#32315
Conversation
|
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. |
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Could we add a SinkCache test to make sure this does not get broken again?
Also, SinkCache is kind of like the sliding window cache, wondering why we don;t use this:
# SlidingWindowCache
if using_sliding_window_cache:
target_length = max(sequence_length, self.config.sliding_window)
which we have in mistral!
|
Yes, we could but we would have to also handle cropping attention masks within Oke, I will port cropping into |
|
But the sequence length is given and we create the attention mask based off sequence_length, target_length. causal_mask = torch.full(
(sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
)
exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
if self.config.sliding_window is not None:
if not using_sliding_window_cache or sequence_length > self.config.sliding_window:
exclude_mask.bitwise_or_(
torch.arange(target_length, device=device)
<= (cache_position.reshape(-1, 1) - self.config.sliding_window)
) |
|
now, if slicing the mask is the only thing we need to use potentially siding window attention (or SinkCache) on models, LGTM! |
|
(It means we can refactor the mistral piece of code as well) |
|
Yes, cropping right now is the only way to make SinkCache happy. I just realized moving the crop to I propose to leave it in |
We have a couple slow sink cache tests failing 🙃 I'd argue that we have so many failing slow tests that we stopped paying attention to new failures (myself included) :P |
gante
left a comment
gante
left a comment
There was a problem hiding this comment.
I think the current changes will add more fuel to the fire, rather than untangle the use of the different Cache classes. Added a comment below explaining my vision for this problem, for discussion 🤗
| logger.warning_once( | ||
| "`get_max_length` for SinkCache is deprecated and will return 0 after v4.47. ", | ||
| "Use past_key_values.window_length to get the maximum window length of the cache.", | ||
| ) |
There was a problem hiding this comment.
I dislike this change -- self.window_length is the maximum length of the cache, returning 0 (or None) would not be correct. I realize the use of the method is currently inconsistent throughout our cache classes.
We should add two new attributes to the Cache classes, to fully specify a cache:
- a boolean attribute to determine whether the cache grows or has a static shape (discussed in other PRs, to be used in compilation checks)
- a boolean to determine whether it uses a sliding window (to be used e.g. in this PR)
With get_max_length and these two attributes we could be able to define all types of cache, and control the modeling code accordingly 🤗
There was a problem hiding this comment.
Hmm, makes sense. I used the none following SlidingwindownCache which actually is almost same as SinkCache and it does not return any max length. So we need two attributes on each cache class: _is_sliding and _is_static right? What about HybridCache, will it be sliding or not?
There was a problem hiding this comment.
Those two attribute names look good 👍 (btw, probably should be done in a separate PR, for better future reference)
HybridCache is a good corner case indeed 🤔 I think we can treat it as a NON sliding cache, as common model variables (e.g. attention_mask) need to be prepared according to the non-sliding cache layers.
There was a problem hiding this comment.
Got you! I will add those two attributes in a second PR.
Just so we're on the same page, even after adding attributes we'll have to crop inputs to the max sliding window length. I guess the attr will make checks look nicer, like if cache._is_static: length = cache.get_max_length()
There was a problem hiding this comment.
Yes, on the same page as you: we need to crop things with sliding windows :)
|
close as sink cache will not be maintained anymore |
4 participants
What does this PR do?
Fixes #32233. Previous PR completely broke SinkCache and now it is not working as supposed. This PR reverts the change regarding
max_cache_length, because for sink cache we need to get the latestmax lengthmasks and have position ida not exceeding themax lengthI am aware that SinkCache is buggy and doesn't work for some edge cases and for multi-turn dialogues. But I think we can try to keep the previous working version for one-turn generation, until we figure out how to make SinkCache happy