-
Notifications
You must be signed in to change notification settings - Fork 26.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Cache: Static cache as a standalone object #30476
Changes from 1 commit
fa0b05e
56cf73b
97aa785
a6853a1
326a1c7
9de9556
90771df
cd6aecd
b8caa5e
16e01f2
2d1eab8
bb159d2
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -732,27 +732,6 @@ def _init_weights(self, module): | |
if module.padding_idx is not None: | ||
module.weight.data[module.padding_idx].zero_() | ||
|
||
def _setup_cache(self, cache_cls, max_batch_size, max_cache_len: Optional[int] = None): | ||
if self.config._attn_implementation == "flash_attention_2" and cache_cls == StaticCache: | ||
raise ValueError( | ||
"`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " | ||
"make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" | ||
) | ||
|
||
for layer in self.model.layers: | ||
device = layer.input_layernorm.weight.device | ||
if hasattr(self.config, "_pre_quantization_dtype"): | ||
dtype = self.config._pre_quantization_dtype | ||
else: | ||
dtype = layer.self_attn.o_proj.weight.dtype | ||
layer.self_attn.past_key_value = cache_cls( | ||
self.config, max_batch_size, max_cache_len, device=device, dtype=dtype | ||
) | ||
|
||
def _reset_cache(self): | ||
for layer in self.model.layers: | ||
layer.self_attn.past_key_value = None | ||
|
||
|
||
COHERE_INPUTS_DOCSTRING = r""" | ||
Args: | ||
|
@@ -980,7 +959,7 @@ def _update_causal_mask( | |
attention_mask: torch.Tensor, | ||
input_tensor: torch.Tensor, | ||
cache_position: torch.Tensor, | ||
past_seen_tokens: int, | ||
past_key_values: Cache, | ||
): | ||
# TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static | ||
# KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. | ||
|
@@ -992,9 +971,12 @@ def _update_causal_mask( | |
return attention_mask | ||
return None | ||
|
||
if self.config._attn_implementation == "sdpa": | ||
# For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, | ||
# in order to dispatch on Flash Attention 2. | ||
# For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in | ||
# order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail | ||
# to infer the attention mask. | ||
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 | ||
using_static_cache = isinstance(past_key_values, StaticCache) | ||
if self.config._attn_implementation == "sdpa" and not using_static_cache: | ||
Comment on lines
+976
to
+981
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. this is new, and since we pass cahce position, let's use cache_position[0] There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agreed in theory, can't do in practice: breaks torch.fx tests 💔 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. yeah thought so |
||
if AttentionMaskConverter._ignore_causal_mask_sdpa( | ||
attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens | ||
): | ||
|
@@ -1003,9 +985,9 @@ def _update_causal_mask( | |
dtype, device = input_tensor.dtype, input_tensor.device | ||
min_dtype = torch.finfo(dtype).min | ||
sequence_length = input_tensor.shape[1] | ||
if hasattr(getattr(self.layers[0], "self_attn", {}), "past_key_value"): # static cache | ||
target_length = self.config.max_position_embeddings | ||
else: # dynamic cache | ||
if using_static_cache: | ||
target_length = past_key_values.get_max_length() | ||
Comment on lines
+990
to
+991
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. can't we always use get_max_length()? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It should be There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. but alright |
||
else: | ||
target_length = ( | ||
attention_mask.shape[-1] | ||
if isinstance(attention_mask, torch.Tensor) | ||
|
@@ -1027,6 +1009,10 @@ def _update_causal_mask( | |
# backwards compatibility: we allow passing a 4D attention mask shorter than the input length with | ||
# cache. In that case, the 4D attention mask attends to the newest tokens only. | ||
if attention_mask.shape[-2] < cache_position[0] + sequence_length: | ||
logger.warning_once( | ||
"Passing a 4d mask shorter than the input length is deprecated and will be removed in " | ||
"transformers v4.42.0" | ||
) | ||
offset = cache_position[0] | ||
else: | ||
offset = 0 | ||
|
@@ -1184,13 +1170,6 @@ def prepare_inputs_for_generation( | |
use_cache=True, | ||
**kwargs, | ||
): | ||
# With static cache, the `past_key_values` is None | ||
# TODO joao: standardize interface for the different Cache classes and remove of this if | ||
has_static_cache = False | ||
if past_key_values is None: | ||
past_key_values = getattr(getattr(self.model.layers[0], "self_attn", {}), "past_key_value", None) | ||
has_static_cache = past_key_values is not None | ||
|
||
past_length = 0 | ||
if past_key_values is not None: | ||
if isinstance(past_key_values, Cache): | ||
|
@@ -1208,8 +1187,7 @@ def prepare_inputs_for_generation( | |
|
||
# Keep only the unprocessed tokens: | ||
# 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where | ||
# some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as | ||
# input) | ||
# some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as input) | ||
if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: | ||
input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] | ||
# 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard | ||
|
@@ -1249,9 +1227,6 @@ def prepare_inputs_for_generation( | |
elif use_cache: | ||
cache_position = cache_position[-input_length:] | ||
|
||
if has_static_cache: | ||
past_key_values = None | ||
|
||
model_inputs.update( | ||
{ | ||
"position_ids": position_ids, | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
might be faster?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
setting to a new tensor produces a graph break 💔 (I'm assuming you meant
self.key_cache[layer_idx] = torch.zeros(...)
)There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No no, I think just filling them with zeros should work
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That would result in
TypeError: 'float' object is not subscriptable
when indexing the cache :DBut filling with zeros with
tensor.zero_()
works 👍There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ok 👍🏻 let's go with that then!