diff --git a/src/transformers/models/mt5/modeling_mt5.py b/src/transformers/models/mt5/modeling_mt5.py index e936b18ebed8..f33c9f0cc3b3 100644 --- a/src/transformers/models/mt5/modeling_mt5.py +++ b/src/transformers/models/mt5/modeling_mt5.py @@ -16,6 +16,7 @@ import copy import math +from collections.abc import Callable from typing import Optional, Union import torch @@ -37,13 +38,17 @@ Seq2SeqSequenceClassifierOutput, TokenClassifierOutput, ) -from ...modeling_utils import PreTrainedModel +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...utils import ( DUMMY_INPUTS, DUMMY_MASK, + TransformersKwargs, auto_docstring, + can_return_tuple, logging, ) +from ...utils.generic import OutputRecorder, check_model_inputs from .configuration_mt5 import MT5Config @@ -148,6 +153,36 @@ def forward(self, hidden_states): return hidden_states +# Copied from transformers.models.t5.modeling_t5.eager_attention_forward +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: Optional[float] = None, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + if scaling is None: + scaling = query.size(-1) ** -0.5 + + # Take the dot product between "query" and "key" to get the raw attention scores. + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + + if attention_mask is not None: + attention_mask = attention_mask[:, :, :, : key.shape[-2]] + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + # Copied from transformers.models.t5.modeling_t5.T5Attention with T5->MT5 class MT5Attention(nn.Module): def __init__( @@ -155,17 +190,21 @@ def __init__( config: MT5Config, has_relative_attention_bias=False, layer_idx: Optional[int] = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal + self.scaling = 1.0 self.has_relative_attention_bias = has_relative_attention_bias self.relative_attention_num_buckets = config.relative_attention_num_buckets self.relative_attention_max_distance = config.relative_attention_max_distance self.d_model = config.d_model - self.key_value_proj_dim = config.d_kv + self.head_dim = config.d_kv self.n_heads = config.num_heads self.dropout = config.dropout_rate - self.inner_dim = self.n_heads * self.key_value_proj_dim + self.inner_dim = self.n_heads * self.head_dim self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -260,34 +299,34 @@ def forward( position_bias=None, past_key_values=None, query_length=None, - use_cache=False, - output_attentions=False, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): - """ - Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). - """ - # Input is (batch_size, seq_length, dim) - # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder) - batch_size, seq_length = hidden_states.shape[:2] - - # if key_value_states are provided this layer is used as a cross-attention layer for the decoder + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder is_cross_attention = key_value_states is not None - query_states = self.q(hidden_states) - query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + + q_input_shape = (bsz, tgt_len, -1, self.head_dim) + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # get query proj + query_states = self.q(hidden_states).view(*q_input_shape).transpose(1, 2) - # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache` is_updated = False - if isinstance(past_key_values, EncoderDecoderCache): - is_updated = past_key_values.is_updated.get(self.layer_idx) - if is_cross_attention: - # after the first generated id, we can subsequently re-use all key/value_states from cache - curr_past_key_values = past_key_values.cross_attention_cache + if past_key_values is not None: + if isinstance(past_key_values, EncoderDecoderCache): + is_updated = past_key_values.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + curr_past_key_values = past_key_values.cross_attention_cache + else: + curr_past_key_values = past_key_values.self_attention_cache else: - curr_past_key_values = past_key_values.self_attention_cache - else: - curr_past_key_values = past_key_values + curr_past_key_values = past_key_values current_states = key_value_states if is_cross_attention else hidden_states if is_cross_attention and past_key_values is not None and is_updated: @@ -297,8 +336,8 @@ def forward( else: key_states = self.k(current_states) value_states = self.v(current_states) - key_states = key_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) - value_states = value_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) + key_states = key_states.view(*kv_input_shape).transpose(1, 2) + value_states = value_states.view(*kv_input_shape).transpose(1, 2) if past_key_values is not None: # save all key/value_states to cache to be re-used for fast auto-regressive generation @@ -310,47 +349,56 @@ def forward( if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): past_key_values.is_updated[self.layer_idx] = True - # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 - scores = torch.matmul(query_states, key_states.transpose(3, 2)) - if position_bias is None: - key_length = key_states.shape[-2] + q_length = query_states.shape[-2] + kv_length = key_states.shape[-2] # cache position is 0-indexed so we add 1 to get the real length of queries (aka with past) real_seq_length = query_length if query_length is not None else cache_position[-1] + 1 if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, self.n_heads, seq_length, key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, q_length, kv_length), device=query_states.device, dtype=query_states.dtype ) if self.gradient_checkpointing and self.training: position_bias.requires_grad = True else: position_bias = self.compute_bias( - real_seq_length, key_length, device=scores.device, cache_position=cache_position + real_seq_length, kv_length, device=query_states.device, cache_position=cache_position ) - position_bias = position_bias[:, :, -seq_length:, :] + position_bias = position_bias[:, :, -q_length:, :] if mask is not None: - causal_mask = mask[:, :, :, : key_states.shape[-2]] - position_bias = position_bias + causal_mask - - position_bias_masked = position_bias - scores += position_bias_masked + if mask.dtype == torch.bool: + min_dtype = torch.finfo(query_states.dtype).min + # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type) + mask = torch.where( + mask, torch.tensor(0.0, device=mask.device, dtype=query_states.dtype), min_dtype + ) - # (batch_size, n_heads, seq_length, key_length) - attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) - attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + mask = mask[:, :, :, : key_states.shape[-2]] + position_bias = position_bias + mask - attn_output = torch.matmul(attn_weights, value_states) + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + f"Relative bias attention is only supported by `eager` and `sdpa` but found {self.config._attn_implementation}." + ) + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + position_bias, # position bias acts as mask + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) - attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.view(batch_size, -1, self.inner_dim) + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() attn_output = self.o(attn_output) - - outputs = (attn_output, position_bias) - - if output_attentions: - outputs = outputs + (attn_weights,) - return outputs + return attn_output, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->MT5 @@ -358,7 +406,10 @@ class MT5LayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None): super().__init__() self.SelfAttention = MT5Attention( - config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx + config, + has_relative_attention_bias=has_relative_attention_bias, + layer_idx=layer_idx, + is_causal=config.is_decoder, ) self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -369,30 +420,32 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attn_output, position_bias, attn_weights = self.SelfAttention( normed_hidden_states, mask=attention_mask, position_bias=position_bias, past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, cache_position=cache_position, + **kwargs, ) - hidden_states = hidden_states + self.dropout(attention_output[0]) - outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them - return outputs + hidden_states = hidden_states + self.dropout(attn_output) + return hidden_states, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->MT5 class MT5LayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: Optional[int] = None): super().__init__() - self.EncDecAttention = MT5Attention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.EncDecAttention = MT5Attention( + config, + has_relative_attention_bias=False, + layer_idx=layer_idx, + is_causal=False, + ) self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -403,26 +456,23 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, query_length=None, - output_attentions=False, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attn_output, position_bias, attn_weights = self.EncDecAttention( normed_hidden_states, mask=attention_mask, key_value_states=key_value_states, position_bias=position_bias, past_key_values=past_key_values, - use_cache=use_cache, query_length=query_length, - output_attentions=output_attentions, cache_position=cache_position, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + hidden_states = hidden_states + self.dropout(attn_output) + return hidden_states, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5Block with T5->MT5 @@ -448,22 +498,18 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - return_dict=True, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): - self_attention_outputs = self.layer[0]( + self_attention_output, self_attention_position_bias, _ = self.layer[0]( hidden_states, attention_mask=attention_mask, position_bias=position_bias, past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, cache_position=cache_position, + **kwargs, ) - hidden_states = self_attention_outputs[0] - attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights + hidden_states = self_attention_output # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -475,18 +521,18 @@ def forward( hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) do_cross_attention = self.is_decoder and encoder_hidden_states is not None + cross_attention_position_bias = None if do_cross_attention: - cross_attention_outputs = self.layer[1]( + cross_attention_output, cross_attention_position_bias, _ = self.layer[1]( hidden_states, key_value_states=encoder_hidden_states, attention_mask=encoder_attention_mask, position_bias=encoder_decoder_position_bias, past_key_values=past_key_values, query_length=cache_position[-1] + 1, - use_cache=use_cache, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = cross_attention_outputs[0] + hidden_states = cross_attention_output # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -497,9 +543,6 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - # Keep cross-attention outputs and relative position weights - attention_outputs = attention_outputs + cross_attention_outputs[1:] - # Apply Feed Forward layer hidden_states = self.layer[-1](hidden_states) @@ -512,11 +555,7 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - outputs = (hidden_states,) - - return ( - outputs + attention_outputs - ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + return hidden_states, self_attention_position_bias, cross_attention_position_bias # Copied from transformers.models.t5.modeling_t5.T5ClassificationHead with T5->MT5 @@ -544,11 +583,21 @@ class MT5PreTrainedModel(PreTrainedModel): config: MT5Config base_model_prefix = "transformer" supports_gradient_checkpointing = True + _supports_attention_backend = True + _supports_sdpa = True + _supports_flash_attn = False + _supports_flex_attn = False _can_compile_fullgraph = True _no_split_modules = ["MT5Block"] _keep_in_fp32_modules = ["wo"] + _can_record_outputs = { + "hidden_states": OutputRecorder(MT5Block, index=0), + "attentions": OutputRecorder(MT5LayerSelfAttention, index=-1), + "cross_attentions": OutputRecorder(MT5LayerCrossAttention, index=-1), + } + @property def dummy_inputs(self): input_ids = torch.tensor(DUMMY_INPUTS) @@ -658,6 +707,7 @@ def __init__(self, config): def set_input_embeddings(self, new_embeddings): self.embed_tokens = new_embeddings + @check_model_inputs def forward( self, input_ids=None, @@ -667,50 +717,37 @@ def forward( inputs_embeds=None, past_key_values=None, use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, cache_position=None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - use_cache = use_cache if use_cache is not None else self.config.use_cache - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if input_ids is not None and inputs_embeds is not None: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError( - f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time" - ) - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache else: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") + use_cache = False - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: - if self.embed_tokens is None: - raise ValueError("You have to initialize the model with valid token embeddings") inputs_embeds = self.embed_tokens(input_ids) - batch_size, seq_length = input_shape + if use_cache and past_key_values is None: + past_key_values = ( + EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + if self.config.is_encoder_decoder + else DynamicCache(config=self.config) + ) + elif not self.is_decoder: + # do not pass cache object down the line for encoder stack + # it messes indexing later in decoder-stack because cache object is modified in-place + past_key_values = None - if use_cache is True: - if not self.is_decoder: - raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder") + device = inputs_embeds.device + seq_length = inputs_embeds.shape[1] + + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + if cache_position is None: + cache_position = torch.arange(past_key_values_length, past_key_values_length + seq_length, device=device) if self.is_decoder: if use_cache and past_key_values is None: @@ -757,19 +794,13 @@ def forward( encoder_hidden_states=encoder_hidden_states, ) - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - all_cross_attentions = () if (output_attentions and self.is_decoder) else None position_bias = None encoder_decoder_position_bias = None hidden_states = self.dropout(inputs_embeds) for layer_module in self.block: - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_outputs = layer_module( + hidden_states, self_attention_position_bias, cross_attention_position_bias = layer_module( hidden_states, attention_mask, position_bias, @@ -777,51 +808,21 @@ def forward( encoder_extended_attention_mask, encoder_decoder_position_bias, # as a positional argument for gradient checkpointing past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - return_dict=return_dict, cache_position=cache_position, + **kwargs, ) - hidden_states = layer_outputs[0] - - # We share the position biases between the layers - the first layer store them - # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), - # (cross-attention position bias), (cross-attention weights) - position_bias = layer_outputs[1] + # We share the position biases between the layers + position_bias = self_attention_position_bias if self.is_decoder and encoder_hidden_states is not None: - encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2] - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[2],) - if self.is_decoder: - all_cross_attentions = all_cross_attentions + (layer_outputs[4],) + encoder_decoder_position_bias = cross_attention_position_bias hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.dropout(hidden_states) - # Add last layer - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - past_key_values, - all_hidden_states, - all_attentions, - all_cross_attentions, - ] - if v is not None - ) return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=past_key_values, - hidden_states=all_hidden_states, - attentions=all_attentions, - cross_attentions=all_cross_attentions, ) @@ -882,6 +883,7 @@ def set_input_embeddings(self, new_embeddings): self.encoder.set_input_embeddings(new_embeddings) self.decoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring # Copied from transformers.models.t5.modeling_t5.T5Model.forward with google-t5/->google/, T5->MT5, t5->mt5 def forward( @@ -894,12 +896,8 @@ def forward( past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.Tensor] = None, decoder_inputs_embeds: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], Seq2SeqModelOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -950,27 +948,16 @@ def forward( >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) >>> last_hidden_states = outputs.last_hidden_state ```""" - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # Encode if needed (training, first prediction pass) if encoder_outputs is None: encoder_outputs = self.encoder( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + return_dict=True, + **kwargs, ) - - hidden_states = encoder_outputs[0] + hidden_states = encoder_outputs.last_hidden_state # Decode decoder_outputs = self.decoder( @@ -980,16 +967,11 @@ def forward( past_key_values=past_key_values, encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, + return_dict=True, + **kwargs, ) - if not return_dict: - return decoder_outputs + encoder_outputs - return Seq2SeqModelOutput( last_hidden_state=decoder_outputs.last_hidden_state, past_key_values=decoder_outputs.past_key_values, @@ -1065,6 +1047,7 @@ def set_input_embeddings(self, new_embeddings): self.encoder.set_input_embeddings(new_embeddings) self.decoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.forward with google-t5/->google/, T5->MT5, t5->mt5 def forward( @@ -1078,12 +1061,8 @@ def forward( inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1140,9 +1119,6 @@ def forward( >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) >>> # studies have shown that owning a dog is good for you. ```""" - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # Encode if needed (training, first prediction pass) if encoder_outputs is None: # Convert encoder inputs in embeddings if needed @@ -1150,18 +1126,10 @@ def forward( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + return_dict=True, + **kwargs, ) - - hidden_states = encoder_outputs[0] + hidden_states = encoder_outputs.last_hidden_state if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: # get decoder inputs from shifting lm labels to the right @@ -1175,14 +1143,12 @@ def forward( past_key_values=past_key_values, encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, + return_dict=True, + **kwargs, ) - sequence_output = decoder_outputs[0] + sequence_output = decoder_outputs.last_hidden_state if self.config.tie_word_embeddings: sequence_output = sequence_output * (self.model_dim**-0.5) @@ -1196,10 +1162,6 @@ def forward( labels = labels.to(lm_logits.device) loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - if not return_dict: - output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs - return ((loss,) + output) if loss is not None else output - return Seq2SeqLMOutput( loss=loss, logits=lm_logits, @@ -1261,6 +1223,7 @@ def set_input_embeddings(self, new_embeddings): self.shared = new_embeddings self.encoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.forward with google-t5/->google/, T5->MT5, t5->mt5 def forward( @@ -1268,10 +1231,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], BaseModelOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1296,15 +1256,11 @@ def forward( >>> outputs = model(input_ids=input_ids) >>> last_hidden_states = outputs.last_hidden_state ```""" - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - encoder_outputs = self.encoder( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + **kwargs, ) return encoder_outputs @@ -1328,6 +1284,7 @@ def __init__(self, config: MT5Config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.forward with T5->MT5, t5->mt5 def forward( @@ -1341,10 +1298,7 @@ def forward( decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple, Seq2SeqSequenceClassifierOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1377,7 +1331,6 @@ def forward( Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if labels is not None: use_cache = False @@ -1406,11 +1359,10 @@ def forward( inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) - sequence_output = outputs[0] + sequence_output = outputs.last_hidden_state eos_mask = input_ids.eq(self.config.eos_token_id).to(sequence_output.device) @@ -1443,9 +1395,6 @@ def forward( elif self.config.problem_type == "multi_label_classification": loss_fct = BCEWithLogitsLoss() loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output return Seq2SeqSequenceClassifierOutput( loss=loss, @@ -1474,6 +1423,7 @@ def __init__(self, config: MT5Config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring # Copied from transformers.models.t5.modeling_t5.T5ForTokenClassification.forward with T5->MT5 def forward( @@ -1482,10 +1432,7 @@ def forward( attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.Tensor], TokenClassifierOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1501,18 +1448,15 @@ def forward( labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - outputs = self.transformer( input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state hidden_states = self.dropout(hidden_states) logits = self.classifier(hidden_states) @@ -1521,10 +1465,6 @@ def forward( loss_fct = CrossEntropyLoss() loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - if not return_dict: - output = (logits, outputs[2:-1]) - return ((loss,) + output) if loss is not None else output - return TokenClassifierOutput( loss=loss, logits=logits, @@ -1576,6 +1516,7 @@ def set_input_embeddings(self, new_embeddings): self.encoder.set_input_embeddings(new_embeddings) self.decoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.forward def forward( @@ -1590,10 +1531,7 @@ def forward( inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], Seq2SeqQuestionAnsweringModelOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1623,7 +1561,6 @@ def forward( Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict use_cache = use_cache if use_cache is not None else self.config.use_cache if start_positions is not None and end_positions is not None: use_cache = False @@ -1640,27 +1577,17 @@ def forward( ) decoder_input_ids = self._shift_right(input_ids) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # Encode if needed (training, first prediction pass) if encoder_outputs is None: encoder_outputs = self.encoder( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + return_dict=True, + **kwargs, ) - hidden_states = encoder_outputs[0] + hidden_states = encoder_outputs.last_hidden_state # Decode decoder_outputs = self.decoder( @@ -1671,9 +1598,8 @@ def forward( encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1700,10 +1626,6 @@ def forward( end_loss = loss_fct(end_logits, end_positions) total_loss = (start_loss + end_loss) / 2 - if not return_dict: - output = (start_logits, end_logits) + decoder_outputs[1:] + encoder_outputs - return ((total_loss,) + output) if total_loss is not None else output - return Seq2SeqQuestionAnsweringModelOutput( loss=total_loss, start_logits=start_logits, diff --git a/src/transformers/models/t5/modeling_t5.py b/src/transformers/models/t5/modeling_t5.py index bc772b443cbe..f008aa008038 100644 --- a/src/transformers/models/t5/modeling_t5.py +++ b/src/transformers/models/t5/modeling_t5.py @@ -16,6 +16,7 @@ import copy import math +from collections.abc import Callable from typing import Optional, Union import torch @@ -37,13 +38,17 @@ Seq2SeqSequenceClassifierOutput, TokenClassifierOutput, ) -from ...modeling_utils import PreTrainedModel +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...utils import ( DUMMY_INPUTS, DUMMY_MASK, + TransformersKwargs, auto_docstring, + can_return_tuple, logging, ) +from ...utils.generic import OutputRecorder, check_model_inputs from .configuration_t5 import T5Config @@ -157,23 +162,57 @@ def forward(self, hidden_states): return hidden_states +# Copied from transformers.models.bert.modeling_bert.eager_attention_forward +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: Optional[float] = None, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + if scaling is None: + scaling = query.size(-1) ** -0.5 + + # Take the dot product between "query" and "key" to get the raw attention scores. + attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling + + if attention_mask is not None: + attention_mask = attention_mask[:, :, :, : key.shape[-2]] + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + class T5Attention(nn.Module): def __init__( self, config: T5Config, has_relative_attention_bias=False, layer_idx: Optional[int] = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal + self.scaling = 1.0 self.has_relative_attention_bias = has_relative_attention_bias self.relative_attention_num_buckets = config.relative_attention_num_buckets self.relative_attention_max_distance = config.relative_attention_max_distance self.d_model = config.d_model - self.key_value_proj_dim = config.d_kv + self.head_dim = config.d_kv self.n_heads = config.num_heads self.dropout = config.dropout_rate - self.inner_dim = self.n_heads * self.key_value_proj_dim + self.inner_dim = self.n_heads * self.head_dim self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -268,34 +307,34 @@ def forward( position_bias=None, past_key_values=None, query_length=None, - use_cache=False, - output_attentions=False, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): - """ - Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). - """ - # Input is (batch_size, seq_length, dim) - # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder) - batch_size, seq_length = hidden_states.shape[:2] - - # if key_value_states are provided this layer is used as a cross-attention layer for the decoder + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder is_cross_attention = key_value_states is not None - query_states = self.q(hidden_states) - query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len + + q_input_shape = (bsz, tgt_len, -1, self.head_dim) + kv_input_shape = (bsz, src_len, -1, self.head_dim) + + # get query proj + query_states = self.q(hidden_states).view(*q_input_shape).transpose(1, 2) - # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache` is_updated = False - if isinstance(past_key_values, EncoderDecoderCache): - is_updated = past_key_values.is_updated.get(self.layer_idx) - if is_cross_attention: - # after the first generated id, we can subsequently re-use all key/value_states from cache - curr_past_key_values = past_key_values.cross_attention_cache + if past_key_values is not None: + if isinstance(past_key_values, EncoderDecoderCache): + is_updated = past_key_values.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + curr_past_key_values = past_key_values.cross_attention_cache + else: + curr_past_key_values = past_key_values.self_attention_cache else: - curr_past_key_values = past_key_values.self_attention_cache - else: - curr_past_key_values = past_key_values + curr_past_key_values = past_key_values current_states = key_value_states if is_cross_attention else hidden_states if is_cross_attention and past_key_values is not None and is_updated: @@ -305,8 +344,8 @@ def forward( else: key_states = self.k(current_states) value_states = self.v(current_states) - key_states = key_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) - value_states = value_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) + key_states = key_states.view(*kv_input_shape).transpose(1, 2) + value_states = value_states.view(*kv_input_shape).transpose(1, 2) if past_key_values is not None: # save all key/value_states to cache to be re-used for fast auto-regressive generation @@ -318,54 +357,66 @@ def forward( if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): past_key_values.is_updated[self.layer_idx] = True - # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 - scores = torch.matmul(query_states, key_states.transpose(3, 2)) - if position_bias is None: - key_length = key_states.shape[-2] + q_length = query_states.shape[-2] + kv_length = key_states.shape[-2] # cache position is 0-indexed so we add 1 to get the real length of queries (aka with past) real_seq_length = query_length if query_length is not None else cache_position[-1] + 1 if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, self.n_heads, seq_length, key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, q_length, kv_length), device=query_states.device, dtype=query_states.dtype ) if self.gradient_checkpointing and self.training: position_bias.requires_grad = True else: position_bias = self.compute_bias( - real_seq_length, key_length, device=scores.device, cache_position=cache_position + real_seq_length, kv_length, device=query_states.device, cache_position=cache_position ) - position_bias = position_bias[:, :, -seq_length:, :] + position_bias = position_bias[:, :, -q_length:, :] if mask is not None: - causal_mask = mask[:, :, :, : key_states.shape[-2]] - position_bias = position_bias + causal_mask - - position_bias_masked = position_bias - scores += position_bias_masked + if mask.dtype == torch.bool: + min_dtype = torch.finfo(query_states.dtype).min + # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type) + mask = torch.where( + mask, torch.tensor(0.0, device=mask.device, dtype=query_states.dtype), min_dtype + ) - # (batch_size, n_heads, seq_length, key_length) - attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) - attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + mask = mask[:, :, :, : key_states.shape[-2]] + position_bias = position_bias + mask - attn_output = torch.matmul(attn_weights, value_states) + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + f"Relative bias attention is only supported by `eager` and `sdpa` but found {self.config._attn_implementation}." + ) + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + position_bias, # position bias acts as mask + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) - attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.view(batch_size, -1, self.inner_dim) + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() attn_output = self.o(attn_output) - - outputs = (attn_output, position_bias) - - if output_attentions: - outputs = outputs + (attn_weights,) - return outputs + return attn_output, position_bias, attn_weights class T5LayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None): super().__init__() self.SelfAttention = T5Attention( - config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx + config, + has_relative_attention_bias=has_relative_attention_bias, + layer_idx=layer_idx, + is_causal=config.is_decoder, ) self.layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -376,29 +427,28 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attn_output, position_bias, attn_weights = self.SelfAttention( normed_hidden_states, mask=attention_mask, position_bias=position_bias, past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, cache_position=cache_position, + **kwargs, ) - hidden_states = hidden_states + self.dropout(attention_output[0]) - outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them - return outputs + hidden_states = hidden_states + self.dropout(attn_output) + return hidden_states, position_bias, attn_weights class T5LayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: Optional[int] = None): super().__init__() - self.EncDecAttention = T5Attention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.EncDecAttention = T5Attention( + config, has_relative_attention_bias=False, layer_idx=layer_idx, is_causal=False + ) self.layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -409,26 +459,23 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, query_length=None, - output_attentions=False, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attn_output, position_bias, attn_weights = self.EncDecAttention( normed_hidden_states, mask=attention_mask, key_value_states=key_value_states, position_bias=position_bias, past_key_values=past_key_values, - use_cache=use_cache, query_length=query_length, - output_attentions=output_attentions, cache_position=cache_position, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + hidden_states = hidden_states + self.dropout(attn_output) + return hidden_states, position_bias, attn_weights class T5Block(GradientCheckpointingLayer): @@ -453,22 +500,18 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - return_dict=True, cache_position=None, + **kwargs: Unpack[TransformersKwargs], ): - self_attention_outputs = self.layer[0]( + self_attention_output, self_attention_position_bias, _ = self.layer[0]( hidden_states, attention_mask=attention_mask, position_bias=position_bias, past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, cache_position=cache_position, + **kwargs, ) - hidden_states = self_attention_outputs[0] - attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights + hidden_states = self_attention_output # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -480,18 +523,18 @@ def forward( hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) do_cross_attention = self.is_decoder and encoder_hidden_states is not None + cross_attention_position_bias = None if do_cross_attention: - cross_attention_outputs = self.layer[1]( + cross_attention_output, cross_attention_position_bias, _ = self.layer[1]( hidden_states, key_value_states=encoder_hidden_states, attention_mask=encoder_attention_mask, position_bias=encoder_decoder_position_bias, past_key_values=past_key_values, query_length=cache_position[-1] + 1, - use_cache=use_cache, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = cross_attention_outputs[0] + hidden_states = cross_attention_output # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -502,9 +545,6 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - # Keep cross-attention outputs and relative position weights - attention_outputs = attention_outputs + cross_attention_outputs[1:] - # Apply Feed Forward layer hidden_states = self.layer[-1](hidden_states) @@ -517,11 +557,7 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - outputs = (hidden_states,) - - return ( - outputs + attention_outputs - ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + return hidden_states, self_attention_position_bias, cross_attention_position_bias class T5ClassificationHead(nn.Module): @@ -547,11 +583,21 @@ class T5PreTrainedModel(PreTrainedModel): config: T5Config base_model_prefix = "transformer" supports_gradient_checkpointing = True + _supports_attention_backend = True + _supports_sdpa = True + _supports_flash_attn = False + _supports_flex_attn = False _can_compile_fullgraph = True _no_split_modules = ["T5Block"] _keep_in_fp32_modules = ["wo"] + _can_record_outputs = { + "hidden_states": OutputRecorder(T5Block, index=0), + "attentions": OutputRecorder(T5LayerSelfAttention, index=-1), + "cross_attentions": OutputRecorder(T5LayerCrossAttention, index=-1), + } + @property def dummy_inputs(self): input_ids = torch.tensor(DUMMY_INPUTS) @@ -660,6 +706,7 @@ def __init__(self, config): def set_input_embeddings(self, new_embeddings): self.embed_tokens = new_embeddings + @check_model_inputs def forward( self, input_ids=None, @@ -669,69 +716,37 @@ def forward( inputs_embeds=None, past_key_values=None, use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, cache_position=None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - use_cache = use_cache if use_cache is not None else self.config.use_cache - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if input_ids is not None and inputs_embeds is not None: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError( - f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time" - ) - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache else: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") + use_cache = False - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: - if self.embed_tokens is None: - raise ValueError("You have to initialize the model with valid token embeddings") inputs_embeds = self.embed_tokens(input_ids) - batch_size, seq_length = input_shape - - if use_cache is True: - if not self.is_decoder: - raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder") - - if self.is_decoder: - if use_cache and past_key_values is None: - if self.config.is_encoder_decoder: - past_key_values = EncoderDecoderCache( - DynamicCache(config=self.config), DynamicCache(config=self.config) - ) - else: - past_key_values = DynamicCache(config=self.config) + if use_cache and past_key_values is None: + past_key_values = ( + EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + if self.config.is_encoder_decoder + else DynamicCache(config=self.config) + ) elif not self.is_decoder: # do not pass cache object down the line for encoder stack # it messes indexing later in decoder-stack because cache object is modified in-place past_key_values = None + device = inputs_embeds.device + seq_length = inputs_embeds.shape[1] + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 if cache_position is None: - cache_position = torch.arange( - past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device - ) + cache_position = torch.arange(past_key_values_length, past_key_values_length + seq_length, device=device) if self.config.is_decoder: attention_mask = create_causal_mask( @@ -759,19 +774,11 @@ def forward( encoder_hidden_states=encoder_hidden_states, ) - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - all_cross_attentions = () if (output_attentions and self.is_decoder) else None position_bias = None encoder_decoder_position_bias = None - hidden_states = self.dropout(inputs_embeds) - for layer_module in self.block: - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_outputs = layer_module( + hidden_states, self_attention_position_bias, cross_attention_position_bias = layer_module( hidden_states, attention_mask, position_bias, @@ -779,51 +786,21 @@ def forward( encoder_extended_attention_mask, encoder_decoder_position_bias, # as a positional argument for gradient checkpointing past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - return_dict=return_dict, cache_position=cache_position, + **kwargs, ) - hidden_states = layer_outputs[0] - - # We share the position biases between the layers - the first layer store them - # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), - # (cross-attention position bias), (cross-attention weights) - position_bias = layer_outputs[1] + # We share the position biases between the layers + position_bias = self_attention_position_bias if self.is_decoder and encoder_hidden_states is not None: - encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2] - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[2],) - if self.is_decoder: - all_cross_attentions = all_cross_attentions + (layer_outputs[4],) + encoder_decoder_position_bias = cross_attention_position_bias hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.dropout(hidden_states) - # Add last layer - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - past_key_values, - all_hidden_states, - all_attentions, - all_cross_attentions, - ] - if v is not None - ) return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=past_key_values, - hidden_states=all_hidden_states, - attentions=all_attentions, - cross_attentions=all_cross_attentions, ) @@ -864,6 +841,7 @@ def set_input_embeddings(self, new_embeddings): self.encoder.set_input_embeddings(new_embeddings) self.decoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring def forward( self, @@ -875,12 +853,8 @@ def forward( past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.Tensor] = None, decoder_inputs_embeds: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], Seq2SeqModelOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -931,27 +905,16 @@ def forward( >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) >>> last_hidden_states = outputs.last_hidden_state ```""" - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # Encode if needed (training, first prediction pass) if encoder_outputs is None: encoder_outputs = self.encoder( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] + hidden_states = encoder_outputs.last_hidden_state # Decode decoder_outputs = self.decoder( @@ -961,16 +924,11 @@ def forward( past_key_values=past_key_values, encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, + return_dict=True, + **kwargs, ) - if not return_dict: - return decoder_outputs + encoder_outputs - return Seq2SeqModelOutput( last_hidden_state=decoder_outputs.last_hidden_state, past_key_values=decoder_outputs.past_key_values, @@ -1029,6 +987,7 @@ def set_input_embeddings(self, new_embeddings): self.encoder.set_input_embeddings(new_embeddings) self.decoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring def forward( self, @@ -1041,12 +1000,8 @@ def forward( inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1103,9 +1058,6 @@ def forward( >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) >>> # studies have shown that owning a dog is good for you. ```""" - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # Encode if needed (training, first prediction pass) if encoder_outputs is None: # Convert encoder inputs in embeddings if needed @@ -1113,18 +1065,10 @@ def forward( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] + hidden_states = encoder_outputs.last_hidden_state if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: # get decoder inputs from shifting lm labels to the right @@ -1138,14 +1082,12 @@ def forward( past_key_values=past_key_values, encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, + return_dict=True, + **kwargs, ) - sequence_output = decoder_outputs[0] + sequence_output = decoder_outputs.last_hidden_state if self.config.tie_word_embeddings: sequence_output = sequence_output * (self.model_dim**-0.5) @@ -1159,10 +1101,6 @@ def forward( labels = labels.to(lm_logits.device) loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - if not return_dict: - output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs - return ((loss,) + output) if loss is not None else output - return Seq2SeqLMOutput( loss=loss, logits=lm_logits, @@ -1203,16 +1141,14 @@ def set_input_embeddings(self, new_embeddings): self.shared = new_embeddings self.encoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], BaseModelOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1237,15 +1173,11 @@ def forward( >>> outputs = model(input_ids=input_ids) >>> last_hidden_states = outputs.last_hidden_state ```""" - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - encoder_outputs = self.encoder( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + **kwargs, ) return encoder_outputs @@ -1268,6 +1200,7 @@ def __init__(self, config: T5Config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def forward( self, @@ -1280,10 +1213,7 @@ def forward( decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple, Seq2SeqSequenceClassifierOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1316,7 +1246,6 @@ def forward( Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if labels is not None: use_cache = False @@ -1345,11 +1274,10 @@ def forward( inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) - sequence_output = outputs[0] + sequence_output = outputs.last_hidden_state eos_mask = input_ids.eq(self.config.eos_token_id).to(sequence_output.device) @@ -1382,9 +1310,6 @@ def forward( elif self.config.problem_type == "multi_label_classification": loss_fct = BCEWithLogitsLoss() loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output return Seq2SeqSequenceClassifierOutput( loss=loss, @@ -1412,6 +1337,7 @@ def __init__(self, config: T5Config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def forward( self, @@ -1419,10 +1345,7 @@ def forward( attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.Tensor], TokenClassifierOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1438,18 +1361,15 @@ def forward( labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - outputs = self.transformer( input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state hidden_states = self.dropout(hidden_states) logits = self.classifier(hidden_states) @@ -1458,10 +1378,6 @@ def forward( loss_fct = CrossEntropyLoss() loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - if not return_dict: - output = (logits, outputs[2:-1]) - return ((loss,) + output) if loss is not None else output - return TokenClassifierOutput( loss=loss, logits=logits, @@ -1510,6 +1426,7 @@ def set_input_embeddings(self, new_embeddings): self.encoder.set_input_embeddings(new_embeddings) self.decoder.set_input_embeddings(new_embeddings) + @can_return_tuple @auto_docstring def forward( self, @@ -1523,10 +1440,7 @@ def forward( inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple[torch.FloatTensor], Seq2SeqQuestionAnsweringModelOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): @@ -1556,7 +1470,6 @@ def forward( Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict use_cache = use_cache if use_cache is not None else self.config.use_cache if start_positions is not None and end_positions is not None: use_cache = False @@ -1573,27 +1486,16 @@ def forward( ) decoder_input_ids = self._shift_right(input_ids) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # Encode if needed (training, first prediction pass) if encoder_outputs is None: encoder_outputs = self.encoder( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + return_dict=True, + **kwargs, ) - - hidden_states = encoder_outputs[0] + hidden_states = encoder_outputs.last_hidden_state # Decode decoder_outputs = self.decoder( @@ -1604,9 +1506,8 @@ def forward( encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1633,10 +1534,6 @@ def forward( end_loss = loss_fct(end_logits, end_positions) total_loss = (start_loss + end_loss) / 2 - if not return_dict: - output = (start_logits, end_logits) + decoder_outputs[1:] + encoder_outputs - return ((total_loss,) + output) if total_loss is not None else output - return Seq2SeqQuestionAnsweringModelOutput( loss=total_loss, start_logits=start_logits, diff --git a/tests/models/mt5/test_modeling_mt5.py b/tests/models/mt5/test_modeling_mt5.py index cd0ab3677643..5482f4ad39bb 100644 --- a/tests/models/mt5/test_modeling_mt5.py +++ b/tests/models/mt5/test_modeling_mt5.py @@ -739,6 +739,10 @@ def test_model_from_pretrained(self): def test_model_base_model_prefix(self): pass + @unittest.skip(reason="MT5 always needs to include a relative bias into the mask") + def test_sdpa_can_dispatch_on_flash(self): + pass + # Copied from tests.models.t5.test_modeling_t5.T5EncoderOnlyModelTester with T5->MT5 class MT5EncoderOnlyModelTester: @@ -924,6 +928,10 @@ def is_pipeline_test_to_skip( return False + @unittest.skip(reason="MT5 always needs to include a relative bias into the mask") + def test_sdpa_can_dispatch_on_flash(self): + pass + @require_torch @require_sentencepiece @@ -943,7 +951,7 @@ def test_small_integration_test(self): >>> score = t5_model.score(inputs=["Hello there"], targets=["Hi I am"], vocabulary=vocab) """ - model = AutoModelForSeq2SeqLM.from_pretrained("google/mt5-small", return_dict=True).to(torch_device) + model = AutoModelForSeq2SeqLM.from_pretrained("google/mt5-small", attn_implementation="eager").to(torch_device) tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") input_ids = tokenizer("Hello there", return_tensors="pt").input_ids diff --git a/tests/models/t5/test_modeling_t5.py b/tests/models/t5/test_modeling_t5.py index 29f99c367510..7cffab9b0a5d 100644 --- a/tests/models/t5/test_modeling_t5.py +++ b/tests/models/t5/test_modeling_t5.py @@ -748,6 +748,10 @@ def test_model_from_pretrained(self): def test_model_base_model_prefix(self): pass + @unittest.skip(reason="T5 always needs to include a relative bias into the mask") + def test_sdpa_can_dispatch_on_flash(self): + pass + class T5EncoderOnlyModelTester: def __init__( @@ -931,6 +935,10 @@ def is_pipeline_test_to_skip( return False + @unittest.skip(reason="T5 always needs to include a relative bias into the mask") + def test_sdpa_can_dispatch_on_flash(self): + pass + def use_task_specific_params(model, task): task_params = model.config.task_specific_params[task] @@ -1066,7 +1074,9 @@ def test_small_integration_test(self): >>> score = t5_model.score(inputs=["Hello there"], targets=["Hi I am"], vocabulary=vocab) """ - model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small").to(torch_device) + model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small", attn_implementation="eager").to( + torch_device + ) tokenizer = T5Tokenizer.from_pretrained("google-t5/t5-small") input_ids = tokenizer("Hello there", return_tensors="pt").input_ids @@ -1097,7 +1107,9 @@ def test_small_v1_1_integration_test(self): >>> score = t5_model.score(inputs=["Hello there"], targets=["Hi I am"], vocabulary=vocab) """ - model = T5ForConditionalGeneration.from_pretrained("google/t5-v1_1-small").to(torch_device) + model = T5ForConditionalGeneration.from_pretrained("google/t5-v1_1-small", attn_implementation="eager").to( + torch_device + ) tokenizer = T5Tokenizer.from_pretrained("google/t5-v1_1-small") input_ids = tokenizer("Hello there", return_tensors="pt").input_ids @@ -1121,7 +1133,9 @@ def test_small_byt5_integration_test(self): >>> score = t5_model.score(inputs=["Hello there"], targets=["Hi I am"], vocabulary=vocab) """ - model = T5ForConditionalGeneration.from_pretrained("google/byt5-small").to(torch_device) + model = T5ForConditionalGeneration.from_pretrained("google/byt5-small", attn_implementation="eager").to( + torch_device + ) tokenizer = ByT5Tokenizer.from_pretrained("google/byt5-small") input_ids = tokenizer("Hello there", return_tensors="pt").input_ids