diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 5da4688f8f4e..439556217bde 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -200,6 +200,14 @@ def __init__(self, sliding_window: int, **kwargs): self.sliding_window = sliding_window self.cumulative_length = 0 self._sliding_window_tensor = torch.tensor(self.sliding_window, dtype=torch.long) + self.record_past = False + + def activate_past_recording(self): + """ + Calling this function will activate past state recording, meaning that a call to `update` will wait for a call to `crop` + before restricting the size of the `k/v_states` to `sliding_window`, to be able to retrieve previous full states. + """ + self.record_past = True def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: super().lazy_initialization(key_states, value_states) @@ -228,8 +236,13 @@ def update( full_key_states = torch.cat([self.keys, key_states], dim=-2) full_value_states = torch.cat([self.values, value_states], dim=-2) # Only cache the last `self.sliding_window - 1` tokens (or all of them if lower than that) - self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :] - self.values = full_value_states[:, :, -self.sliding_window + 1 :, :] + if not self.record_past: + self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :] + self.values = full_value_states[:, :, -self.sliding_window + 1 :, :] + # If we record the past, we keep them all for now, and they'll be restricted to the window size in `crop` + else: + self.keys = full_key_states + self.values = full_value_states # Return the full states return full_key_states, full_value_states @@ -256,16 +269,31 @@ def get_max_length(self) -> int: def crop(self, max_length: int) -> None: """ - Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be - negative to remove `max_length` tokens. + Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length` + tokens. """ + # If we are beyond the sliding window, we need to be more careful if self.get_seq_length() >= self.sliding_window: - raise ValueError( - "Cannot `crop` a `DynamicSlidingWindowLayer` after it has seen more tokens than its" - "sliding window (otherwise some states are lost)" - ) - super().crop(max_length) - self.cumulative_length = self.keys.shape[-2] + if not self.record_past: + raise RuntimeError( + "`crop` was called, but the current layer does not track past states, and the sliding window size was already " + "reached. Call `activate_past_recording` before `crop` to be able to rollback the cache." + ) + if max_length > 0: + raise RuntimeError( + "Once the sliding window size has been reached, `DynamicSlidingWindowLayer` can only be cropped by passing a " + "negative int, to specify how many tokens to remove" + ) + tokens_to_remove = abs(max_length) + # We crop, and restrict the size back to the sliding window if still larger + self.keys = self.keys[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] + self.values = self.values[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] + self.cumulative_length = self.cumulative_length - tokens_to_remove + + # If we did not reach the sliding window, we can do the same as for a full attention layer + else: + super().crop(max_length) + self.cumulative_length = self.keys.shape[-2] class DynamicIndexedLayer(DynamicLayer): diff --git a/src/transformers/generation/candidate_generator.py b/src/transformers/generation/candidate_generator.py index 5776e2ccec98..71edaa53a393 100644 --- a/src/transformers/generation/candidate_generator.py +++ b/src/transformers/generation/candidate_generator.py @@ -188,9 +188,6 @@ def __init__( processor for processor in self.logits_processor if not isinstance(processor, MinLengthLogitsProcessor) ] - # We need to roll back the cache in assisted generation, only DynamicCache is supported - self.generation_config.cache_implementation = "dynamic_full" - if ( is_sklearn_available() and self.assistant_generation_config.assistant_confidence_threshold @@ -295,8 +292,8 @@ def _update_past_and_masks( """Update past key values and attention masks for subsequent generation rounds.""" has_past_key_values = self.assistant_kwargs.get("past_key_values", None) is not None if has_past_key_values: - new_cache_size = input_ids.shape[-1] - 1 - remove_from_pkv - self.assistant_kwargs["past_key_values"].crop(new_cache_size - num_added_tokens) + tokens_to_remove = remove_from_pkv + num_added_tokens + self.assistant_kwargs["past_key_values"].crop(-tokens_to_remove) self.assistant_kwargs = _prepare_attention_mask( self.assistant_kwargs, input_ids.shape[-1], self.assistant_model.config.is_encoder_decoder ) @@ -305,10 +302,6 @@ def _update_past_and_masks( ) self.assistant_kwargs = _prepare_token_type_ids(self.assistant_kwargs, input_ids.shape[-1]) - # This unsets `dynamic_full`, needed to initialize a new cache for the assistant. After the first forward - # pass on each generation, we reuse the cache instead. - self.generation_config.cache_implementation = None - return has_past_key_values def _prepare_generation_args(self, input_ids: torch.LongTensor, min_new_tokens: int, max_new_tokens: int) -> dict: @@ -1434,13 +1427,7 @@ def __init__( model_kwargs: dict[str, Any], logits_processor: Optional["LogitsProcessorList"] = None, ): - from ..cache_utils import ( - DynamicLayer, - DynamicSlidingWindowLayer, - LinearAttentionAndFullAttentionLayer, - LinearAttentionAndSlidingWindowAttentionLayer, - MtpCache, - ) + from ..cache_utils import MtpCache from ..modeling_layers import MtpModel self.num_mtp_layers = getattr(main_model.config.get_text_config(), "num_mtp_layers", None) @@ -1456,18 +1443,7 @@ def __init__( self.mtp_model = MtpModel.from_pretrained(main_model, device_map={"": self.device}) # Create the mtp cache and allow it to keep its past before we crop it - mtp_cache = MtpCache(config=main_model.config.get_mtp_config()) - # Use full attention for now on sliding layers, same as main_model cache - mtp_cache.layers = [ - DynamicLayer() if type(layer) is DynamicSlidingWindowLayer else layer for layer in mtp_cache.layers - ] - mtp_cache.layers = [ - LinearAttentionAndFullAttentionLayer(number_of_states=layer.number_of_states) - if type(layer) is LinearAttentionAndSlidingWindowAttentionLayer - else layer - for layer in mtp_cache.layers - ] - self.mtp_cache = mtp_cache + self.mtp_cache = MtpCache(config=main_model.config.get_mtp_config()) self.mtp_cache.activate_past_recording() # Save those to know how to decode mtp tokens diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index ad2847f18eba..2eeb3739797c 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/src/transformers/generation/configuration_utils.py @@ -44,7 +44,7 @@ logger = logging.get_logger(__name__) METADATA_FIELDS = ("_from_model_config", "_commit_hash", "_original_object_hash", "transformers_version") STATIC_CACHE_IMPLEMENTATIONS = ("static", "offloaded_static") -DYNAMIC_CACHE_IMPLEMENTATIONS = ("dynamic", "dynamic_full", "offloaded", "quantized") +DYNAMIC_CACHE_IMPLEMENTATIONS = ("dynamic", "offloaded", "quantized") # All the following are redundant and deprecated, but kept for BC DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS = ( "sliding_window", diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 54e70ee16485..3ac3c3a1dad8 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1955,20 +1955,7 @@ def _prepare_cache_for_generation( return # Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation` - - # Assisted decoding and contrastive search require cache rollback, which is incompatible with sliding layers. - # To handle this, we skip passing the model config to DynamicCache (forcing a full-layer cache). - # The "dynamic_full" option is a shortcut for generate() users to avoid sliding layers on their own. - if generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH): - if generation_config.cache_implementation is not None: - logger.warning_once( - "An assistant model is provided, using a dynamic cache instead of a cache of type=" - f"'{generation_config.cache_implementation}'." - ) - generation_config.cache_implementation = "dynamic_full" - - dynamic_cache_kwargs = {} - dynamic_cache_kwargs["config"] = self.config.get_text_config(decoder=True) + dynamic_cache_kwargs = {"config": self.config.get_text_config(decoder=True)} if generation_config.cache_implementation == "offloaded": dynamic_cache_kwargs["offloading"] = True @@ -2005,30 +1992,9 @@ def _prepare_cache_for_generation( cache_config.setdefault("config", self.config.get_text_config(decoder=True)) backend = cache_config.pop("backend", "quanto") model_kwargs[cache_name] = QuantizedCache(backend=backend, **cache_config) - # i.e. `cache_implementation` in [None, "dynamic", "offloaded", "dynamic_full"] - # TODO: prepare linear cache from a single API, instead of creating in modeling code + # i.e. `cache_implementation` in [None, "dynamic", "offloaded"] else: - cache = DynamicCache(**dynamic_cache_kwargs) - # Replace sliding by full - if generation_config.cache_implementation == "dynamic_full": - from ..cache_utils import ( - DynamicLayer, - DynamicSlidingWindowLayer, - LinearAttentionAndFullAttentionLayer, - LinearAttentionAndSlidingWindowAttentionLayer, - ) - - cache.layers = [ - DynamicLayer() if type(layer) is DynamicSlidingWindowLayer else layer for layer in cache.layers - ] - cache.layers = [ - LinearAttentionAndFullAttentionLayer(number_of_states=layer.number_of_states) - if type(layer) is LinearAttentionAndSlidingWindowAttentionLayer - else layer - for layer in cache.layers - ] - - model_kwargs[cache_name] = cache + model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs) if ( self.config.is_encoder_decoder @@ -2040,6 +2006,10 @@ def _prepare_cache_for_generation( DynamicCache(**dynamic_cache_kwargs), # cross-attention cache ) + # If we just created a cache for an assistant model, mark it for past recording, as we will need to rollback it + if generation_config.is_assistant: + model_kwargs[cache_name].activate_past_recording() + def _supports_logits_to_keep(self: "GenerativePreTrainedModel") -> bool: """ Return True if the current model supports the keyword argument `logits_to_keep` in forward() diff --git a/src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py b/src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py index 714802956af8..d0b7d0140b0d 100644 --- a/src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py +++ b/src/transformers/models/diffusion_gemma/generation_diffusion_gemma.py @@ -925,9 +925,7 @@ def _prepare_cache_for_generation( # Dynamic Caches else: - dynamic_cache_kwargs = {} - if generation_config.cache_implementation != "dynamic_full": - dynamic_cache_kwargs["config"] = self.config.get_text_config(decoder=True) + dynamic_cache_kwargs = {"config": self.config.get_text_config(decoder=True)} if generation_config.cache_implementation == "offloaded": dynamic_cache_kwargs["offloading"] = True past_key_values = DynamicCache(**dynamic_cache_kwargs) diff --git a/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py index b04737837609..a9b38c1ae2e9 100644 --- a/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py @@ -79,11 +79,12 @@ def batch_select_indices(self, indices: torch.Tensor) -> None: self.idx_keys = self.idx_keys[indices, ...] def crop(self, max_length: int) -> None: - super().crop(max_length) + # Important to get the seq_len before the call to `super`, as it will be changed inside otherwise if max_length < 0: max_length = self.get_seq_length() - abs(max_length) if self.idx_keys is not None and self.idx_keys.shape[-2] > max_length: self.idx_keys = self.idx_keys[..., :max_length, :] + super().crop(max_length) class MiniMaxM3VLSparseStaticCacheLayer(StaticLayer): diff --git a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py index 438316b593c9..a7d92104c87b 100644 --- a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py @@ -257,11 +257,12 @@ def batch_select_indices(self, indices: torch.Tensor) -> None: self.idx_keys = self.idx_keys[indices, ...] def crop(self, max_length: int) -> None: - super().crop(max_length) + # Important to get the seq_len before the call to `super`, as it will be changed inside otherwise if max_length < 0: max_length = self.get_seq_length() - abs(max_length) if self.idx_keys is not None and self.idx_keys.shape[-2] > max_length: self.idx_keys = self.idx_keys[..., :max_length, :] + super().crop(max_length) class MiniMaxM3VLSparseStaticCacheLayer(StaticLayer): diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index f239e09d29ff..f1469d715090 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -699,8 +699,6 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type): # the assistant model is correct # c) there are at least two forward passes in the main model, to ensure the input preparation of # the main model is correct - # d) use a cache type compatible with rollbacks (only dynamic cache atm). Otherwise, there may be - # differences vs model-specific default cache generation_kwargs = { "eos_token_id": -1, # see a) "max_new_tokens": 4, # see c) @@ -712,7 +710,6 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type): "output_attentions": self.has_attentions, "return_dict_in_generate": True, "use_cache": True, - "cache_implementation": "dynamic_full", # see d) } logits_processor_kwargs = self._get_logits_processor_kwargs(config=model.config) @@ -797,8 +794,6 @@ def test_prompt_lookup_decoding_matches_greedy_search(self): # prompt lookup is correct # c) there are at least two forward passes in the main model, to ensure the input preparation of # the main model is correct - # d) use a cache type compatible with rollbacks (only dynamic cache atm). Otherwise, there may be - # differences vs model-specific default cache generation_kwargs = { "eos_token_id": -1, # see a) "max_new_tokens": 4, # see c) @@ -810,7 +805,6 @@ def test_prompt_lookup_decoding_matches_greedy_search(self): "output_attentions": self.has_attentions, "return_dict_in_generate": True, "use_cache": True, - "cache_implementation": "dynamic_full", # see d) } logits_processor_kwargs = self._get_logits_processor_kwargs(config=model.config) @@ -872,8 +866,6 @@ def test_assisted_decoding_sample(self): # the assistant model is correct # c) there are at least two forward passes in the main model, to ensure the input preparation of # the main model is correct - # d) use a cache type compatible with rollbacks (only dynamic cache atm). Otherwise, there may be - # differences vs model-specific default cache assistant_model = model assistant_model.generation_config.num_assistant_tokens = 2 # see b) assistant_model.generation_config.num_assistant_tokens_schedule = "constant" # see b) @@ -889,7 +881,6 @@ def test_assisted_decoding_sample(self): "output_attentions": self.has_attentions, "return_dict_in_generate": True, "use_cache": True, - "cache_implementation": "dynamic_full", # see d) } logits_processor_kwargs = self._get_logits_processor_kwargs(config=model.config) output_assisted = model.generate(**generation_kwargs, **inputs_dict, **logits_processor_kwargs)