T5 compile compatibilty #34089

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

Merged

zucchini-nlp merged 34 commits into huggingface:main from zucchini-nlp:t5-compile

Oct 22, 2024

src/transformers/cache_utils.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -1475,11 +1475,7 @@ def from_legacy_cache( @@
         def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
             """Returns the sequence length of the cached states. A layer index can be optionally passed."""
             # check if empty list because in case of static cache it will be a tensors and we can't check `if not torch.Tensor`
-            if self.self_attention_cache.key_cache == []:
-                return 0
-            if len(self.self_attention_cache.key_cache) > 1 and self.self_attention_cache.key_cache[layer_idx] == []:
-                return 0
-            return (self.self_attention_cache.key_cache[layer_idx][0, 0].any(dim=-1)).sum()
+            return self.self_attention_cache.get_seq_length(layer_idx)
         def reset(self):
             if hasattr(self.self_attention_cache, "reset"):
@@ Expand Down @@

src/transformers/generation/utils.py

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -1535,8 +1535,12 @@ def _prepare_generation_config(
  
        def _get_initial_cache_position(self, input_ids, model_kwargs):

            """Calculates `cache_position` for the pre-fill stage based on `input_ids` and optionally past length"""

            # `torch.compile`-friendly `torch.arange` from a shape -- the lines below are equivalent to `torch.arange`

            if "inputs_embeds" in model_kwargs:

            if "inputs_embeds" in model_kwargs and not self.config.is_encoder_decoder:

                cache_position = torch.ones_like(model_kwargs["inputs_embeds"][0, :, 0], dtype=torch.int64).cumsum(0) - 1

            elif "decoder_inputs_embeds" in model_kwargs and self.config.is_encoder_decoder:

                cache_position = (

                    torch.ones_like(model_kwargs["decoder_inputs_embeds"][0, :, 0], dtype=torch.int64).cumsum(0) - 1

                )

            else:

                cache_position = torch.ones_like(input_ids[0, :], dtype=torch.int64).cumsum(0) - 1

    @@ -1633,7 +1637,7 @@ def get_layer_device_map(execution_device_map: Optional[dict] = None):
  
                cache_kwargs = {

                    "config": self.config.get_text_config(),

                    "max_batch_size": batch_size,

                    "batch_size": batch_size,

                    "max_cache_len": max_cache_len,

                    "device": device,

                    "dtype": cache_dtype,

src/transformers/models/longt5/configuration_longt5.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -79,7 +79,12 @@ class LongT5Config(PretrainedConfig): @@
         model_type = "longt5"
         keys_to_ignore_at_inference = ["past_key_values"]
-        attribute_map = {"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}
+        attribute_map = {
+            "hidden_size": "d_model",
+            "num_attention_heads": "num_heads",
+            "num_hidden_layers": "num_layers",
+            "head_dim": "d_kv",
+        }
         def __init__(
             self,
@@ Expand Down @@

src/transformers/models/longt5/modeling_longt5.py