Merge branch 'main' into fix-t5-gemma-config

src/transformers/generation/utils.py

@@ -2001,11 +2001,17 @@ class GenerationMixin(ContinuousMixin):
         # Use DynamicCache() instance by default. This will avoid back and forth from legacy format that
         # keeps copying the cache thus using much more memory
         else:
-            model_kwargs[cache_name] = (
-                DynamicCache(**dynamic_cache_kwargs)
-                if not requires_cross_attention_cache
-                else EncoderDecoderCache(DynamicCache(**dynamic_cache_kwargs), DynamicCache(**dynamic_cache_kwargs))
-            )
+            if not requires_cross_attention_cache:
+                model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs)
+            else:
+                # For encoder-decoder models, we need to use separate configs for encoder and decoder
+                decoder_cache_kwargs = {}
+                encoder_cache_kwargs = {}
+                decoder_cache_kwargs["config"] = self.config.get_text_config(decoder=True)
+                encoder_cache_kwargs["config"] = self.config.get_text_config(encoder=True)
+                model_kwargs[cache_name] = EncoderDecoderCache(
+                    DynamicCache(**decoder_cache_kwargs), DynamicCache(**encoder_cache_kwargs)
+                )
 
     def _supports_logits_to_keep(self) -> bool:
         """

src/transformers/models/t5gemma/configuration_t5gemma.py

@@ -323,9 +323,26 @@ class T5GemmaConfig(PretrainedConfig):
             setattr(self.decoder, key, value)
         super().__setattr__(key, value)
 
-    def get_text_config(self, *args, **kwargs):
-        # Always return self, regardless of the decoder option.
-        return self
+    def get_text_config(self, decoder=None, encoder=None):
+        """
+        Returns the text config related to the text input (encoder) or text output (decoder) of the model.
+
+        Args:
+            decoder (`Optional[bool]`, *optional*):
+                If set to `True`, returns the decoder config.
+            encoder (`Optional[bool]`, *optional*):
+                If set to `True`, returns the encoder config.
+        """
+        if decoder is True and encoder is False:
+            return self.decoder
+        elif encoder is True and decoder is False:
+            return self.encoder
+        elif decoder is None and encoder is None:
+            # Default case - return decoder for generation compatibility
+            return self.decoder
+        else:
+            # For any other case (both True, both False, etc.), return self
+            return self
 
 
 __all__ = ["T5GemmaConfig", "T5GemmaModuleConfig"]

src/transformers/models/t5gemma/modular_t5gemma.py

@@ -207,9 +207,26 @@ class T5GemmaConfig(PretrainedConfig):
             setattr(self.decoder, key, value)
         super().__setattr__(key, value)
 
-    def get_text_config(self, *args, **kwargs):
-        # Always return self, regardless of the decoder option.
-        return self
+    def get_text_config(self, decoder=None, encoder=None):
+        """
+        Returns the text config related to the text input (encoder) or text output (decoder) of the model.
+
+        Args:
+            decoder (`Optional[bool]`, *optional*):
+                If set to `True`, returns the decoder config.
+            encoder (`Optional[bool]`, *optional*):
+                If set to `True`, returns the encoder config.
+        """
+        if decoder is True and encoder is False:
+            return self.decoder
+        elif encoder is True and decoder is False:
+            return self.encoder
+        elif decoder is None and encoder is None:
+            # Default case - return decoder for generation compatibility
+            return self.decoder
+        else:
+            # For any other case (both True, both False, etc.), return self
+            return self
 
 
 class T5GemmaRMSNorm(Gemma2RMSNorm):