variable_cache.py compatibility for v4.57.2 / python3.12

modeling_decilm.py

@@ -27,7 +27,7 @@ import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers import GenerationConfig
-from transformers.generation.utils import NEED_SETUP_CACHE_CLASSES_MAPPING, GenerationMixin, GenerateOutput
+from transformers.generation.utils import GenerationMixin, GenerateOutput
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 from transformers.utils import (
@@ -809,8 +809,9 @@ class DeciLMPreTrainedModel(PreTrainedModel):
     ) -> tuple[GenerationConfig, dict]:
         # DeciLM-specific code
         generation_config, model_kwargs = super()._prepare_generation_config(generation_config, *args, **kwargs)
-        generation_config.cache_implementation = "variable"
-        NEED_SETUP_CACHE_CLASSES_MAPPING["variable"] = VariableCache
+        generation_config.disable_compile = True
+        #generation_config.cache_implementation = "variable"
+        #NEED_SETUP_CACHE_CLASSES_MAPPING["variable"] = VariableCache
         return generation_config, model_kwargs
 
 
@@ -1133,6 +1134,9 @@ class DeciLMForCausalLM(DeciLMPreTrainedModel, GenerationMixin):
     def get_decoder(self):
         return self.model
 
+    def getVariableCache(self, batch_size=1, max_cache_len=4096, dtype=torch.bfloat16):
+        return VariableCache(config=self.config, batch_size=batch_size, max_batch_size=batch_size, max_cache_len=max_cache_len, dtype=dtype)
+
     @add_start_docstrings_to_model_forward(DECILM_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(

variable_cache.py

@@ -31,6 +31,9 @@ class VariableCache(Cache_4_44_2, Cache):
     The default implementation for the layer caches is StaticCache.
     The cache of each layer is allocated to the same gpu as the layer itself.
     """
+    
+    max_batch_size = None
+    max_cache_len = None
 
     def __init__(
             self,
@@ -50,7 +53,7 @@ class VariableCache(Cache_4_44_2, Cache):
         self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len
         self.dtype = dtype
 
-        self.layer_caches: list[Cache_4_44_2 | None] = [None] * config.num_hidden_layers
+        self.layers: list[Cache_4_44_2 | None] = [None] * config.num_hidden_layers
         self.layer_devices: list[torch.device | None] = [None] * config.num_hidden_layers
 
     def update(
@@ -60,11 +63,11 @@ class VariableCache(Cache_4_44_2, Cache):
             layer_idx: int,
             cache_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if self.layer_caches[layer_idx] is None:
+        if self.layers[layer_idx] is None:
             self.layer_devices[layer_idx] = key_states.device
             self._init_layer_cache(layer_idx)
 
-        layer_cache = self.layer_caches[layer_idx]
+        layer_cache = self.layers[layer_idx]
         assert layer_cache is not None, f"Trying to update the cache of a cache-less layer: {layer_idx=}"
 
         k_out, v_out = layer_cache.update(key_states=key_states,
@@ -93,37 +96,37 @@ class VariableCache(Cache_4_44_2, Cache):
         if attention_config.window_length is not None:
             if not attention_config.is_sink:
                 config.sliding_window = attention_config.window_length
-                self.layer_caches[layer_idx] = SlidingWindowCache(config=config,
+                self.layers[layer_idx] = SlidingWindowCache(config=config,
                                                                   max_batch_size=self.max_batch_size,
                                                                   max_cache_len=self.max_cache_len,
                                                                   device=device,
                                                                   dtype=self.dtype)
                 return
             elif not attention_config.unshifted_sink:
-                self.layer_caches[layer_idx] = SinkCache(window_length=attention_config.window_length,
+                self.layers[layer_idx] = SinkCache(window_length=attention_config.window_length,
                                                          num_sink_tokens=attention_config.num_sink_tokens)
                 return
 
-        self.layer_caches[layer_idx] = StaticCache(config=config,
+        self.layers[layer_idx] = StaticCache(config=config,
                                                    max_batch_size=self.max_batch_size,
                                                    max_cache_len=self.max_cache_len,
                                                    device=device,
                                                    dtype=self.dtype)
 
     def _get_first_real_cache(self) -> Cache:
-        for layer_cache in self.layer_caches:
+        for layer_cache in self.layers:
             if layer_cache is not None:
                 return layer_cache
         raise ValueError(f"No real cache found, all layer caches are None.")
 
     def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
-        if layer_idx == 0 and self.layer_caches[0] is None:
+        if layer_idx == 0 and self.layers[0] is None:
             try:
                 layer_cache = self._get_first_real_cache()
             except ValueError:
                 return 0
         else:
-            layer_cache = self.layer_caches[layer_idx]
+            layer_cache = self.layers[layer_idx]
         return layer_cache.get_seq_length()
 
     def get_max_length(self) -> Optional[int]:
@@ -131,9 +134,12 @@ class VariableCache(Cache_4_44_2, Cache):
         return self.max_cache_len
 
     def reset(self):
-        for layer_idx in range(len(self.layer_caches)):
-            layer_cache = self.layer_caches[layer_idx]
+        for layer_idx in range(len(self.layers)):
+            layer_cache = self.layers[layer_idx]
             if hasattr(layer_cache, "reset"):
                 layer_cache.reset()
             else:
                 self._init_layer_cache(layer_idx)
+
+    def is_compileable(self) -> bool:
+        return False