diff --git a/src/transformers/models/longt5/modeling_longt5.py b/src/transformers/models/longt5/modeling_longt5.py index b3e082789e9c..016895e6cfa0 100644 --- a/src/transformers/models/longt5/modeling_longt5.py +++ b/src/transformers/models/longt5/modeling_longt5.py @@ -15,6 +15,7 @@ import copy import math +from collections.abc import Callable from typing import Any import torch @@ -33,10 +34,12 @@ Seq2SeqLMOutput, Seq2SeqModelOutput, ) -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, is_torchdynamo_compiling, logging, @@ -47,6 +50,35 @@ logger = logging.get_logger(__name__) +# 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: torch.Tensor | None, + scaling: float | None = 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: + 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 + + # TODO: Update before the merge @@ -303,9 +335,12 @@ def __init__( config: LongT5Config, has_relative_attention_bias=False, layer_idx: int | None = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal 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 @@ -314,6 +349,8 @@ def __init__( self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # LongT5 folds the relative position bias into the attention scores and does not scale the query/key dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -404,8 +441,7 @@ def forward( key_value_states=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): """ Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). @@ -452,43 +488,60 @@ def forward( 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] if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, query_states.shape[1], input_shape[1], key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, input_shape[1], key_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( - input_shape[1], key_length, device=scores.device, past_seen_tokens=past_seen_tokens + input_shape[1], key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + # `sdpa` may materialize a boolean mask (True = keep). Turn it into an additive float mask so it + # can be folded into the relative position bias, just like the `eager` float mask. + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) position_bias = position_bias + causal_mask - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) + # LongT5 uses a relative attention bias that is added to the attention scores. This is passed as the additive + # attention mask so that it works with the different attention implementations. As it is always non-`None`, + # `is_causal` is never inferred, so the causal behavior is fully encoded in the bias itself. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "LongT5 adds a relative position bias on top of the attention scores, which is only supported by the " + f"`eager` and `sdpa` attention implementations, but got `{self.config._attn_implementation}`." + ) + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - attn_output = torch.matmul(attn_weights, value_states) + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + position_bias, + 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.reshape(*input_shape, -1) + attn_output = attn_output.reshape(*input_shape, -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 LongT5LocalAttention(nn.Module): @@ -907,7 +960,10 @@ class LongT5LayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): super().__init__() self.SelfAttention = LongT5Attention( - 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 = LongT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -918,22 +974,18 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attention_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, + **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(attention_output) + return hidden_states, position_bias, attn_weights class LongT5LayerLocalSelfAttention(nn.Module): @@ -1000,7 +1052,9 @@ def forward( class LongT5LayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: int | None = None): super().__init__() - self.EncDecAttention = LongT5Attention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.EncDecAttention = LongT5Attention( + config, has_relative_attention_bias=False, layer_idx=layer_idx, is_causal=False + ) self.layer_norm = LongT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -1011,21 +1065,19 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attention_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, - output_attentions=output_attentions, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + layer_output = hidden_states + self.dropout(attention_output) + return layer_output, position_bias, attn_weights class LongT5Block(GradientCheckpointingLayer): @@ -1348,11 +1400,12 @@ def forward( 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) + # The self- and cross-attention wrappers now always return their attention weights, so the decoder layer + # output layout is: hidden-states, self-attention position bias, self-attention weights, + # cross-attention position bias, cross-attention weights. position_bias = layer_outputs[1] if self.is_decoder and encoder_hidden_states is not None: - encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2] + encoder_decoder_position_bias = layer_outputs[3] if output_attentions: all_attentions = all_attentions + (layer_outputs[2],) diff --git a/src/transformers/models/mt5/modeling_mt5.py b/src/transformers/models/mt5/modeling_mt5.py index 96896c90da35..645f2762428e 100644 --- a/src/transformers/models/mt5/modeling_mt5.py +++ b/src/transformers/models/mt5/modeling_mt5.py @@ -15,6 +15,7 @@ import copy import math +from collections.abc import Callable import torch from torch import nn @@ -35,14 +36,46 @@ Seq2SeqSequenceClassifierOutput, TokenClassifierOutput, ) -from ...modeling_utils import PreTrainedModel -from ...utils import DUMMY_INPUTS, DUMMY_MASK, auto_docstring, logging, torch_compilable_check +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import DUMMY_INPUTS, DUMMY_MASK, TransformersKwargs, auto_docstring, logging, torch_compilable_check +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_mt5 import MT5Config logger = logging.get_logger(__name__) +# 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: torch.Tensor | None, + scaling: float | None = 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: + 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.T5LayerNorm with T5->MT5 class MT5LayerNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): @@ -148,9 +181,12 @@ def __init__( config: MT5Config, has_relative_attention_bias=False, layer_idx: int | None = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal 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 @@ -159,6 +195,8 @@ def __init__( self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # MT5 folds the relative position bias into the attention scores and does not scale the query/key dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -249,8 +287,7 @@ def forward( key_value_states=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): """ Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). @@ -297,43 +334,60 @@ def forward( 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] if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, query_states.shape[1], input_shape[1], key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, input_shape[1], key_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( - input_shape[1], key_length, device=scores.device, past_seen_tokens=past_seen_tokens + input_shape[1], key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + # `sdpa` may materialize a boolean mask (True = keep). Turn it into an additive float mask so it + # can be folded into the relative position bias, just like the `eager` float mask. + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) position_bias = position_bias + causal_mask - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) - - attn_output = torch.matmul(attn_weights, value_states) + # MT5 uses a relative attention bias that is added to the attention scores. This is passed as the additive + # attention mask so that it works with the different attention implementations. As it is always non-`None`, + # `is_causal` is never inferred, so the causal behavior is fully encoded in the bias itself. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "MT5 adds a relative position bias on top of the attention scores, which is only supported by the " + f"`eager` and `sdpa` attention implementations, but got `{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, + 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.reshape(*input_shape, -1) + attn_output = attn_output.reshape(*input_shape, -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 @@ -341,7 +395,10 @@ class MT5LayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = 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) @@ -352,29 +409,27 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attention_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, + **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(attention_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: int | None = 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) @@ -385,21 +440,19 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attention_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, - output_attentions=output_attentions, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + layer_output = hidden_states + self.dropout(attention_output) + return layer_output, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5Block with T5->MT5 @@ -425,21 +478,15 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - return_dict=True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - self_attention_outputs = self.layer[0]( + hidden_states, self_attn_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, + **kwargs, ) - hidden_states = self_attention_outputs[0] - attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -450,17 +497,17 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + cross_attn_position_bias = None do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: - cross_attention_outputs = self.layer[1]( + hidden_states, cross_attn_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, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = cross_attention_outputs[0] # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -471,9 +518,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) @@ -486,11 +530,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_attn_position_bias, cross_attn_position_bias # Copied from transformers.models.t5.modeling_t5.T5ClassificationHead with T5->MT5 @@ -523,6 +563,17 @@ class MT5PreTrainedModel(PreTrainedModel): _no_split_modules = ["MT5Block"] _keep_in_fp32_modules = ["wo"] + _supports_attention_backend = True + _supports_flash_attn = False + _supports_flex_attn = False + _supports_sdpa = True + + _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) @@ -632,6 +683,9 @@ def __init__(self, config): def set_input_embeddings(self, new_embeddings): self.embed_tokens = new_embeddings + @merge_with_config_defaults + @capture_outputs + @auto_docstring def forward( self, input_ids=None, @@ -641,17 +695,9 @@ def forward( inputs_embeds=None, past_key_values=None, use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=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.return_dict if input_ids is not None and inputs_embeds is not None: err_msg_prefix = "decoder_" if self.is_decoder else "" @@ -679,8 +725,6 @@ def forward( 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") @@ -699,6 +743,10 @@ def forward( past_key_values = None if self.config.is_decoder: + # MT5 folds the relative position bias into the attention scores and always feeds it through the + # `attention_mask` argument, so `sdpa` can never rely on its `is_causal` shortcut. The dummy mask function + # is a no-op but forces the causal mask to always be materialized instead of being skipped. + dummy_and_mask_function = lambda *args: torch.tensor(True, dtype=torch.bool) # noqa: E731 attention_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, @@ -706,6 +754,7 @@ def forward( past_key_values=past_key_values.self_attention_cache if isinstance(past_key_values, EncoderDecoderCache) else past_key_values, + and_mask_function=dummy_and_mask_function, ) else: attention_mask = create_bidirectional_mask( @@ -723,19 +772,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, @@ -743,50 +786,20 @@ 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, + **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 - the first layer stores them + 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, ) @@ -845,6 +858,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( @@ -858,11 +872,8 @@ def forward( inputs_embeds: torch.Tensor | None = None, decoder_inputs_embeds: torch.Tensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you @@ -913,7 +924,6 @@ def forward( >>> 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.return_dict # Encode if needed (training, first prediction pass) if encoder_outputs is None: @@ -921,11 +931,9 @@ 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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -943,14 +951,9 @@ 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, + **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, @@ -1024,6 +1027,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, @@ -1037,11 +1041,8 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqLMOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqLMOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you @@ -1098,7 +1099,6 @@ def forward( >>> # 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.return_dict # Encode if needed (training, first prediction pass) if encoder_outputs is None: @@ -1107,11 +1107,9 @@ 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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -1133,9 +1131,7 @@ 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, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1149,10 +1145,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, @@ -1214,6 +1206,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( @@ -1221,11 +1214,8 @@ def forward( input_ids: torch.LongTensor | None = None, attention_mask: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | BaseModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you @@ -1249,15 +1239,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.return_dict - - encoder_outputs = self.encoder( + encoder_outputs: BaseModelOutput = 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 @@ -1281,6 +1267,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( @@ -1294,11 +1281,8 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple | Seq2SeqSequenceClassifierOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqSequenceClassifierOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you @@ -1330,7 +1314,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.return_dict if labels is not None: use_cache = False @@ -1350,7 +1333,7 @@ def forward( ) decoder_input_ids = self._shift_right(input_ids) - outputs = self.transformer( + outputs: Seq2SeqModelOutput = self.transformer( input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids, @@ -1359,9 +1342,7 @@ 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, + **kwargs, ) sequence_output = outputs[0] @@ -1403,9 +1384,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, @@ -1434,6 +1412,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( @@ -1442,11 +1421,8 @@ def forward( attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, labels: torch.Tensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.Tensor] | TokenClassifierOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> TokenClassifierOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you @@ -1461,15 +1437,11 @@ 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.return_dict - - outputs = self.transformer( + outputs: BaseModelOutput = 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, + **kwargs, ) hidden_states = outputs[0] @@ -1481,10 +1453,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, @@ -1534,6 +1502,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( @@ -1548,11 +1517,8 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqQuestionAnsweringModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqQuestionAnsweringModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -1581,7 +1547,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.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 @@ -1598,20 +1563,15 @@ 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.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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -1629,9 +1589,7 @@ 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, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1658,10 +1616,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/pix2struct/modeling_pix2struct.py b/src/transformers/models/pix2struct/modeling_pix2struct.py index e4850ec8cf84..fd5c2d3e082e 100644 --- a/src/transformers/models/pix2struct/modeling_pix2struct.py +++ b/src/transformers/models/pix2struct/modeling_pix2struct.py @@ -14,6 +14,7 @@ """Pix2Struct modeling file""" import math +from collections.abc import Callable import torch from torch import nn @@ -31,12 +32,13 @@ Seq2SeqLMOutput, Seq2SeqModelOutput, ) -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, - is_torchdynamo_compiling, logging, ) from .configuration_pix2struct import Pix2StructConfig, Pix2StructTextConfig, Pix2StructVisionConfig @@ -112,11 +114,15 @@ def forward(self, flattened_patches: torch.Tensor) -> torch.Tensor: class Pix2StructVisionAttention(nn.Module): def __init__(self, config): super().__init__() + self.config = config self.hidden_size = config.hidden_size self.key_value_proj_dim = config.d_kv self.n_heads = config.num_attention_heads self.dropout = config.attention_dropout self.inner_dim = self.n_heads * self.key_value_proj_dim + # Pix2Struct does not scale the q/k dot product (T5-style), so the interface scaling is a no-op. + self.scaling = 1.0 + self.is_causal = False self.query = nn.Linear(self.hidden_size, self.inner_dim, bias=False) self.key = nn.Linear(self.hidden_size, self.inner_dim, bias=False) @@ -129,74 +135,47 @@ def forward( self, hidden_states, attention_mask=None, - position_bias=None, - output_attentions=False, + **kwargs: Unpack[TransformersKwargs], ): """ Self-attention block """ # Input is (batch_size, seq_length, dim) - # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length) batch_size, seq_length = hidden_states.shape[:2] + hidden_shape = (batch_size, seq_length, -1, self.key_value_proj_dim) - def to_projection_shape(states): - """projection""" - return states.contiguous().view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) - - # get query states # (batch_size, n_heads, seq_length, dim_per_head) - query_states = to_projection_shape(self.query(hidden_states)) - - # get key/value states - key_states = to_projection_shape(self.key(hidden_states)) - value_states = to_projection_shape(self.value(hidden_states)) - - # 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: - position_bias = torch.zeros( - (1, self.n_heads, seq_length, seq_length), device=scores.device, dtype=scores.dtype - ) - if self.gradient_checkpointing and self.training: - position_bias.requires_grad = True - - if attention_mask.dim() == 2: - position_bias = position_bias + attention_mask[:, None, None, :].to(position_bias.device) - elif attention_mask is not None: - # (batch_size, n_heads, seq_length, key_length) - position_bias = position_bias + attention_mask.to(position_bias.device) - elif not is_torchdynamo_compiling(): - attention_mask = torch.ones( - (batch_size, seq_length), device=position_bias.device, dtype=position_bias.dtype + query_states = self.query(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.key(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.value(hidden_states).view(hidden_shape).transpose(1, 2) + + # Pix2Struct's vision attention is bidirectional and only relies on the padding mask. As the mask is always + # passed as an additive `attention_mask`, `is_causal` is never inferred. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "Pix2Struct only supports the `eager` and `sdpa` attention implementations, but got " + f"`{self.config._attn_implementation}`." ) - position_bias = position_bias + attention_mask.to(position_bias.device) - - position_bias = 1 - position_bias - - position_bias_masked = position_bias.masked_fill(position_bias == 1, torch.finfo(scores.dtype).min) - scores += position_bias_masked - scores = torch.max(scores, torch.tensor(torch.finfo(scores.dtype).min)) - - # (batch_size, n_heads, seq_length, key_length) - attn_weights = nn.functional.softmax(scores, dim=-1, dtype=torch.float32).type_as(scores) - - # (batch_size, n_heads, seq_length, key_length) - attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) - - attn_output = torch.matmul(attn_weights, value_states) + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) # (batch_size, seq_length, dim) - attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim) - + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.output(attn_output) - outputs = (attn_output,) + (position_bias,) - - if output_attentions: - outputs = outputs + (attn_weights,) - return outputs + return attn_output, attn_weights # Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5DenseGatedActDense->Pix2StructVisionMlp,T5Config->Pix2StructVisionConfig,config.d_model->config.hidden_size,dropout_rate->dropout_rate @@ -243,20 +222,18 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - output_attentions: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor]: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: residual = hidden_states # in Pix2StructVision, layernorm is applied before self-attention hidden_states = self.pre_attention_layer_norm(hidden_states) - self_attention_outputs = self.attention( + attention_output, attn_weights = self.attention( hidden_states, attention_mask=attention_mask, - output_attentions=output_attentions, + **kwargs, ) - attention_output = self_attention_outputs[0] - outputs = self_attention_outputs[1:] # add self attentions if we output attention weights # first residual connection hidden_states = attention_output + residual @@ -265,9 +242,7 @@ def forward( layer_output = self.pre_mlp_layer_norm(hidden_states) layer_output = self.mlp(layer_output) + hidden_states # second residual connection - outputs = (layer_output,) + outputs - - return outputs + return layer_output, attn_weights class Pix2StructVisionEncoder(nn.Module): @@ -283,27 +258,23 @@ def forward( attention_mask: torch.Tensor | None = None, output_attentions: bool = False, output_hidden_states: bool = False, - return_dict: bool = True, - ) -> tuple | BaseModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None - for i, layer_module in enumerate(self.layer): + for layer_module in self.layer: if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - layer_outputs = layer_module(hidden_states, attention_mask, output_attentions) - - hidden_states = layer_outputs[0] + hidden_states, attn_weights = layer_module(hidden_states, attention_mask, **kwargs) if output_attentions: - all_self_attentions = all_self_attentions + (layer_outputs[1],) + all_self_attentions = all_self_attentions + (attn_weights,) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if not return_dict: - return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) return BaseModelOutput( last_hidden_state=hidden_states, hidden_states=all_hidden_states, @@ -317,6 +288,8 @@ class Pix2StructPreTrainedModel(PreTrainedModel): input_modalities = ("image", "text") _can_compile_fullgraph = False + _supports_attention_backend = True + _supports_sdpa = True @property def dummy_inputs(self): @@ -506,12 +479,17 @@ def forward( embedding_output = self.embeddings(flattened_patches) + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + ) + encoder_outputs = self.encoder( embedding_output, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, ) sequence_output = encoder_outputs[0] sequence_output = self.layernorm(sequence_output) @@ -573,9 +551,46 @@ 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: torch.Tensor | None, + scaling: float | None = 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: + 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 Pix2StructTextAttention(nn.Module): - def __init__(self, config: Pix2StructTextConfig, has_relative_attention_bias=False, layer_idx: int | None = None): + def __init__( + self, + config: Pix2StructTextConfig, + has_relative_attention_bias=False, + layer_idx: int | None = None, + is_causal: bool = False, + ): super().__init__() + self.config = config + self.is_causal = is_causal 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 @@ -584,6 +599,8 @@ def __init__(self, config: Pix2StructTextConfig, has_relative_attention_bias=Fal self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # Pix2Struct folds the relative position bias into the attention scores and does not scale the q/k dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None: logger.warning_once( @@ -677,15 +694,15 @@ def forward( key_value_states=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **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) or (batch_size, 1, seq_length, key_length) (causal decoder) - batch_size, seq_length = hidden_states.shape[:2] + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.key_value_proj_dim) past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0 # We clone here for StaticCache, as we get the value before updating it, but use it after and it's the same ref past_seen_tokens = past_seen_tokens.clone() if isinstance(past_seen_tokens, torch.Tensor) else past_seen_tokens @@ -693,11 +710,11 @@ def forward( # 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.query(hidden_states) - query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) + query_states = self.query(hidden_states).view(hidden_shape).transpose(1, 2) # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache` - if past_key_values is not None and isinstance(past_key_values, EncoderDecoderCache): + 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 @@ -708,60 +725,75 @@ def forward( 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 and is_updated: + if is_cross_attention and past_key_values is not None and is_updated: # reuse k,v, cross_attentions key_states = curr_past_key_values.layers[self.layer_idx].keys value_states = curr_past_key_values.layers[self.layer_idx].values else: - key_states = self.key(current_states) - value_states = self.value(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) + kv_shape = (*current_states.shape[:-1], -1, self.key_value_proj_dim) + key_states = self.key(current_states).view(kv_shape).transpose(1, 2) + value_states = self.value(current_states).view(kv_shape).transpose(1, 2) if past_key_values is not None: key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx) # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls - if is_cross_attention: + 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] 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, input_shape[1], key_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( - seq_length, key_length, device=scores.device, past_seen_tokens=past_seen_tokens + input_shape[1], key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + # `sdpa` may materialize a boolean mask (True = keep). Turn it into an additive float mask so it + # can be folded into the relative position bias, just like the `eager` float mask. + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) position_bias = position_bias + causal_mask - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) - - attn_output = torch.matmul(attn_weights, value_states) + # Pix2Struct uses a relative attention bias that is added to the attention scores. This is passed as the + # additive attention mask so that it works with the different attention implementations. As it is always + # non-`None`, `is_causal` is never inferred, so the causal behavior is fully encoded in the bias itself. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "Pix2Struct adds a relative position bias on top of the attention scores, which is only supported " + f"by the `eager` and `sdpa` attention implementations, but got `{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, + 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(*input_shape, -1).contiguous() attn_output = self.output(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 T5LayerNorm->Pix2StructLayerNorm,T5Attention->Pix2StructTextAttention,T5LayerSelfAttention->Pix2StructTextLayerSelfAttention,self.SelfAttention->self.attention,config.d_model->config.hidden_size @@ -769,7 +801,10 @@ class Pix2StructTextLayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): super().__init__() self.attention = Pix2StructTextAttention( - 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 = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -780,29 +815,27 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.attention( + attention_output, position_bias, attn_weights = self.attention( normed_hidden_states, mask=attention_mask, position_bias=position_bias, past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, + **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(attention_output) + return hidden_states, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5LayerNorm->Pix2StructLayerNorm,T5Attention->Pix2StructTextAttention,T5LayerCrossAttention->Pix2StructTextLayerCrossAttention,self.EncDecAttention->self.attention,config.d_model->config.hidden_size class Pix2StructTextLayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: int | None = None): super().__init__() - self.attention = Pix2StructTextAttention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.attention = Pix2StructTextAttention( + config, has_relative_attention_bias=False, layer_idx=layer_idx, is_causal=False + ) self.layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -813,21 +846,19 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.attention( + attention_output, position_bias, attn_weights = self.attention( normed_hidden_states, mask=attention_mask, key_value_states=key_value_states, position_bias=position_bias, past_key_values=past_key_values, - output_attentions=output_attentions, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + layer_output = hidden_states + self.dropout(attention_output) + return layer_output, position_bias, attn_weights class Pix2StructTextBlock(GradientCheckpointingLayer): @@ -1031,11 +1062,16 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: + # Pix2Struct folds the relative position bias into the attention scores and always feeds it through the + # `attention_mask` argument, so `sdpa` can never rely on its `is_causal` shortcut. The dummy mask function + # is a no-op but forces the causal mask to always be materialized instead of being skipped. + dummy_and_mask_function = lambda *args: torch.tensor(True, dtype=torch.bool) # noqa: E731 causal_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + and_mask_function=dummy_and_mask_function, ) else: causal_mask = attention_mask[:, None, None, :] @@ -1077,11 +1113,12 @@ def forward( 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) + # The self- and cross-attention wrappers now always return their attention weights, so the layer output + # layout is: hidden-states, self-attention position bias, self-attention weights, + # cross-attention position bias, cross-attention weights. position_bias = layer_outputs[1] if encoder_hidden_states is not None: - encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2] + encoder_decoder_position_bias = layer_outputs[3] if output_attentions: all_attentions = all_attentions + (layer_outputs[2],) diff --git a/src/transformers/models/pop2piano/modeling_pop2piano.py b/src/transformers/models/pop2piano/modeling_pop2piano.py index 312891401bbc..7bedb5d4c84b 100644 --- a/src/transformers/models/pop2piano/modeling_pop2piano.py +++ b/src/transformers/models/pop2piano/modeling_pop2piano.py @@ -15,6 +15,7 @@ import copy import math +from collections.abc import Callable import torch from torch import nn @@ -29,14 +30,46 @@ from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput -from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, is_torchdynamo_compiling, logging +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling, logging +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_pop2piano import Pop2PianoConfig logger = logging.get_logger(__name__) +# 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: torch.Tensor | None, + scaling: float | None = 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: + 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.T5LayerNorm with T5->Pop2Piano class Pop2PianoLayerNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): @@ -142,9 +175,12 @@ def __init__( config: Pop2PianoConfig, has_relative_attention_bias=False, layer_idx: int | None = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal 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 @@ -153,6 +189,8 @@ def __init__( self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # Pop2Piano folds the relative position bias into the attention scores and does not scale the query/key dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -243,8 +281,7 @@ def forward( key_value_states=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): """ Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). @@ -291,43 +328,60 @@ def forward( 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] if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, query_states.shape[1], input_shape[1], key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, input_shape[1], key_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( - input_shape[1], key_length, device=scores.device, past_seen_tokens=past_seen_tokens + input_shape[1], key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + # `sdpa` may materialize a boolean mask (True = keep). Turn it into an additive float mask so it + # can be folded into the relative position bias, just like the `eager` float mask. + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) position_bias = position_bias + causal_mask - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) - - attn_output = torch.matmul(attn_weights, value_states) + # Pop2Piano uses a relative attention bias that is added to the attention scores. This is passed as the additive + # attention mask so that it works with the different attention implementations. As it is always non-`None`, + # `is_causal` is never inferred, so the causal behavior is fully encoded in the bias itself. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "Pop2Piano adds a relative position bias on top of the attention scores, which is only supported by the " + f"`eager` and `sdpa` attention implementations, but got `{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, + 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.reshape(*input_shape, -1) + attn_output = attn_output.reshape(*input_shape, -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->Pop2Piano,t5->pop2piano @@ -335,7 +389,10 @@ class Pop2PianoLayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): super().__init__() self.SelfAttention = Pop2PianoAttention( - 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 = Pop2PianoLayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -346,29 +403,27 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attention_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, + **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(attention_output) + return hidden_states, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->Pop2Piano,t5->pop2piano class Pop2PianoLayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: int | None = None): super().__init__() - self.EncDecAttention = Pop2PianoAttention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.EncDecAttention = Pop2PianoAttention( + config, has_relative_attention_bias=False, layer_idx=layer_idx, is_causal=False + ) self.layer_norm = Pop2PianoLayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -379,21 +434,19 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attention_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, - output_attentions=output_attentions, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + layer_output = hidden_states + self.dropout(attention_output) + return layer_output, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5Block with T5->Pop2Piano,t5->pop2piano @@ -421,21 +474,15 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - return_dict=True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - self_attention_outputs = self.layer[0]( + hidden_states, self_attn_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, + **kwargs, ) - hidden_states = self_attention_outputs[0] - attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -446,17 +493,17 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + cross_attn_position_bias = None do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: - cross_attention_outputs = self.layer[1]( + hidden_states, cross_attn_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, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = cross_attention_outputs[0] # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -467,9 +514,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) @@ -482,11 +526,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_attn_position_bias, cross_attn_position_bias @auto_docstring @@ -500,6 +540,17 @@ class Pop2PianoPreTrainedModel(PreTrainedModel): _no_split_modules = ["Pop2PianoBlock"] _keep_in_fp32_modules = ["wo"] + _supports_attention_backend = True + _supports_flash_attn = False + _supports_flex_attn = False + _supports_sdpa = True + + _can_record_outputs = { + "hidden_states": OutputRecorder(Pop2PianoBlock, index=0), + "attentions": OutputRecorder(Pop2PianoLayerSelfAttention, index=-1), + "cross_attentions": OutputRecorder(Pop2PianoLayerCrossAttention, index=-1), + } + @torch.no_grad() def _init_weights(self, module): """Initialize the weights""" @@ -586,6 +637,8 @@ def __init__(self, config): def set_input_embeddings(self, new_embeddings): self.embed_tokens = new_embeddings + @merge_with_config_defaults + @capture_outputs def forward( self, input_ids=None, @@ -595,17 +648,9 @@ def forward( inputs_embeds=None, past_key_values=None, use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=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.return_dict if input_ids is not None and inputs_embeds is not None: err_msg_prefix = "decoder_" if self.is_decoder else "" @@ -659,11 +704,16 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: + # T5 folds the relative position bias into the attention scores and always feeds it through the + # `attention_mask` argument, so `sdpa` can never rely on its `is_causal` shortcut. The dummy mask function + # is a no-op but forces the causal mask to always be materialized instead of being skipped. + dummy_and_mask_function = lambda *args: torch.tensor(True, dtype=torch.bool) # noqa: E731 causal_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + and_mask_function=dummy_and_mask_function, ) else: causal_mask = attention_mask[:, None, None, :] @@ -678,19 +728,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 i, layer_module in enumerate(self.block): - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_outputs = layer_module( + for layer_module in self.block: + hidden_states, self_attention_position_bias, cross_attention_position_bias = layer_module( hidden_states, causal_mask, position_bias, @@ -698,49 +742,20 @@ def forward( encoder_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, + **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 - the first layer stores them + 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, ) @@ -851,6 +866,7 @@ def get_mel_conditioner_outputs( return input_features, None + @can_return_tuple @auto_docstring def forward( self, @@ -865,11 +881,8 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqLMOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqLMOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Pop2Piano is a model with relative position embeddings @@ -892,7 +905,6 @@ def forward( labels in `[0, ..., config.vocab_size]` """ 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.return_dict if inputs_embeds is not None and input_features is not None: raise ValueError("Both `inputs_embeds` and `input_features` received! Please provide only one of them") @@ -906,11 +918,9 @@ 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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -932,9 +942,7 @@ 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, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -949,10 +957,6 @@ def forward( loss_fct = CrossEntropyLoss(ignore_index=-100) 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, diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index fca1a97a7efc..e008418e844e 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -20,6 +20,7 @@ import copy import math +from collections.abc import Callable import torch import torch.nn as nn @@ -37,7 +38,7 @@ Seq2SeqMoEModelOutput, Seq2SeqMoEOutput, ) -from ...modeling_utils import PreTrainedModel +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling, logging from ...utils.generic import can_return_tuple, merge_with_config_defaults @@ -222,15 +223,46 @@ def forward(self, hidden_states, **kwargs): return output +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float | None = 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: + 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 SwitchTransformersAttention(nn.Module): def __init__( self, config: SwitchTransformersConfig, has_relative_attention_bias=False, layer_idx: int | None = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal 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 @@ -239,6 +271,8 @@ def __init__( self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # SWITCH_TRANSFORMERS folds the relative position bias into the attention scores and does not scale the query/key dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -329,8 +363,7 @@ def forward( key_value_states=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): """ Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). @@ -377,50 +410,70 @@ def forward( 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] if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, query_states.shape[1], input_shape[1], key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, input_shape[1], key_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( - input_shape[1], key_length, device=scores.device, past_seen_tokens=past_seen_tokens + input_shape[1], key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + # `sdpa` may materialize a boolean mask (True = keep). Turn it into an additive float mask so it + # can be folded into the relative position bias, just like the `eager` float mask. + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) position_bias = position_bias + causal_mask - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) - - attn_output = torch.matmul(attn_weights, value_states) + # SWITCH_TRANSFORMERS uses a relative attention bias that is added to the attention scores. This is passed as the additive + # attention mask so that it works with the different attention implementations. As it is always non-`None`, + # `is_causal` is never inferred, so the causal behavior is fully encoded in the bias itself. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "SWITCH_TRANSFORMERS adds a relative position bias on top of the attention scores, which is only supported by the " + f"`eager` and `sdpa` attention implementations, but got `{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, + 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.reshape(*input_shape, -1) + attn_output = attn_output.reshape(*input_shape, -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 SwitchTransformersLayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): super().__init__() self.SelfAttention = SwitchTransformersAttention( - 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 = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -431,29 +484,25 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attention_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, + **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(attention_output) + return hidden_states, position_bias, attn_weights class SwitchTransformersLayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: int | None = None): super().__init__() self.EncDecAttention = SwitchTransformersAttention( - config, has_relative_attention_bias=False, layer_idx=layer_idx + config, has_relative_attention_bias=False, layer_idx=layer_idx, is_causal=False ) self.layer_norm = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -465,21 +514,19 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attention_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, - output_attentions=output_attentions, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + layer_output = hidden_states + self.dropout(attention_output) + return layer_output, position_bias, attn_weights class SwitchTransformersBlock(GradientCheckpointingLayer): @@ -507,43 +554,58 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - hidden_states, _ = self.layer[0]( + hidden_states, self_attn_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, + **kwargs, ) # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + cross_attn_position_bias = None do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: - hidden_states, _ = self.layer[1]( + hidden_states, cross_attn_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, + **kwargs, ) # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) hidden_states = self.layer[-1](hidden_states) + # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - return hidden_states + + return hidden_states, self_attn_position_bias, cross_attn_position_bias @auto_docstring @@ -718,8 +780,8 @@ def forward( hidden_states = self.dropout(inputs_embeds) - for i, layer_module in enumerate(self.block): - hidden_states = layer_module( + for layer_module in self.block: + hidden_states, self_attention_position_bias, cross_attention_position_bias = layer_module( hidden_states, causal_mask, position_bias, @@ -727,10 +789,14 @@ def forward( encoder_attention_mask, encoder_decoder_position_bias, past_key_values=past_key_values, - use_cache=use_cache, **kwargs, ) + # We share the position biases between the layers - the first layer stores them + position_bias = self_attention_position_bias + if self.is_decoder and encoder_hidden_states is not None: + encoder_decoder_position_bias = cross_attention_position_bias + hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.dropout(hidden_states) diff --git a/src/transformers/models/switch_transformers/modular_switch_transformers.py b/src/transformers/models/switch_transformers/modular_switch_transformers.py index cc1ea6be960e..78885e38ccb5 100644 --- a/src/transformers/models/switch_transformers/modular_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modular_switch_transformers.py @@ -287,43 +287,58 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - hidden_states, _ = self.layer[0]( + hidden_states, self_attn_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, + **kwargs, ) # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + cross_attn_position_bias = None do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: - hidden_states, _ = self.layer[1]( + hidden_states, cross_attn_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, + **kwargs, ) # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) hidden_states = self.layer[-1](hidden_states) + # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any(): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - return hidden_states + + return hidden_states, self_attn_position_bias, cross_attn_position_bias @auto_docstring @@ -498,8 +513,8 @@ def forward( hidden_states = self.dropout(inputs_embeds) - for i, layer_module in enumerate(self.block): - hidden_states = layer_module( + for layer_module in self.block: + hidden_states, self_attention_position_bias, cross_attention_position_bias = layer_module( hidden_states, causal_mask, position_bias, @@ -507,10 +522,14 @@ def forward( encoder_attention_mask, encoder_decoder_position_bias, past_key_values=past_key_values, - use_cache=use_cache, **kwargs, ) + # We share the position biases between the layers - the first layer stores them + position_bias = self_attention_position_bias + if self.is_decoder and encoder_hidden_states is not None: + encoder_decoder_position_bias = cross_attention_position_bias + hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.dropout(hidden_states) diff --git a/src/transformers/models/t5/modeling_t5.py b/src/transformers/models/t5/modeling_t5.py index c48fe1ef8bd3..8118ffe4d773 100644 --- a/src/transformers/models/t5/modeling_t5.py +++ b/src/transformers/models/t5/modeling_t5.py @@ -15,6 +15,7 @@ import copy import math +from collections.abc import Callable import torch from torch import nn @@ -35,8 +36,11 @@ Seq2SeqSequenceClassifierOutput, TokenClassifierOutput, ) -from ...modeling_utils import PreTrainedModel -from ...utils import DUMMY_INPUTS, DUMMY_MASK, auto_docstring, logging, torch_compilable_check +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import DUMMY_INPUTS, DUMMY_MASK, TransformersKwargs, auto_docstring, logging, torch_compilable_check +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_t5 import T5Config @@ -137,15 +141,47 @@ 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: torch.Tensor | None, + scaling: float | None = 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: + 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: int | None = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal 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 @@ -154,6 +190,8 @@ def __init__( self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # T5 folds the relative position bias into the attention scores and does not scale the query/key dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -244,8 +282,7 @@ def forward( key_value_states=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): """ Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). @@ -292,50 +329,70 @@ def forward( 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] if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, query_states.shape[1], input_shape[1], key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, input_shape[1], key_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( - input_shape[1], key_length, device=scores.device, past_seen_tokens=past_seen_tokens + input_shape[1], key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + # `sdpa` may materialize a boolean mask (True = keep). Turn it into an additive float mask so it + # can be folded into the relative position bias, just like the `eager` float mask. + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) position_bias = position_bias + causal_mask - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) - - attn_output = torch.matmul(attn_weights, value_states) + # T5 uses a relative attention bias that is added to the attention scores. This is passed as the additive + # attention mask so that it works with the different attention implementations. As it is always non-`None`, + # `is_causal` is never inferred, so the causal behavior is fully encoded in the bias itself. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "T5 adds a relative position bias on top of the attention scores, which is only supported by the " + f"`eager` and `sdpa` attention implementations, but got `{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, + 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.reshape(*input_shape, -1) + attn_output = attn_output.reshape(*input_shape, -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: int | None = 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) @@ -346,28 +403,26 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attention_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, + **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(attention_output) + return hidden_states, position_bias, attn_weights class T5LayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: int | None = 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) @@ -378,21 +433,19 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attention_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, - output_attentions=output_attentions, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + layer_output = hidden_states + self.dropout(attention_output) + return layer_output, position_bias, attn_weights class T5Block(GradientCheckpointingLayer): @@ -417,21 +470,15 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - return_dict=True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - self_attention_outputs = self.layer[0]( + hidden_states, self_attn_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, + **kwargs, ) - hidden_states = self_attention_outputs[0] - attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -442,17 +489,17 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + cross_attn_position_bias = None do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: - cross_attention_outputs = self.layer[1]( + hidden_states, cross_attn_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, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = cross_attention_outputs[0] # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -463,9 +510,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) @@ -478,11 +522,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_attn_position_bias, cross_attn_position_bias class T5ClassificationHead(nn.Module): @@ -513,6 +553,17 @@ class T5PreTrainedModel(PreTrainedModel): _no_split_modules = ["T5Block"] _keep_in_fp32_modules = ["wo"] + _supports_attention_backend = True + _supports_flash_attn = False + _supports_flex_attn = False + _supports_sdpa = True + + _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) @@ -621,6 +672,9 @@ def __init__(self, config): def set_input_embeddings(self, new_embeddings): self.embed_tokens = new_embeddings + @merge_with_config_defaults + @capture_outputs + @auto_docstring def forward( self, input_ids=None, @@ -630,17 +684,9 @@ def forward( inputs_embeds=None, past_key_values=None, use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=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.return_dict if input_ids is not None and inputs_embeds is not None: err_msg_prefix = "decoder_" if self.is_decoder else "" @@ -668,8 +714,6 @@ def forward( 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") @@ -688,6 +732,10 @@ def forward( past_key_values = None if self.config.is_decoder: + # T5 folds the relative position bias into the attention scores and always feeds it through the + # `attention_mask` argument, so `sdpa` can never rely on its `is_causal` shortcut. The dummy mask function + # is a no-op but forces the causal mask to always be materialized instead of being skipped. + dummy_and_mask_function = lambda *args: torch.tensor(True, dtype=torch.bool) # noqa: E731 attention_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, @@ -695,6 +743,7 @@ def forward( past_key_values=past_key_values.self_attention_cache if isinstance(past_key_values, EncoderDecoderCache) else past_key_values, + and_mask_function=dummy_and_mask_function, ) else: attention_mask = create_bidirectional_mask( @@ -712,19 +761,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, @@ -732,50 +775,20 @@ 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, + **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 - the first layer stores them + 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, ) @@ -814,6 +827,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, @@ -826,11 +840,8 @@ def forward( inputs_embeds: torch.Tensor | None = None, decoder_inputs_embeds: torch.Tensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -881,7 +892,6 @@ def forward( >>> 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.return_dict # Encode if needed (training, first prediction pass) if encoder_outputs is None: @@ -889,11 +899,9 @@ 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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -911,14 +919,9 @@ 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, + **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, @@ -975,6 +978,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, @@ -988,11 +992,8 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqLMOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqLMOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -1049,7 +1050,6 @@ def forward( >>> # 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.return_dict # Encode if needed (training, first prediction pass) if encoder_outputs is None: @@ -1058,11 +1058,9 @@ 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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -1084,9 +1082,7 @@ 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, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1103,10 +1099,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, @@ -1147,17 +1139,15 @@ 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: torch.LongTensor | None = None, attention_mask: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | BaseModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -1181,15 +1171,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.return_dict - - encoder_outputs = self.encoder( + encoder_outputs: BaseModelOutput = 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 @@ -1212,6 +1198,7 @@ def __init__(self, config: T5Config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def forward( self, @@ -1224,11 +1211,8 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple | Seq2SeqSequenceClassifierOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqSequenceClassifierOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -1260,7 +1244,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.return_dict if labels is not None: use_cache = False @@ -1280,7 +1263,7 @@ def forward( ) decoder_input_ids = self._shift_right(input_ids) - outputs = self.transformer( + outputs: Seq2SeqModelOutput = self.transformer( input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids, @@ -1289,9 +1272,7 @@ 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, + **kwargs, ) sequence_output = outputs[0] @@ -1333,9 +1314,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, @@ -1363,6 +1341,7 @@ def __init__(self, config: T5Config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def forward( self, @@ -1370,11 +1349,8 @@ def forward( attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, labels: torch.Tensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.Tensor] | TokenClassifierOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> TokenClassifierOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -1389,15 +1365,11 @@ 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.return_dict - - outputs = self.transformer( + outputs: BaseModelOutput = 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, + **kwargs, ) hidden_states = outputs[0] @@ -1409,10 +1381,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, @@ -1459,6 +1427,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, @@ -1472,11 +1441,8 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqQuestionAnsweringModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqQuestionAnsweringModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -1505,7 +1471,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.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 @@ -1522,20 +1487,15 @@ 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.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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -1553,9 +1513,7 @@ 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, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1582,10 +1540,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/udop/modeling_udop.py b/src/transformers/models/udop/modeling_udop.py index 758d3f2943fa..b5fc7f77179d 100644 --- a/src/transformers/models/udop/modeling_udop.py +++ b/src/transformers/models/udop/modeling_udop.py @@ -18,7 +18,7 @@ import math import random from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Callable, Sequence from copy import deepcopy from dataclasses import dataclass from typing import Any @@ -39,17 +39,50 @@ from ...generation import GenerationMixin from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_utils import PreTrainedModel +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...utils import ( ModelOutput, + TransformersKwargs, auto_docstring, is_torchdynamo_compiling, ) +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs logger = logging.getLogger(__name__) +# 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: torch.Tensor | None, + scaling: float | None = 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: + 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 + + @auto_docstring( custom_intro=""" Class for the model's outputs that may also contain a past key/values (to speed up sequential decoding). Includes @@ -252,6 +285,10 @@ class UdopPreTrainedModel(PreTrainedModel): _can_compile_fullgraph = False _keep_in_fp32_modules = ["wo"] + _supports_attention_backend = True + _supports_flash_attn = False + _supports_flex_attn = False + _supports_sdpa = True @torch.no_grad() def _init_weights(self, module): @@ -434,9 +471,12 @@ def __init__( config: UdopConfig, has_relative_attention_bias=False, layer_idx: int | None = None, + is_causal: bool = False, ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal 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 @@ -445,6 +485,8 @@ def __init__( self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # Udop folds the relative position bias into the attention scores and does not scale the query/key dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -535,8 +577,7 @@ def forward( key_value_states=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): """ Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). @@ -583,43 +624,60 @@ def forward( 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] if not self.has_relative_attention_bias: position_bias = torch.zeros( - (1, query_states.shape[1], input_shape[1], key_length), device=scores.device, dtype=scores.dtype + (1, self.n_heads, input_shape[1], key_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( - input_shape[1], key_length, device=scores.device, past_seen_tokens=past_seen_tokens + input_shape[1], key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + # `sdpa` may materialize a boolean mask (True = keep). Turn it into an additive float mask so it + # can be folded into the relative position bias, just like the `eager` float mask. + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) position_bias = position_bias + causal_mask - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) - - attn_output = torch.matmul(attn_weights, value_states) + # Udop uses a relative attention bias that is added to the attention scores. This is passed as the additive + # attention mask so that it works with the different attention implementations. As it is always non-`None`, + # `is_causal` is never inferred, so the causal behavior is fully encoded in the bias itself. + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "Udop adds a relative position bias on top of the attention scores, which is only supported by the " + f"`eager` and `sdpa` attention implementations, but got `{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, + 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.reshape(*input_shape, -1) + attn_output = attn_output.reshape(*input_shape, -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->Udop @@ -627,7 +685,10 @@ class UdopLayerSelfAttention(nn.Module): def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): super().__init__() self.SelfAttention = UdopAttention( - 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 = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -638,29 +699,27 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( + attention_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, + **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(attention_output) + return hidden_states, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->Udop class UdopLayerCrossAttention(nn.Module): def __init__(self, config, layer_idx: int | None = None): super().__init__() - self.EncDecAttention = UdopAttention(config, has_relative_attention_bias=False, layer_idx=layer_idx) + self.EncDecAttention = UdopAttention( + config, has_relative_attention_bias=False, layer_idx=layer_idx, is_causal=False + ) self.layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon) self.dropout = nn.Dropout(config.dropout_rate) @@ -671,21 +730,19 @@ def forward( attention_mask=None, position_bias=None, past_key_values=None, - output_attentions=False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( + attention_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, - output_attentions=output_attentions, + **kwargs, ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs + layer_output = hidden_states + self.dropout(attention_output) + return layer_output, position_bias, attn_weights # Copied from transformers.models.t5.modeling_t5.T5Block with T5->Udop @@ -713,21 +770,15 @@ def forward( encoder_attention_mask=None, encoder_decoder_position_bias=None, past_key_values=None, - use_cache=False, - output_attentions=False, - return_dict=True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): - self_attention_outputs = self.layer[0]( + hidden_states, self_attn_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, + **kwargs, ) - hidden_states = self_attention_outputs[0] - attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -738,17 +789,17 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + cross_attn_position_bias = None do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: - cross_attention_outputs = self.layer[1]( + hidden_states, cross_attn_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, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = cross_attention_outputs[0] # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -759,9 +810,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) @@ -774,11 +822,16 @@ def forward( ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - outputs = (hidden_states,) + return hidden_states, self_attn_position_bias, cross_attn_position_bias - return ( - outputs + attention_outputs - ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + +# `UdopBlock`, `UdopLayerSelfAttention` and `UdopLayerCrossAttention` are defined after `UdopPreTrainedModel`, so the +# output recorders are registered here (they are read lazily when a model instance is created). +UdopPreTrainedModel._can_record_outputs = { + "hidden_states": OutputRecorder(UdopBlock, index=0), + "attentions": OutputRecorder(UdopLayerSelfAttention, index=-1), + "cross_attentions": OutputRecorder(UdopLayerCrossAttention, index=-1), +} class UdopCellEmbeddings(nn.Module): @@ -792,6 +845,9 @@ def __init__(self, max_2d_position_embeddings=501, hidden_size=1024): def forward(self, bbox): bbox = torch.clip(bbox, 0.0, 1.0) bbox = (bbox * (self.max_2d_position_embeddings - 1)).long() + # In low precision (fp16/bf16) the multiplication above can round up to `max_2d_position_embeddings`, which is + # out of range for the embedding tables, so we clamp the indices to the valid range. + bbox = bbox.clamp(0, self.max_2d_position_embeddings - 1) left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0]) upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1]) right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2]) @@ -881,7 +937,9 @@ def get_bucket(self, attention_mask: Tensor | None = None, bbox: dict[str, Any] def get_relative_position(self, positions): context_position = positions[:, :, None] memory_position = positions[:, None, :] - relative_position = memory_position - context_position + # Upcast to float32: the positions are scaled by `scaling_factor`, which can overflow the range of fp16/bf16 + # for large bbox coordinates. float32 matches the default precision and keeps the bucketing exact. + relative_position = (memory_position - context_position).float() if self.augmentation and self.training: relative_position *= random.uniform(*AUGMENTATION_RANGE) relative_position *= self.scaling_factor @@ -1061,6 +1119,8 @@ def get_output_embeddings(self): def set_input_embeddings(self, new_embeddings): self.embed_tokens = new_embeddings + @merge_with_config_defaults + @capture_outputs def forward( self, input_ids=None, @@ -1075,17 +1135,9 @@ def forward( position_bias=None, past_key_values=None, use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - **kwargs, - ) -> tuple | BaseModelOutputWithAttentionMask: + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithAttentionMask: 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.return_dict # input embeddings processing @@ -1162,11 +1214,16 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: + # Udop folds the relative position bias into the attention scores and always feeds it through the + # `attention_mask` argument, so `sdpa` can never rely on its `is_causal` shortcut. The dummy mask function + # is a no-op but forces the causal mask to always be materialized instead of being skipped. + dummy_and_mask_function = lambda *args: torch.tensor(True, dtype=torch.bool) # noqa: E731 causal_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + and_mask_function=dummy_and_mask_function, ) else: causal_mask = attention_mask[:, None, None, :] @@ -1183,10 +1240,6 @@ def forward( else: encoder_extended_attention_mask = None - 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 - if self.is_decoder: # modified lines position_bias = None else: @@ -1198,11 +1251,8 @@ def forward( hidden_states = self.dropout(hidden_states) - for i, layer_module in enumerate(self.block): - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_outputs = layer_module( + for layer_module in self.block: + hidden_states, self_attention_position_bias, cross_attention_position_bias = layer_module( hidden_states, causal_mask, position_bias, @@ -1210,53 +1260,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, + **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 weights), - # (self-attention position bias), (cross-attention weights), (cross-attention position bias) - - position_bias = layer_outputs[1] + # We share the position biases between the layers - the first layer stores them + 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],) # We keep only self-attention weights for now - 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, - attention_mask, - past_key_values, - all_hidden_states, - all_attentions, - all_cross_attentions, - ] - if v is not None - ) - return BaseModelOutputWithAttentionMask( last_hidden_state=hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, - hidden_states=all_hidden_states, - attentions=all_attentions, - cross_attentions=all_cross_attentions, ) @@ -1297,6 +1315,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, @@ -1312,11 +1331,8 @@ def forward( past_key_values: Cache | None = None, decoder_inputs_embeds: Tensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple | Seq2SeqModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqModelOutput: r""" bbox (`torch.LongTensor` of shape `({0}, 4)`, *optional*): Bounding boxes of each input sequence tokens. Selected in the range `[0, @@ -1370,7 +1386,6 @@ def forward( [1, 1, 1024] ```""" 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.return_dict # Encode if needed (training, first prediction pass) if encoder_outputs is None: @@ -1381,13 +1396,11 @@ def forward( pixel_values=pixel_values, visual_bbox=visual_bbox, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + **kwargs, ) hidden_states = encoder_outputs[0] - encoder_attention_mask = encoder_outputs.attention_mask if return_dict else encoder_outputs[1] + encoder_attention_mask = encoder_outputs.attention_mask # Decode decoder_outputs = self.decoder( @@ -1398,17 +1411,9 @@ def forward( encoder_hidden_states=hidden_states, encoder_attention_mask=encoder_attention_mask, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + **kwargs, ) - if not return_dict: - # we filter out the attention mask - decoder_outputs = tuple(value for idx, value in enumerate(decoder_outputs) if idx != 1) - encoder_outputs = tuple(value for idx, value in enumerate(encoder_outputs) if idx != 1) - return decoder_outputs + encoder_outputs - return Seq2SeqModelOutput( last_hidden_state=decoder_outputs.last_hidden_state, past_key_values=decoder_outputs.past_key_values, @@ -1471,6 +1476,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, @@ -1486,12 +1492,9 @@ def forward( past_key_values: Cache | None = None, decoder_inputs_embeds: Tensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, labels: Tensor | None = None, - **kwargs, - ) -> tuple | Seq2SeqLMOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqLMOutput: r""" bbox (`torch.LongTensor` of shape `({0}, 4)`, *optional*): Bounding boxes of each input sequence tokens. Selected in the range `[0, @@ -1550,7 +1553,6 @@ def forward( ```""" 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.return_dict if decoder_input_ids is None and labels is not None: decoder_input_ids = self._shift_right(labels) @@ -1564,13 +1566,11 @@ def forward( pixel_values=pixel_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + **kwargs, ) hidden_states = encoder_outputs[0] - encoder_attention_mask = encoder_outputs.attention_mask if return_dict else encoder_outputs[1] + encoder_attention_mask = encoder_outputs.attention_mask # Decode decoder_outputs = self.decoder( @@ -1581,9 +1581,7 @@ def forward( encoder_hidden_states=hidden_states, encoder_attention_mask=encoder_attention_mask, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1598,10 +1596,6 @@ def forward( loss_fct = CrossEntropyLoss(ignore_index=-100) loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - if not return_dict: - output = (lm_logits,) + decoder_outputs[2:] + (encoder_outputs[0],) + encoder_outputs[2:] - return ((loss,) + output) if loss is not None else output - return Seq2SeqLMOutput( loss=loss, logits=lm_logits, @@ -1647,6 +1641,7 @@ 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, @@ -1656,11 +1651,8 @@ def forward( pixel_values: Tensor | None = None, visual_bbox: dict[str, Any] | None = None, inputs_embeds: Tensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | BaseModelOutputWithAttentionMask: + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithAttentionMask: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you @@ -1706,12 +1698,6 @@ def forward( >>> outputs = model(**encoding) >>> last_hidden_states = outputs.last_hidden_state ```""" - 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.return_dict - encoder_outputs = self.encoder( input_ids=input_ids, bbox=bbox, @@ -1719,9 +1705,7 @@ def forward( pixel_values=pixel_values, 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 diff --git a/src/transformers/models/umt5/modeling_umt5.py b/src/transformers/models/umt5/modeling_umt5.py index e1f2d647b52f..546c0eb1f3d4 100644 --- a/src/transformers/models/umt5/modeling_umt5.py +++ b/src/transformers/models/umt5/modeling_umt5.py @@ -15,6 +15,7 @@ import copy import math +from collections.abc import Callable import torch from torch import nn @@ -35,15 +36,18 @@ 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, - is_torchdynamo_compiling, logging, torch_compilable_check, ) +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_umt5 import UMT5Config @@ -148,14 +152,47 @@ 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: torch.Tensor | None, + scaling: float | None = 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: + 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 UMT5Attention(nn.Module): """ T5's attention using relative_attention_bias. """ - def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None): + def __init__( + self, config, has_relative_attention_bias=False, layer_idx: int | None = None, is_causal: bool = False + ): super().__init__() + self.config = config self.is_decoder = config.is_decoder + self.is_causal = is_causal 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 @@ -164,6 +201,8 @@ def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | N self.n_heads = config.num_heads self.dropout = config.dropout_rate self.inner_dim = self.n_heads * self.key_value_proj_dim + # UMT5 folds the relative position bias into the attention scores and does not scale the query/key dot product. + self.scaling = 1.0 self.layer_idx = layer_idx if layer_idx is None and self.is_decoder: logger.warning_once( @@ -295,33 +334,48 @@ def forward( 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)) - key_length = key_states.shape[-2] 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, seq_length, key_length), device=query_states.device, dtype=query_states.dtype ) else: position_bias = self.compute_bias( - seq_length, key_length, device=scores.device, past_seen_tokens=past_seen_tokens + seq_length, key_length, device=query_states.device, past_seen_tokens=past_seen_tokens ) if attention_mask is not None: - position_bias = position_bias + attention_mask - - position_bias_masked = position_bias - scores += position_bias_masked - - # (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) - - attn_output = torch.matmul(attn_weights, value_states) + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + if causal_mask.dtype == torch.bool: + causal_mask = torch.where( + causal_mask, + torch.tensor(0.0, device=causal_mask.device, dtype=position_bias.dtype), + torch.finfo(position_bias.dtype).min, + ) + position_bias = position_bias + causal_mask - attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.view(batch_size, seq_length, -1) + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation != "sdpa": + raise ValueError( + "UMT5 adds a relative position bias on top of the attention scores, which is only supported by " + "the `eager` and `sdpa` attention implementations, but got " + f"`{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, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o(attn_output) return attn_output, attn_weights @@ -345,6 +399,7 @@ def forward( normed_hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, + **kwargs, ) hidden_states = hidden_states + self.dropout(attention_output[0]) outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them @@ -372,6 +427,7 @@ def forward( encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, + **kwargs, ) layer_output = hidden_states + self.dropout(attention_output[0]) outputs = (layer_output,) + attention_output[1:] # add attentions if we output them @@ -396,14 +452,13 @@ def forward( encoder_hidden_states=None, encoder_attention_mask=None, past_key_values=None, - use_cache=False, - output_attentions=False, **kwargs, ): - hidden_states, self_attn_weights = self.layer[0]( + hidden_states, _ = self.layer[0]( hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, + **kwargs, ) # clamp inf values to enable fp16 training @@ -413,14 +468,14 @@ def forward( hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) # Cross-Attention Block - cross_attn_weights = None do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: - hidden_states, cross_attn_weights = self.layer[1]( + hidden_states, _ = self.layer[1]( hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=encoder_attention_mask, past_key_values=past_key_values, + **kwargs, ) # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: @@ -437,12 +492,7 @@ def forward( clamp_value = torch.where(torch.isinf(hidden_states).any(), max_dtype - 1000, max_dtype) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights, cross_attn_weights) - - return outputs + return hidden_states # Copied from transformers.models.t5.modeling_t5.T5ClassificationHead with T5->UMT5 @@ -473,6 +523,16 @@ class UMT5PreTrainedModel(PreTrainedModel): _can_compile_fullgraph = True _no_split_modules = ["UMT5Block"] _keep_in_fp32_modules = ["wo"] + _supports_attention_backend = True + _supports_flash_attn = False + _supports_flex_attn = False + _supports_sdpa = True + + _can_record_outputs = { + "hidden_states": OutputRecorder(UMT5Block, index=0), + "attentions": OutputRecorder(UMT5LayerSelfAttention, index=-1), + "cross_attentions": OutputRecorder(UMT5LayerCrossAttention, index=-1), + } @property def dummy_inputs(self): @@ -590,6 +650,9 @@ def __init__(self, config): def set_input_embeddings(self, new_embeddings): self.embed_tokens = new_embeddings + @merge_with_config_defaults + @capture_outputs + @auto_docstring def forward( self, input_ids=None, @@ -599,17 +662,9 @@ def forward( inputs_embeds=None, past_key_values=None, use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPastAndCrossAttentions: 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.return_dict if input_ids is not None and inputs_embeds is not None: err_msg_prefix = "decoder_" if self.is_decoder else "" @@ -637,8 +692,6 @@ def forward( 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") @@ -657,27 +710,28 @@ def forward( # it messes indexing later in decoder-stack because cache object is modified in-place past_key_values = None - past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - if attention_mask is None and not is_torchdynamo_compiling(): - # required mask seq length can be calculated via length of past cache - mask_seq_length = past_key_values_length + seq_length - attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) - if self.is_decoder: + # UMT5 folds the relative position bias into the attention scores and always routes the mask through the + # `attention_mask` argument, so `sdpa` must never take its `is_causal` shortcut. The dummy mask function is + # a no-op that keeps the causal mask materialized so it can be added to the position bias. + dummy_and_mask_function = lambda *args: torch.tensor(True, dtype=torch.bool) # noqa: E731 causal_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, - past_key_values=past_key_values, + past_key_values=past_key_values.self_attention_cache + if isinstance(past_key_values, EncoderDecoderCache) + else past_key_values, + and_mask_function=dummy_and_mask_function, ) - elif attention_mask is not None: - causal_mask = attention_mask[:, None, None, :] - causal_mask = causal_mask.to(dtype=inputs_embeds.dtype) - causal_mask = (1.0 - causal_mask) * torch.finfo(inputs_embeds.dtype).min else: - causal_mask = None + causal_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) - if self.is_decoder and encoder_attention_mask is not None: + if self.is_decoder and encoder_hidden_states is not None: encoder_extended_attention_mask = create_bidirectional_mask( config=self.config, inputs_embeds=inputs_embeds, @@ -687,58 +741,24 @@ def forward( else: encoder_extended_attention_mask = None - 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 - hidden_states = self.dropout(inputs_embeds) - for i, layer_module in enumerate(self.block): - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_outputs = layer_module( + for layer_module in self.block: + hidden_states = layer_module( hidden_states, causal_mask, encoder_hidden_states, # as a positional argument for gradient checkpointing encoder_attention_mask=encoder_extended_attention_mask, past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = layer_outputs[0] - - if output_attentions: - all_attentions += (layer_outputs[1],) - if self.is_decoder: - all_cross_attentions += (layer_outputs[2],) - 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, ) @@ -795,6 +815,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, @@ -807,11 +828,8 @@ def forward( inputs_embeds: torch.Tensor | None = None, decoder_inputs_embeds: torch.Tensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. UMT5 is a model with relative position embeddings so @@ -862,7 +880,6 @@ def forward( >>> 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.return_dict # Encode if needed (training, first prediction pass) if encoder_outputs is None: @@ -870,11 +887,9 @@ 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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -892,14 +907,9 @@ 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, + **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, @@ -972,6 +982,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, @@ -985,11 +996,8 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqLMOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqLMOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. UMT5 is a model with relative position embeddings so @@ -1043,7 +1051,6 @@ def forward( >>> tokenizer.decode(outputs[0], skip_special_tokens=True) ```""" 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.return_dict # Encode if needed (training, first prediction pass) if encoder_outputs is None: @@ -1052,11 +1059,9 @@ 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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -1078,9 +1083,7 @@ 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, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1099,10 +1102,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, @@ -1163,6 +1162,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 T5->UMT5, google-t5/t5-small->google/umt5-small, t5#training->umt5#training def forward( @@ -1170,11 +1170,8 @@ def forward( input_ids: torch.LongTensor | None = None, attention_mask: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | BaseModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. UMT5 is a model with relative position embeddings so you @@ -1198,15 +1195,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.return_dict - - encoder_outputs = self.encoder( + encoder_outputs: BaseModelOutput = 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 @@ -1230,6 +1223,7 @@ def __init__(self, config: UMT5Config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def forward( self, @@ -1242,11 +1236,8 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple | Seq2SeqSequenceClassifierOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqSequenceClassifierOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. UMT5 is a model with relative position embeddings so @@ -1278,7 +1269,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.return_dict if labels is not None: use_cache = False @@ -1307,9 +1297,7 @@ 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, + **kwargs, ) sequence_output = outputs[0] @@ -1351,9 +1339,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, @@ -1384,6 +1369,7 @@ def __init__(self, config: UMT5Config): # 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->UMT5, t5->umt5 def forward( @@ -1392,11 +1378,8 @@ def forward( attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, labels: torch.Tensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.Tensor] | TokenClassifierOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> TokenClassifierOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. UMT5 is a model with relative position embeddings so you @@ -1411,15 +1394,11 @@ 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.return_dict - - outputs = self.transformer( + outputs: BaseModelOutput = 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, + **kwargs, ) hidden_states = outputs[0] @@ -1431,10 +1410,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, @@ -1482,6 +1457,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, @@ -1495,11 +1471,8 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor] | Seq2SeqQuestionAnsweringModelOutput: + **kwargs: Unpack[TransformersKwargs], + ) -> Seq2SeqQuestionAnsweringModelOutput: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. UMT5 is a model with relative position embeddings so @@ -1528,7 +1501,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.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 @@ -1545,20 +1517,15 @@ 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.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, + **kwargs, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + elif 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, @@ -1576,9 +1543,7 @@ 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, + **kwargs, ) sequence_output = decoder_outputs[0] @@ -1605,10 +1570,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/udop/test_modeling_udop.py b/tests/models/udop/test_modeling_udop.py index 4d098822402a..9b9c24d1f2ea 100644 --- a/tests/models/udop/test_modeling_udop.py +++ b/tests/models/udop/test_modeling_udop.py @@ -280,6 +280,8 @@ class UdopModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin test_cpu_offload = False # The small UDOP model needs higher percentages for CPU/MP tests model_split_percents = [0.8, 0.9] + # UDOP requires `bbox` for its 2D relative position bias, so it must be forwarded by the generic tests + additional_model_inputs = ["bbox"] def setUp(self): self.model_tester = UdopModelTester(self) @@ -546,6 +548,8 @@ class UdopEncoderOnlyModelTest(ModelTesterMixin, unittest.TestCase): all_model_classes = (UdopEncoderModel,) if is_torch_available() else () test_resize_embeddings = False + # UDOP requires `bbox` for its 2D relative position bias, so it must be forwarded by the generic tests + additional_model_inputs = ["bbox"] def setUp(self): self.model_tester = UdopEncoderOnlyModelTester(self)