diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index a79672d76e05..cb8df79f8079 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -7,6 +7,7 @@ from .configuration_utils import PreTrainedConfig from .utils import ( is_hqq_available, + is_optimum_quanto_available, is_quanto_greater, is_torch_greater_or_equal, is_torchdynamo_compiling, @@ -584,7 +585,12 @@ def __init__( ) # We need to import quanto here to avoid circular imports due to optimum/quanto/models/transformers_models.py - if is_quanto_greater("0.2.5", accept_dev=True): + if not is_optimum_quanto_available(): + raise ImportError( + "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto " + "backend. Please install it via with `pip install optimum-quanto`" + ) + elif is_quanto_greater("0.2.5", accept_dev=True): from optimum.quanto import MaxOptimizer, qint2, qint4 else: raise ImportError( @@ -634,7 +640,10 @@ def __init__( ) if not is_hqq_available(): - raise ImportError("You need to install `hqq` to use `HQQQuantizedLayer`") + raise ImportError( + "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. " + "Please install it via with `pip install hqq`" + ) if self.nbits not in [1, 2, 3, 4, 8]: raise ValueError( diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index c929f88f184a..8d6fdc01fd53 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -48,8 +48,6 @@ ModelOutput, TransformersKwargs, is_accelerate_available, - is_hqq_available, - is_optimum_quanto_available, is_torchdynamo_exporting, logging, ) @@ -1858,20 +1856,19 @@ def _get_initial_cache_position(self, seq_length, device, model_kwargs): model_kwargs["cache_position"] = cache_position return model_kwargs - def _get_cache(self, cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs) -> Cache: + def _prepare_static_cache( + self, cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs + ) -> Cache: """ Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a new `generate` call requires a larger cache or uses a different batch size. Returns the resulting cache object. """ - requires_cross_attention_cache = ( - self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None - ) offload_cache = "offloaded" in cache_implementation if hasattr(self, "_cache"): - cache_to_check = self._cache.self_attention_cache if requires_cross_attention_cache else self._cache + cache_to_check = self._cache.self_attention_cache if self.config.is_encoder_decoder else self._cache need_new_cache = ( not hasattr(self, "_cache") @@ -1880,7 +1877,7 @@ def _get_cache(self, cache_implementation: str, batch_size: int, max_cache_len: or cache_to_check.max_cache_len < max_cache_len ) - if requires_cross_attention_cache and hasattr(self, "_cache"): + if self.config.is_encoder_decoder and hasattr(self, "_cache"): need_new_cache = ( need_new_cache or self._cache.cross_attention_cache.max_cache_len != model_kwargs["encoder_outputs"][0].shape[1] @@ -1893,7 +1890,7 @@ def _get_cache(self, cache_implementation: str, batch_size: int, max_cache_len: "offloading": offload_cache, } self._cache = StaticCache(**self_attention_cache_kwargs) - if requires_cross_attention_cache: + if self.config.is_encoder_decoder: cross_attention_cache_kwargs = { "config": self.config.get_text_config(decoder=True), "max_cache_len": model_kwargs["encoder_outputs"][0].shape[1], @@ -1936,12 +1933,9 @@ def _prepare_cache_for_generation( instantiated, writes it to `model_kwargs`, under the name expected by the model. """ - is_hybrid_cache = any(class_name in self.__class__.__name__.lower() for class_name in ["mamba", "falconh1"]) - cache_name = "past_key_values" if not is_hybrid_cache else "cache_params" - - requires_cross_attention_cache = ( - self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None - ) + # TODO @raushan, unify cache arg naming for all models + is_linear_attn_cache = "mamba" in self.__class__.__name__.lower() + cache_name = "past_key_values" if not is_linear_attn_cache else "cache_params" # Quick escape route 1: if the user specifies a cache, we only need to check for conflicting `generate` arguments user_defined_cache = model_kwargs.get(cache_name) @@ -1973,76 +1967,55 @@ def _prepare_cache_for_generation( # Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation` - # TODO(joao): support static caches in assisted generation. assisted generation needs to roll back caches, - # which is only supported in dynamic caches atm - if ( - generation_mode == GenerationMode.ASSISTED_GENERATION - and 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 = None - # 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) - or generation_config.cache_implementation == "dynamic_full" - ): - dynamic_cache_kwargs = {} - else: - dynamic_cache_kwargs = {"config": self.config.get_text_config(decoder=True)} - if generation_config.cache_implementation is not None: - if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS: - if generation_config.cache_implementation in DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS: - logger.warning_once( - f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated. " - f"Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, and the layer structure will be " - "inferred automatically." - ) - model_kwargs[cache_name] = self._get_cache( - cache_implementation=generation_config.cache_implementation, - batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size, - max_cache_len=max_cache_length, - model_kwargs=model_kwargs, + 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}'." ) - elif generation_config.cache_implementation == "quantized": - if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache(): - raise ValueError( - "This model does not support the quantized cache. If you want your model to support quantized " - "cache, please open an issue and tag @zucchini-nlp." - ) + generation_config.cache_implementation = "dynamic_full" + + dynamic_cache_kwargs = {} + if generation_config.cache_implementation != "dynamic_full": + dynamic_cache_kwargs["config"] = self.config.get_text_config(decoder=True) + + if generation_config.cache_implementation == "offloaded": + dynamic_cache_kwargs["offloading"] = True + + if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS: + if generation_config.cache_implementation in DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS: + logger.warning_once( + f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated " + f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, " + "and the layer structure will be inferred automatically." + ) + model_kwargs["past_key_values"] = self._prepare_static_cache( + cache_implementation=generation_config.cache_implementation, + batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size, + max_cache_len=max_cache_length, + model_kwargs=model_kwargs, + ) + elif generation_config.cache_implementation == "quantized": + if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache(): + raise ValueError( + "This model does not support the quantized cache. If you want your model to support quantized " + "cache, please open an issue and tag @zucchini-nlp." + ) + + cache_config = generation_config.cache_config if generation_config.cache_config is not None else {} + cache_config.setdefault("config", self.config.get_text_config(decoder=True)) + backend = cache_config.pop("backend", "quanto") + model_kwargs["past_key_values"] = 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 + else: + model_kwargs["past_key_values"] = DynamicCache(**dynamic_cache_kwargs) - cache_config = generation_config.cache_config if generation_config.cache_config is not None else {} - # Add the config if it was not provided, as it's a required argument - if "config" not in cache_config: - cache_config["config"] = self.config.get_text_config() - # Pop the backend from the config (defaults to quanto if not defined) - backend = cache_config.pop("backend", "quanto") - - if backend == "quanto" and not is_optimum_quanto_available(): - raise ImportError( - "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto " - "backend. Please install it via with `pip install optimum-quanto`" - ) - elif backend == "HQQ" and not is_hqq_available(): - raise ImportError( - "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. " - "Please install it via with `pip install hqq`" - ) - model_kwargs[cache_name] = QuantizedCache(backend=backend, **cache_config) - elif generation_config.cache_implementation == "offloaded": - model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs, offloading=True) - elif "dynamic" in generation_config.cache_implementation: - model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs) - - # TODO (joao): this logic is incomplete, e.g. `offloaded` should apply to both caches. Refactor this function - # to correctly pass parameterization to both caches. if ( - requires_cross_attention_cache + self.config.is_encoder_decoder and "past_key_values" in model_kwargs and not isinstance(model_kwargs["past_key_values"], EncoderDecoderCache) ): diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index b71cad55f602..aa59dee08352 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -576,7 +576,7 @@ def forward( raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if use_cache and past_key_values is None: - past_key_values = DynamicCache() + past_key_values = DynamicCache(config=self.config) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index 3e5c5f66e621..32bfe9159a7d 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -397,7 +397,7 @@ def forward( raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if use_cache and past_key_values is None: - past_key_values = DynamicCache() + past_key_values = DynamicCache(config=self.config) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py index 8cc3e28bbdf0..2f9da1f0091e 100644 --- a/src/transformers/models/blt/modeling_blt.py +++ b/src/transformers/models/blt/modeling_blt.py @@ -930,7 +930,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = DynamicCache() + past_key_values = DynamicCache(config=self.config) if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 @@ -1229,8 +1229,15 @@ def forward( if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if use_cache and past_key_values is None: - past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + if use_cache: + if past_key_values is None: + past_key_values = EncoderDecoderCache( + DynamicCache(config=self.config), DynamicCache(config=self.config) + ) + elif not isinstance(past_key_values, EncoderDecoderCache): + # BLT uses an encoder-decoder cache even though it is not en encoder-decoder model. Create a cross-cache + # if not yet created by the user + past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config)) # Extract input embeddings as early as possible if inputs_embeds is not None: @@ -1377,7 +1384,7 @@ class BltForCausalLM(BltPreTrainedModel, GenerationMixin): _tied_weights_keys = {"model.local_encoder.embed_tokens.weight": "lm_head.weight"} def __init__(self, config: BltConfig): - super().__init__(config.get_text_config()) + super().__init__(config) self.text_config = config.get_text_config() self.vocab_size = config.vocab_size self.model = BltModel(config) @@ -1433,8 +1440,8 @@ def forward( ```python >>> from transformers import AutoTokenizer, BltForCausalLM - >>> model = BltForCausalLM.from_pretrained("Llama-3.2-11B-Vision") - >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision") + >>> model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf") >>> prompt = "If I had to write a haiku, it would be:" >>> inputs = tokenizer(prompt, return_tensors="pt") diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py index 22fef55c96d2..e5b721927fb1 100644 --- a/src/transformers/models/blt/modular_blt.py +++ b/src/transformers/models/blt/modular_blt.py @@ -22,17 +22,17 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...generation import GenerationMixin from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ...utils.generic import OutputRecorder, check_model_inputs, maybe_autocast from ..cohere2.modeling_cohere2 import rotate_half # noqa: F401 from ..llama.modeling_llama import LlamaRotaryEmbedding from ..mllama.modeling_mllama import ( - MllamaForCausalLM, MllamaPreTrainedModel, MllamaSelfAttentionDecoderLayer, MllamaTextCrossAttention, @@ -818,7 +818,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = DynamicCache() + past_key_values = DynamicCache(config=self.config) if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 @@ -954,8 +954,15 @@ def forward( if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if use_cache and past_key_values is None: - past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + if use_cache: + if past_key_values is None: + past_key_values = EncoderDecoderCache( + DynamicCache(config=self.config), DynamicCache(config=self.config) + ) + elif not isinstance(past_key_values, EncoderDecoderCache): + # BLT uses an encoder-decoder cache even though it is not en encoder-decoder model. Create a cross-cache + # if not yet created by the user + past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config)) # Extract input embeddings as early as possible if inputs_embeds is not None: @@ -1090,7 +1097,12 @@ def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) -> return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1 -class BltForCausalLM(MllamaForCausalLM): +@auto_docstring( + custom_intro=""" + The Blt Text Model with a language modeling head on top. + """ +) +class BltForCausalLM(BltPreTrainedModel, GenerationMixin): config: BltConfig _can_compile_fullgraph = False base_model_prefix = "model" @@ -1098,11 +1110,15 @@ class BltForCausalLM(MllamaForCausalLM): def __init__(self, config: BltConfig): super().__init__(config) + self.text_config = config.get_text_config() self.vocab_size = config.vocab_size self.model = BltModel(config) self.lm_head = nn.Linear(config.decoder_config.hidden_size, config.vocab_size, bias=False) + self.post_init() + @can_return_tuple + @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, @@ -1119,6 +1135,50 @@ def forward( logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], ) -> tuple | CausalLMOutputWithPast: + r""" + cross_attention_states (`torch.FloatTensor`, *optional*): + Output of the vision model, used for cross-attention. This tensor contains the processed image features that + the language model will attend to. + cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): + Cross-attention mask to control the interaction between text tokens and image tiles. + This 4D tensor defines which image tiles each text token should attend to. + + For each text token (in seq_length): + - 1 indicates the token **should attend** to the corresponding image tile + - 0 indicates the token **should not attend** to the corresponding image tile + full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*): + A tuple containing two tensors that mask out rows in the cross-attention mechanism: + - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1. + A value of 0 indicates that the corresponding text token's entire row in the cross-attention + matrix should be masked out (all image tokens ignored). + - The second tensor has the same shape and is used internally to apply the masking during + the forward pass of cross-attention layers. + This mask is derived from the cross_attention_mask and is used to handle cases where a text token + should not attend to any image token. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, BltForCausalLM + + >>> model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf") + + >>> prompt = "If I had to write a haiku, it would be:" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6) + >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(result) + If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful. + I love the idea of snowflakes gently falling, each one + ``` + """ # Call parent forward but exclude cross_attention_states from model call outputs = self.model( input_ids=input_ids, diff --git a/src/transformers/models/dia/generation_dia.py b/src/transformers/models/dia/generation_dia.py index d87eb7879b12..5a94f081af05 100644 --- a/src/transformers/models/dia/generation_dia.py +++ b/src/transformers/models/dia/generation_dia.py @@ -388,26 +388,16 @@ def _main_generate_loop( # Prepare inner 2D logic in generation loop input_ids = input_ids.reshape(-1, input_ids.shape[-1]) - model_kwargs = self._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs) - # prepare model inputs - model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) - - # 10. Prefill - model_inputs.update({"output_attentions": generation_config.output_attentions}) - model_inputs.update({"output_hidden_states": generation_config.output_hidden_states}) - outputs = self(**model_inputs, return_dict=True) - - # 11. expand input_ids with `num_return_sequences` additional sequences per batch + # 10. expand input_ids with `num_return_sequences` additional sequences per batch if generation_config.num_return_sequences > 1: raise ValueError("`num_return_sequences>1` is incompatible with Dia.") - # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`) + # 11. run sample (it degenerates to greedy search when `generation_config.do_sample=False`) return self._sample( input_ids, logits_processor=prepared_logits_processor, stopping_criteria=prepared_stopping_criteria, generation_config=generation_config, - prefill_outputs=outputs, **generation_mode_kwargs, **model_kwargs, ) diff --git a/src/transformers/models/janus/modeling_janus.py b/src/transformers/models/janus/modeling_janus.py index 0c61a42ee790..08098e8e88ab 100644 --- a/src/transformers/models/janus/modeling_janus.py +++ b/src/transformers/models/janus/modeling_janus.py @@ -1314,7 +1314,7 @@ def generate( if model_kwargs.get("past_key_values", None) is None: # Prepare cache if not provided. - model_kwargs["past_key_values"] = self._get_cache( + model_kwargs["past_key_values"] = self._prepare_static_cache( cache_implementation=generation_config.cache_implementation or "static", # batch_size should account for both conditional/unconditional input; hence multiplied by 2. batch_size=batch_size * 2, diff --git a/src/transformers/models/janus/modular_janus.py b/src/transformers/models/janus/modular_janus.py index 882d6de88caa..c4b5f69da0c4 100644 --- a/src/transformers/models/janus/modular_janus.py +++ b/src/transformers/models/janus/modular_janus.py @@ -1255,7 +1255,7 @@ def generate( if model_kwargs.get("past_key_values", None) is None: # Prepare cache if not provided. - model_kwargs["past_key_values"] = self._get_cache( + model_kwargs["past_key_values"] = self._prepare_static_cache( cache_implementation=generation_config.cache_implementation or "static", # batch_size should account for both conditional/unconditional input; hence multiplied by 2. batch_size=batch_size * 2, diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py index 785d60879cf1..094ec2c0948f 100644 --- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py +++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py @@ -1220,7 +1220,7 @@ def _prepare_model_inputs( # Add cache-related methods from GenerationMixin to codec model cache_methods = [ "_prepare_cache_for_generation", - "_get_cache", + "_prepare_static_cache", ] for method in cache_methods: setattr(self.codec_model, method, types.MethodType(getattr(self, method).__func__, self.codec_model)) diff --git a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py index 39bd26f3987d..043663bb026c 100644 --- a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py +++ b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py @@ -351,7 +351,7 @@ def _prepare_model_inputs( # Add cache-related methods from GenerationMixin to codec model cache_methods = [ "_prepare_cache_for_generation", - "_get_cache", + "_prepare_static_cache", ] for method in cache_methods: setattr(self.codec_model, method, types.MethodType(getattr(self, method).__func__, self.codec_model)) diff --git a/tests/models/afmoe/test_modeling_afmoe.py b/tests/models/afmoe/test_modeling_afmoe.py index ed954a4ba687..52af6880b552 100644 --- a/tests/models/afmoe/test_modeling_afmoe.py +++ b/tests/models/afmoe/test_modeling_afmoe.py @@ -32,7 +32,7 @@ def __init__( self, parent, batch_size=4, - seq_length=128, + seq_length=12, is_training=True, use_input_mask=True, use_token_type_ids=False,