diff --git a/examples/modular-transformers/modeling_test_detr.py b/examples/modular-transformers/modeling_test_detr.py index f679fb14d95b..818788a9bb01 100644 --- a/examples/modular-transformers/modeling_test_detr.py +++ b/examples/modular-transformers/modeling_test_detr.py @@ -4,30 +4,93 @@ # the file from the modular. If any change should be done, please apply the change to the # modular_test_detr.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 - import math import warnings +from collections.abc import Callable from dataclasses import dataclass import torch +import torch.nn as nn import torch.nn.functional as F -from torch import Tensor, nn +from torch import Tensor from ... import initialization as init from ...activations import ACT2FN from ...backbone_utils import load_backbone from ...integrations import use_kernel_forward_from_hub -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput -from ...modeling_utils import PreTrainedModel -from ...pytorch_utils import meshgrid -from ...utils import ModelOutput, auto_docstring, is_timm_available, requires_backends, torch_compilable_check +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import compile_compatible_method_lru_cache, meshgrid +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, torch_compilable_check +from ...utils.generic import can_return_tuple, check_model_inputs +from ...utils.output_capturing import OutputRecorder from .configuration_test_detr import TestDetrConfig -if is_timm_available(): - from timm import create_model +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the TEST_DETR decoder. This class adds one attribute to BaseModelOutputWithCrossAttentions, + namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them + gone through a layernorm. This is useful when training the model with auxiliary decoding losses. + """ +) +class TestDetrDecoderOutput(BaseModelOutputWithCrossAttentions): + r""" + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax, + used to compute the weighted average in the cross-attention heads. + intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`): + Stacked intermediate reference points (reference points of each layer of the decoder). + """ + + intermediate_hidden_states: torch.FloatTensor | None = None + + intermediate_reference_points: torch.FloatTensor | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the Deformable DETR encoder-decoder model. + """ +) +class TestDetrModelOutput(ModelOutput): + r""" + init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): + Initial reference points sent through the Transformer decoder. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`): + Stacked intermediate hidden states (output of each layer of the decoder). + intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): + Stacked intermediate reference points (reference points of each layer of the decoder). + enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are + picked as region proposals in the first stage. Output of bounding box binary classification (i.e. + foreground and background). + enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): + Logits of predicted bounding boxes coordinates in the first stage. + """ + + init_reference_points: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + intermediate_hidden_states: torch.FloatTensor | None = None + intermediate_reference_points: torch.FloatTensor | None = None + decoder_hidden_states: tuple[torch.FloatTensor] | None = None + decoder_attentions: tuple[torch.FloatTensor] | None = None + cross_attentions: tuple[torch.FloatTensor] | None = None + encoder_last_hidden_state: torch.FloatTensor | None = None + encoder_hidden_states: tuple[torch.FloatTensor] | None = None + encoder_attentions: tuple[torch.FloatTensor] | None = None + enc_outputs_class: torch.FloatTensor | None = None + enc_outputs_coord_logits: torch.FloatTensor | None = None @use_kernel_forward_from_hub("MultiScaleDeformableAttention") @@ -85,73 +148,6 @@ def forward( return output.transpose(1, 2).contiguous() -@dataclass -@auto_docstring( - custom_intro=""" - Base class for outputs of the TestDetrDecoder. This class adds two attributes to - BaseModelOutputWithCrossAttentions, namely: - - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer) - - a stacked tensor of intermediate reference points. - """ -) -class TestDetrDecoderOutput(ModelOutput): - r""" - intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`): - Stacked intermediate hidden states (output of each layer of the decoder). - intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`): - Stacked intermediate reference points (reference points of each layer of the decoder). - cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): - Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, - sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax, - used to compute the weighted average in the cross-attention heads. - """ - - last_hidden_state: torch.FloatTensor | None = None - intermediate_hidden_states: torch.FloatTensor | None = None - intermediate_reference_points: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor] | None = None - attentions: tuple[torch.FloatTensor] | None = None - cross_attentions: tuple[torch.FloatTensor] | None = None - - -@dataclass -@auto_docstring( - custom_intro=""" - Base class for outputs of the Deformable DETR encoder-decoder model. - """ -) -class TestDetrModelOutput(ModelOutput): - r""" - init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): - Initial reference points sent through the Transformer decoder. - last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`): - Sequence of hidden-states at the output of the last layer of the decoder of the model. - intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`): - Stacked intermediate hidden states (output of each layer of the decoder). - intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): - Stacked intermediate reference points (reference points of each layer of the decoder). - enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): - Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are - picked as region proposals in the first stage. Output of bounding box binary classification (i.e. - foreground and background). - enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): - Logits of predicted bounding boxes coordinates in the first stage. - """ - - init_reference_points: torch.FloatTensor | None = None - last_hidden_state: torch.FloatTensor | None = None - intermediate_hidden_states: torch.FloatTensor | None = None - intermediate_reference_points: torch.FloatTensor | None = None - decoder_hidden_states: tuple[torch.FloatTensor] | None = None - decoder_attentions: tuple[torch.FloatTensor] | None = None - cross_attentions: tuple[torch.FloatTensor] | None = None - encoder_last_hidden_state: torch.FloatTensor | None = None - encoder_hidden_states: tuple[torch.FloatTensor] | None = None - encoder_attentions: tuple[torch.FloatTensor] | None = None - enc_outputs_class: torch.FloatTensor | None = None - enc_outputs_coord_logits: torch.FloatTensor | None = None - - class TestDetrFrozenBatchNorm2d(nn.Module): """ BatchNorm2d where the batch statistics and the affine parameters are fixed. @@ -228,47 +224,25 @@ def __init__(self, config): self.config = config - # For backwards compatibility we have to use the timm library directly instead of the AutoBackbone API - if config.use_timm_backbone: - # We default to values which were previously hard-coded. This enables configurability from the config - # using backbone arguments, while keeping the default behavior the same. - requires_backends(self, ["timm"]) - kwargs = getattr(config, "backbone_kwargs", {}) - kwargs = {} if kwargs is None else kwargs.copy() - out_indices = kwargs.pop("out_indices", (2, 3, 4) if config.num_feature_levels > 1 else (4,)) - num_channels = kwargs.pop("in_chans", config.num_channels) - if config.dilation: - kwargs["output_stride"] = kwargs.get("output_stride", 16) - backbone = create_model( - config.backbone, - pretrained=config.use_pretrained_backbone, - features_only=True, - out_indices=out_indices, - in_chans=num_channels, - **kwargs, - ) - else: - backbone = load_backbone(config) + backbone = load_backbone(config) + self.intermediate_channel_sizes = backbone.channels # replace batch norm by frozen batch norm with torch.no_grad(): replace_batch_norm(backbone) - self.model = backbone - self.intermediate_channel_sizes = ( - self.model.feature_info.channels() if config.use_timm_backbone else self.model.channels - ) - backbone_model_type = None - if config.backbone is not None: - backbone_model_type = config.backbone - elif config.backbone_config is not None: - backbone_model_type = config.backbone_config.model_type - else: - raise ValueError("Either `backbone` or `backbone_config` should be provided in the config") + # We used to load with timm library directly instead of the AutoBackbone API + # so we need to unwrap the `backbone._backbone` module to load weights without mismatch + is_timm_model = False + if hasattr(backbone, "_backbone"): + backbone = backbone._backbone + is_timm_model = True + self.model = backbone + backbone_model_type = config.backbone_config.model_type if "resnet" in backbone_model_type: for name, parameter in self.model.named_parameters(): - if config.use_timm_backbone: + if is_timm_model: if "layer2" not in name and "layer3" not in name and "layer4" not in name: parameter.requires_grad_(False) else: @@ -277,7 +251,9 @@ def __init__(self, config): def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor): # send pixel_values through the model to get list of feature maps - features = self.model(pixel_values) if self.config.use_timm_backbone else self.model(pixel_values).feature_maps + features = self.model(pixel_values) + if isinstance(features, dict): + features = features.feature_maps out = [] for feature_map in features: @@ -287,62 +263,55 @@ def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor): return out -class TestDetrConvModel(nn.Module): - """ - This module adds 2D position embeddings to all intermediate feature maps of the convolutional encoder. - """ - - def __init__(self, conv_encoder, position_embedding): - super().__init__() - self.conv_encoder = conv_encoder - self.position_embedding = position_embedding - - def forward(self, pixel_values, pixel_mask): - # send pixel_values and pixel_mask through backbone to get list of (feature_map, pixel_mask) tuples - out = self.conv_encoder(pixel_values, pixel_mask) - pos = [] - for feature_map, mask in out: - # position encoding - pos.append(self.position_embedding(feature_map, mask).to(feature_map.dtype)) - - return out, pos - - class TestDetrSinePositionEmbedding(nn.Module): """ This is a more standard version of the position embedding, very similar to the one used by the Attention is all you need paper, generalized to work on images. """ - def __init__(self, embedding_dim=64, temperature=10000, normalize=False, scale=None): + def __init__( + self, + num_position_features: int = 64, + temperature: int = 10000, + normalize: bool = False, + scale: float | None = None, + ): super().__init__() - self.embedding_dim = embedding_dim - self.temperature = temperature - self.normalize = normalize if scale is not None and normalize is False: raise ValueError("normalize should be True if scale is passed") - if scale is None: - scale = 2 * math.pi - self.scale = scale + self.num_position_features = num_position_features + self.temperature = temperature + self.normalize = normalize + self.scale = 2 * math.pi if scale is None else scale - def forward(self, pixel_values, pixel_mask): - if pixel_mask is None: - raise ValueError("No pixel mask provided") - y_embed = pixel_mask.cumsum(1, dtype=pixel_values.dtype) - x_embed = pixel_mask.cumsum(2, dtype=pixel_values.dtype) + @compile_compatible_method_lru_cache(maxsize=1) + def forward( + self, + shape: torch.Size, + device: torch.device | str, + dtype: torch.dtype, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if mask is None: + mask = torch.zeros((shape[0], shape[2], shape[3]), device=device, dtype=torch.bool) + y_embed = mask.cumsum(1, dtype=dtype) + x_embed = mask.cumsum(2, dtype=dtype) if self.normalize: eps = 1e-6 y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale - dim_t = torch.arange(self.embedding_dim, dtype=pixel_values.dtype, device=pixel_values.device) - dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.embedding_dim) + dim_t = torch.arange(self.num_position_features, dtype=torch.int64, device=device).to(dtype) + dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_position_features) pos_x = x_embed[:, :, :, None] / dim_t pos_y = y_embed[:, :, :, None] / dim_t pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + # Flatten spatial dimensions and permute to (batch_size, sequence_length, hidden_size) format + # expected by the encoder + pos = pos.flatten(2).permute(0, 2, 1) return pos @@ -356,19 +325,124 @@ def __init__(self, embedding_dim=256): self.row_embeddings = nn.Embedding(50, embedding_dim) self.column_embeddings = nn.Embedding(50, embedding_dim) - def forward(self, pixel_values, pixel_mask=None): - height, width = pixel_values.shape[-2:] - width_values = torch.arange(width, device=pixel_values.device) - height_values = torch.arange(height, device=pixel_values.device) + @compile_compatible_method_lru_cache(maxsize=1) + def forward( + self, + shape: torch.Size, + device: torch.device | str, + dtype: torch.dtype, + mask: torch.Tensor | None = None, + ): + height, width = shape[-2:] + width_values = torch.arange(width, device=device) + height_values = torch.arange(height, device=device) x_emb = self.column_embeddings(width_values) y_emb = self.row_embeddings(height_values) pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1) pos = pos.permute(2, 0, 1) pos = pos.unsqueeze(0) - pos = pos.repeat(pixel_values.shape[0], 1, 1, 1) + pos = pos.repeat(shape[0], 1, 1, 1) + # Flatten spatial dimensions and permute to (batch_size, sequence_length, hidden_size) format + # expected by the encoder + pos = pos.flatten(2).permute(0, 2, 1) return pos +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: + attention_mask = attention_mask[:, :, :, : key.shape[-2]] + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class TestDetrSelfAttention(nn.Module): + """ + Multi-headed self-attention from 'Attention Is All You Need' paper. + + In TEST_DETR, position embeddings are added to both queries and keys (but not values) in self-attention. + """ + + def __init__( + self, + config: TestDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float = 0.0, + bias: bool = True, + ): + super().__init__() + self.config = config + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + self.is_causal = False + + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Position embeddings are added to both queries and keys (but not values). + """ + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_key_input = hidden_states + position_embeddings if position_embeddings is not None else hidden_states + + query_states = self.q_proj(query_key_input).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(query_key_input).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + class TestDetrMultiscaleDeformableAttention(nn.Module): """ Multiscale deformable attention as proposed in Deformable DETR. @@ -406,9 +480,6 @@ def __init__(self, config: TestDetrConfig, num_heads: int, n_points: int): self.disable_custom_kernels = config.disable_custom_kernels - def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Tensor | None): - return tensor if position_embeddings is None else tensor + position_embeddings - def forward( self, hidden_states: torch.Tensor, @@ -420,11 +491,11 @@ def forward( spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, - output_attentions: bool = False, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: # add position embeddings to the hidden states before projecting to queries and keys if position_embeddings is not None: - hidden_states = self.with_pos_embed(hidden_states, position_embeddings) + hidden_states = hidden_states + position_embeddings batch_size, num_queries, _ = hidden_states.shape batch_size, sequence_length, _ = encoder_hidden_states.shape @@ -479,159 +550,48 @@ def forward( return output, attention_weights -class TestDetrMultiheadAttention(nn.Module): - """ - Multi-headed attention from 'Attention Is All You Need' paper. - - Here, we add position embeddings to the queries and keys (as explained in the Deformable DETR paper). - """ - - def __init__( - self, - embed_dim: int, - num_heads: int, - dropout: float = 0.0, - bias: bool = True, - ): +class TestDetrMLP(nn.Module): + def __init__(self, config: TestDetrConfig, hidden_size: int, intermediate_size: int): super().__init__() - self.embed_dim = embed_dim - self.num_heads = num_heads - self.dropout = dropout - self.head_dim = embed_dim // num_heads - if self.head_dim * num_heads != self.embed_dim: - raise ValueError( - f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" - f" {num_heads})." - ) - self.scaling = self.head_dim**-0.5 - - self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - - def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int): - return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() - - def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Tensor | None): - return tensor if position_embeddings is None else tensor + position_embeddings - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_embeddings: torch.Tensor | None = None, - output_attentions: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - """Input shape: Batch x Time x Channel""" - - batch_size, target_len, embed_dim = hidden_states.size() - # add position embeddings to the hidden states before projecting to queries and keys - if position_embeddings is not None: - hidden_states_original = hidden_states - hidden_states = self.with_pos_embed(hidden_states, position_embeddings) - - # get queries, keys and values - query_states = self.q_proj(hidden_states) * self.scaling - key_states = self._shape(self.k_proj(hidden_states), -1, batch_size) - value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size) - - proj_shape = (batch_size * self.num_heads, -1, self.head_dim) - query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape) - key_states = key_states.view(*proj_shape) - value_states = value_states.view(*proj_shape) - - source_len = key_states.size(1) - - attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) - - if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len): - raise ValueError( - f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is" - f" {attn_weights.size()}" - ) - - # expand attention_mask - if attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) - - if attention_mask is not None: - if attention_mask.size() != (batch_size, 1, target_len, source_len): - raise ValueError( - f"Attention mask should be of size {(batch_size, 1, target_len, source_len)}, but is" - f" {attention_mask.size()}" - ) - if attention_mask.dtype == torch.bool: - attention_mask = torch.zeros_like(attention_mask, dtype=attn_weights.dtype).masked_fill_( - attention_mask, -torch.inf - ) - attn_weights = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attention_mask - attn_weights = attn_weights.view(batch_size * self.num_heads, target_len, source_len) - - attn_weights = nn.functional.softmax(attn_weights, dim=-1) - - if output_attentions: - # this operation is a bit awkward, but it's required to - # make sure that attn_weights keeps its gradient. - # In order to do so, attn_weights have to reshaped - # twice and have to be reused in the following - attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len) - attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len) - else: - attn_weights_reshaped = None - - attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) - - attn_output = torch.bmm(attn_probs, value_states) - - if attn_output.size() != ( - batch_size * self.num_heads, - target_len, - self.head_dim, - ): - raise ValueError( - f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is" - f" {attn_output.size()}" - ) - - attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim) - attn_output = attn_output.transpose(1, 2) - attn_output = attn_output.reshape(batch_size, target_len, embed_dim) - - attn_output = self.out_proj(attn_output) + self.fc1 = nn.Linear(hidden_size, intermediate_size) + self.fc2 = nn.Linear(intermediate_size, hidden_size) + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + self.dropout = config.dropout - return attn_output, attn_weights_reshaped + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + return hidden_states class TestDetrEncoderLayer(GradientCheckpointingLayer): def __init__(self, config: TestDetrConfig): super().__init__() - self.embed_dim = config.d_model + self.hidden_size = config.d_model self.self_attn = TestDetrMultiscaleDeformableAttention( config, num_heads=config.encoder_attention_heads, n_points=config.encoder_n_points, ) - self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size) self.dropout = config.dropout - self.activation_fn = ACT2FN[config.activation_function] - self.activation_dropout = config.activation_dropout - self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) - self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) - self.final_layer_norm = nn.LayerNorm(self.embed_dim) + self.mlp = TestDetrMLP(config, self.hidden_size, config.encoder_ffn_dim) + self.final_layer_norm = nn.LayerNorm(self.hidden_size) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - position_embeddings: torch.Tensor | None = None, + spatial_position_embeddings: torch.Tensor | None = None, reference_points=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, - output_attentions: bool = False, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: """ Args: hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): @@ -646,24 +606,18 @@ def forward( Spatial shapes of the backbone feature maps. level_start_index (`torch.LongTensor`, *optional*): Level start index. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. """ residual = hidden_states - - # Apply Multi-scale Deformable Attention Module on the multi-scale feature maps. - hidden_states, attn_weights = self.self_attn( + hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, - position_embeddings=position_embeddings, + position_embeddings=spatial_position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, - output_attentions=output_attentions, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) @@ -671,12 +625,7 @@ def forward( hidden_states = self.self_attn_layer_norm(hidden_states) residual = hidden_states - hidden_states = self.activation_fn(self.fc1(hidden_states)) - hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) - - hidden_states = self.fc2(hidden_states) - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - + hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states hidden_states = self.final_layer_norm(hidden_states) @@ -685,54 +634,44 @@ def forward( clamp_value = torch.finfo(hidden_states.dtype).max - 1000 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - outputs = (hidden_states,) - - if output_attentions: - outputs += (attn_weights,) - - return outputs + return hidden_states class TestDetrDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: TestDetrConfig): super().__init__() - self.embed_dim = config.d_model + self.hidden_size = config.d_model - # self-attention - self.self_attn = TestDetrMultiheadAttention( - embed_dim=self.embed_dim, - num_heads=config.decoder_attention_heads, + self.self_attn = TestDetrSelfAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.decoder_attention_heads, dropout=config.attention_dropout, ) self.dropout = config.dropout - self.activation_fn = ACT2FN[config.activation_function] - self.activation_dropout = config.activation_dropout - self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) - # cross-attention + self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size) self.encoder_attn = TestDetrMultiscaleDeformableAttention( config, num_heads=config.decoder_attention_heads, n_points=config.decoder_n_points, ) - self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim) - # feedforward neural networks - self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim) - self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim) - self.final_layer_norm = nn.LayerNorm(self.embed_dim) + self.encoder_attn_layer_norm = nn.LayerNorm(self.hidden_size) + self.mlp = TestDetrMLP(config, self.hidden_size, config.decoder_ffn_dim) + self.final_layer_norm = nn.LayerNorm(self.hidden_size) def forward( self, hidden_states: torch.Tensor, - position_embeddings: torch.Tensor | None = None, + object_queries_position_embeddings: torch.Tensor | None = None, reference_points=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, - output_attentions: bool | None = False, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: """ Args: hidden_states (`torch.FloatTensor`): @@ -750,60 +689,47 @@ def forward( encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative values. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. """ residual = hidden_states # Self Attention - hidden_states, self_attn_weights = self.self_attn( + hidden_states, _ = self.self_attn( hidden_states=hidden_states, - position_embeddings=position_embeddings, - output_attentions=output_attentions, + position_embeddings=object_queries_position_embeddings, + **kwargs, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) - second_residual = hidden_states + residual = hidden_states # Cross-Attention - cross_attn_weights = None - hidden_states, cross_attn_weights = self.encoder_attn( + hidden_states, _ = self.encoder_attn( hidden_states=hidden_states, attention_mask=encoder_attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, - position_embeddings=position_embeddings, + position_embeddings=object_queries_position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, - output_attentions=output_attentions, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - hidden_states = second_residual + hidden_states + hidden_states = residual + hidden_states hidden_states = self.encoder_attn_layer_norm(hidden_states) # Fully Connected residual = hidden_states - hidden_states = self.activation_fn(self.fc1(hidden_states)) - hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) - hidden_states = self.fc2(hidden_states) - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states hidden_states = self.final_layer_norm(hidden_states) - outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights, cross_attn_weights) - - return outputs + return hidden_states @auto_docstring @@ -818,6 +744,13 @@ class TestDetrPreTrainedModel(PreTrainedModel): r"TestDetrEncoderLayer", r"TestDetrDecoderLayer", ] + _supports_sdpa = True + _supports_flash_attn = True + _supports_attention_backend = True + _supports_flex_attn = True + _keys_to_ignore_on_load_unexpected = [ + r"detr\.model\.backbone\.model\.layer\d+\.0\.downsample\.1\.num_batches_tracked" + ] @torch.no_grad() def _init_weights(self, module): @@ -875,9 +808,13 @@ class TestDetrEncoder(TestDetrPreTrainedModel): config: TestDetrConfig """ + _can_record_outputs = { + "hidden_states": TestDetrEncoderLayer, + "attentions": OutputRecorder(TestDetrMultiscaleDeformableAttention, layer_name="self_attn", index=1), + } + def __init__(self, config: TestDetrConfig): super().__init__(config) - self.gradient_checkpointing = False self.dropout = config.dropout self.layers = nn.ModuleList([TestDetrEncoderLayer(config) for _ in range(config.encoder_layers)]) @@ -885,51 +822,18 @@ def __init__(self, config: TestDetrConfig): # Initialize weights and apply final processing self.post_init() - @staticmethod - def get_reference_points(spatial_shapes_list, valid_ratios, device): - """ - Get reference points for each feature map. Used in decoder. - - Args: - spatial_shapes_list (`list[tuple[int, int]]`): - Spatial shapes of each feature map. - valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`): - Valid ratios of each feature map. - device (`torch.device`): - Device on which to create the tensors. - Returns: - `torch.FloatTensor` of shape `(batch_size, num_queries, num_feature_levels, 2)` - """ - reference_points_list = [] - for level, (height, width) in enumerate(spatial_shapes_list): - ref_y, ref_x = meshgrid( - torch.linspace(0.5, height - 0.5, height, dtype=valid_ratios.dtype, device=device), - torch.linspace(0.5, width - 0.5, width, dtype=valid_ratios.dtype, device=device), - indexing="ij", - ) - # TODO: valid_ratios could be useless here. check https://github.com/fundamentalvision/Deformable-DETR/issues/36 - ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, level, 1] * height) - ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, level, 0] * width) - ref = torch.stack((ref_x, ref_y), -1) - reference_points_list.append(ref) - reference_points = torch.cat(reference_points_list, 1) - reference_points = reference_points[:, :, None] * valid_ratios[:, None] - return reference_points - + @check_model_inputs() def forward( self, inputs_embeds=None, attention_mask=None, - position_embeddings=None, + spatial_position_embeddings=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, valid_ratios=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: r""" Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): @@ -939,64 +843,65 @@ def forward( - 1 for pixel features that are real (i.e. **not masked**), - 0 for pixel features that are padding (i.e. **masked**). [What are attention masks?](../glossary#attention-mask) - position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): - Position embeddings that are added to the queries and keys in each self-attention layer. + spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Spatial position embeddings (2D positional encodings) that are added to the queries and keys in each self-attention layer. spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`): Spatial shapes of each feature map. level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`): Starting index of each feature map. valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`): Ratio of valid area in each feature level. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - output_hidden_states (`bool`, *optional*): - Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors - for more detail. - return_dict (`bool`, *optional*): - Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - hidden_states = inputs_embeds hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - reference_points = self.get_reference_points(spatial_shapes_list, valid_ratios, device=inputs_embeds.device) - - encoder_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - for i, encoder_layer in enumerate(self.layers): - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - layer_outputs = encoder_layer( + + spatial_shapes_tuple = tuple(spatial_shapes_list) + reference_points = self.get_reference_points(spatial_shapes_tuple, valid_ratios, device=inputs_embeds.device) + + for encoder_layer in self.layers: + hidden_states = encoder_layer( hidden_states, attention_mask, - position_embeddings=position_embeddings, + spatial_position_embeddings=spatial_position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, - output_attentions=output_attentions, + **kwargs, ) - hidden_states = layer_outputs[0] + return BaseModelOutput(last_hidden_state=hidden_states) - if output_attentions: - all_attentions = all_attentions + (layer_outputs[1],) - - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) + @staticmethod + def get_reference_points(spatial_shapes_list, valid_ratios, device): + """ + Get reference points for each feature map. Used in decoder. - if not return_dict: - return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) - return BaseModelOutput( - last_hidden_state=hidden_states, - hidden_states=encoder_states, - attentions=all_attentions, - ) + Args: + spatial_shapes_list (`list[tuple[int, int]]`): + Spatial shapes of each feature map. + valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`): + Valid ratios of each feature map. + device (`torch.device`): + Device on which to create the tensors. + Returns: + `torch.FloatTensor` of shape `(batch_size, num_queries, num_feature_levels, 2)` + """ + reference_points_list = [] + for level, (height, width) in enumerate(spatial_shapes_list): + ref_y, ref_x = meshgrid( + torch.linspace(0.5, height - 0.5, height, dtype=valid_ratios.dtype, device=device), + torch.linspace(0.5, width - 0.5, width, dtype=valid_ratios.dtype, device=device), + indexing="ij", + ) + # TODO: valid_ratios could be useless here. check https://github.com/fundamentalvision/Deformable-DETR/issues/36 + ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, level, 1] * height) + ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, level, 0] * width) + ref = torch.stack((ref_x, ref_y), -1) + reference_points_list.append(ref) + reference_points = torch.cat(reference_points_list, 1) + reference_points = reference_points[:, :, None] * valid_ratios[:, None] + return reference_points def inverse_sigmoid(x, eps=1e-5): @@ -1021,12 +926,17 @@ class TestDetrDecoder(TestDetrPreTrainedModel): config: TestDetrConfig """ + _can_record_outputs = { + "hidden_states": TestDetrDecoderLayer, + "attentions": OutputRecorder(TestDetrSelfAttention, layer_name="self_attn", index=1), + "cross_attentions": OutputRecorder(TestDetrMultiscaleDeformableAttention, layer_name="encoder_attn", index=1), + } + def __init__(self, config: TestDetrConfig): super().__init__(config) self.dropout = config.dropout self.layers = nn.ModuleList([TestDetrDecoderLayer(config) for _ in range(config.decoder_layers)]) - self.gradient_checkpointing = False # hack implementation for iterative bounding box refinement and two-stage Deformable DETR self.bbox_embed = None @@ -1035,21 +945,19 @@ def __init__(self, config: TestDetrConfig): # Initialize weights and apply final processing self.post_init() + @check_model_inputs() def forward( self, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None, - position_embeddings=None, + object_queries_position_embeddings=None, reference_points=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, valid_ratios=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ): r""" Args: @@ -1063,8 +971,8 @@ def forward( in `[0, 1]`: - 1 for pixels that are real (i.e. **not masked**), - 0 for pixels that are padding (i.e. **masked**). - position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): - Position embeddings that are added to the queries and keys in each self-attention layer. + object_queries_position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Position embeddings for the object query slots that are added to the queries and keys in each self-attention layer. reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)` is `as_two_stage` else `(batch_size, num_queries, 2)` or , *optional*): Reference point in range `[0, 1]`, top-left (0,0), bottom-right (1, 1), including padding area. spatial_shapes (`torch.FloatTensor` of shape `(num_feature_levels, 2)`): @@ -1074,28 +982,11 @@ def forward( valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`, *optional*): Ratio of valid area in each feature level. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - output_hidden_states (`bool`, *optional*): - Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors - for more detail. - return_dict (`bool`, *optional*): - Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if inputs_embeds is not None: hidden_states = inputs_embeds # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None intermediate = () intermediate_reference_points = () @@ -1110,23 +1001,18 @@ def forward( else: raise ValueError("Reference points' last dimension must be of size 2") - if output_hidden_states: - all_hidden_states += (hidden_states,) - - layer_outputs = decoder_layer( + hidden_states = decoder_layer( hidden_states, - position_embeddings, + object_queries_position_embeddings, reference_points_input, spatial_shapes, spatial_shapes_list, level_start_index, encoder_hidden_states, # as a positional argument for gradient checkpointing encoder_attention_mask, - output_attentions, + **kwargs, ) - hidden_states = layer_outputs[0] - # hack implementation for iterative bounding box refinement if self.bbox_embed is not None: tmp = self.bbox_embed[idx](hidden_states) @@ -1147,56 +1033,17 @@ def forward( intermediate += (hidden_states,) intermediate_reference_points += (reference_points,) - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if encoder_hidden_states is not None: - all_cross_attentions += (layer_outputs[2],) - # Keep batch_size as first dimension intermediate = torch.stack(intermediate, dim=1) intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - intermediate, - intermediate_reference_points, - all_hidden_states, - all_self_attns, - all_cross_attentions, - ] - if v is not None - ) return TestDetrDecoderOutput( last_hidden_state=hidden_states, intermediate_hidden_states=intermediate, intermediate_reference_points=intermediate_reference_points, - hidden_states=all_hidden_states, - attentions=all_self_attns, - cross_attentions=all_cross_attentions, ) -def build_position_encoding(config): - n_steps = config.d_model // 2 - if config.position_embedding_type == "sine": - # TODO find a better way of exposing other arguments - position_embedding = TestDetrSinePositionEmbedding(n_steps, normalize=True) - elif config.position_embedding_type == "learned": - position_embedding = TestDetrLearnedPositionEmbedding(n_steps) - else: - raise ValueError(f"Not supported {config.position_embedding_type}") - - return position_embedding - - @auto_docstring( custom_intro=""" The bare Deformable DETR Model (consisting of a backbone and encoder-decoder Transformer) outputting raw @@ -1207,17 +1054,23 @@ class TestDetrModel(TestDetrPreTrainedModel): def __init__(self, config: TestDetrConfig): super().__init__(config) - # Create backbone + positional encoding - backbone = TestDetrConvEncoder(config) - position_embeddings = build_position_encoding(config) - self.backbone = TestDetrConvModel(backbone, position_embeddings) + # Create backbone + self.backbone = TestDetrConvEncoder(config) + + # Create positional encoding + if config.position_embedding_type == "sine": + self.position_embedding = TestDetrSinePositionEmbedding(config.d_model // 2, normalize=True) + elif config.position_embedding_type == "learned": + self.position_embedding = TestDetrLearnedPositionEmbedding(config.d_model // 2) + else: + raise ValueError(f"Not supported {config.position_embedding_type}") # Create input projection layers if config.num_feature_levels > 1: - num_backbone_outs = len(backbone.intermediate_channel_sizes) + num_backbone_outs = len(self.backbone.intermediate_channel_sizes) input_proj_list = [] for _ in range(num_backbone_outs): - in_channels = backbone.intermediate_channel_sizes[_] + in_channels = self.backbone.intermediate_channel_sizes[_] input_proj_list.append( nn.Sequential( nn.Conv2d(in_channels, config.d_model, kernel_size=1), @@ -1244,7 +1097,7 @@ def __init__(self, config: TestDetrConfig): [ nn.Sequential( nn.Conv2d( - backbone.intermediate_channel_sizes[-1], + self.backbone.intermediate_channel_sizes[-1], config.d_model, kernel_size=1, ), @@ -1272,11 +1125,11 @@ def __init__(self, config: TestDetrConfig): self.post_init() def freeze_backbone(self): - for name, param in self.backbone.conv_encoder.model.named_parameters(): + for name, param in self.backbone.model.named_parameters(): param.requires_grad_(False) def unfreeze_backbone(self): - for name, param in self.backbone.conv_encoder.model.named_parameters(): + for name, param in self.backbone.model.named_parameters(): param.requires_grad_(True) def get_valid_ratio(self, mask, dtype=torch.float32): @@ -1297,15 +1150,18 @@ def get_proposal_pos_embed(self, proposals): temperature = 10000 scale = 2 * math.pi - dim_t = torch.arange(num_pos_feats, dtype=proposals.dtype, device=proposals.device) + # Compute position embeddings in float32 to avoid overflow with large temperature values in fp16 + proposals_dtype = proposals.dtype + dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=proposals.device) dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats) # batch_size, num_queries, 4 - proposals = proposals.sigmoid() * scale + proposals = proposals.sigmoid().to(torch.float32) * scale # batch_size, num_queries, 4, 128 pos = proposals[:, :, :, None] / dim_t # batch_size, num_queries, 4, 64, 2 -> batch_size, num_queries, 512 pos = torch.stack((pos[:, :, :, 0::2].sin(), pos[:, :, :, 1::2].cos()), dim=4).flatten(2) - return pos + # Convert back to target dtype after all computations are done + return pos.to(proposals_dtype) def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes): """Generate the encoder output proposals from encoded enc_output. @@ -1369,6 +1225,7 @@ def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes) return object_query, output_proposals @auto_docstring + @can_return_tuple def forward( self, pixel_values: torch.FloatTensor, @@ -1377,10 +1234,7 @@ def forward( encoder_outputs: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor] | TestDetrModelOutput: r""" decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): @@ -1397,12 +1251,10 @@ def forward( ```python >>> from transformers import AutoImageProcessor, TestDetrModel >>> from PIL import Image - >>> import httpx - >>> from io import BytesIO + >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" - >>> with httpx.stream("GET", url) as response: - ... image = Image.open(BytesIO(response.read())) + >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("SenseTime/deformable-detr") >>> model = TestDetrModel.from_pretrained("SenseTime/deformable-detr") @@ -1415,12 +1267,6 @@ def forward( >>> list(last_hidden_states.shape) [1, 300, 256] ```""" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - batch_size, num_channels, height, width = pixel_values.shape device = pixel_values.device @@ -1430,16 +1276,22 @@ def forward( # Extract multi-scale feature maps of same resolution `config.d_model` (cf Figure 4 in paper) # First, sent pixel_values + pixel_mask through Backbone to obtain the features # which is a list of tuples - features, position_embeddings_list = self.backbone(pixel_values, pixel_mask) + features = self.backbone(pixel_values, pixel_mask) # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default) sources = [] masks = [] + position_embeddings_list = [] for level, (source, mask) in enumerate(features): sources.append(self.input_proj[level](source)) masks.append(mask) if mask is None: raise ValueError("No attention mask was provided") + # Generate position embeddings for this feature level + pos = self.position_embedding(shape=source.shape, device=device, dtype=pixel_values.dtype, mask=mask).to( + source.dtype + ) + position_embeddings_list.append(pos) # Lowest resolution feature maps are obtained via 3x3 stride 2 convolutions on the final stage if self.config.num_feature_levels > len(sources): @@ -1452,7 +1304,9 @@ def forward( mask = nn.functional.interpolate(pixel_mask[None].to(pixel_values.dtype), size=source.shape[-2:]).to( torch.bool )[0] - pos_l = self.backbone.position_embedding(source, mask).to(source.dtype) + pos_l = self.position_embedding( + shape=source.shape, device=device, dtype=pixel_values.dtype, mask=mask + ).to(source.dtype) sources.append(source) masks.append(mask) position_embeddings_list.append(pos_l) @@ -1473,7 +1327,6 @@ def forward( spatial_shapes_list.append(spatial_shape) source = source.flatten(2).transpose(1, 2) mask = mask.flatten(1) - pos_embed = pos_embed.flatten(2).transpose(1, 2) lvl_pos_embed = pos_embed + self.level_embed[level].view(1, 1, -1) lvl_pos_embed_flatten.append(lvl_pos_embed) source_flatten.append(source) @@ -1491,21 +1344,12 @@ def forward( encoder_outputs = self.encoder( inputs_embeds=source_flatten, attention_mask=mask_flatten, - position_embeddings=lvl_pos_embed_flatten, + spatial_position_embeddings=lvl_pos_embed_flatten, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, valid_ratios=valid_ratios, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + **kwargs, ) # Fifth, prepare decoder inputs @@ -1548,7 +1392,7 @@ def forward( decoder_outputs = self.decoder( inputs_embeds=target, - position_embeddings=query_embed, + object_queries_position_embeddings=query_embed, encoder_hidden_states=encoder_outputs[0], encoder_attention_mask=mask_flatten, reference_points=reference_points, @@ -1556,17 +1400,9 @@ def forward( spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, valid_ratios=valid_ratios, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + **kwargs, ) - if not return_dict: - enc_outputs = tuple(value for value in [enc_outputs_class, enc_outputs_coord_logits] if value is not None) - tuple_outputs = (init_reference_points,) + decoder_outputs + encoder_outputs + enc_outputs - - return tuple_outputs - return TestDetrModelOutput( init_reference_points=init_reference_points, last_hidden_state=decoder_outputs.last_hidden_state, diff --git a/src/transformers/modeling_attn_mask_utils.py b/src/transformers/modeling_attn_mask_utils.py old mode 100755 new mode 100644 index 8eeb3eccd183..08288876b38f --- a/src/transformers/modeling_attn_mask_utils.py +++ b/src/transformers/modeling_attn_mask_utils.py @@ -22,9 +22,19 @@ import torch +from .utils import logging from .utils.import_utils import is_torchdynamo_compiling, is_tracing +logger = logging.get_logger(__name__) + + +DEPRECATION_MESSAGE = ( + "The attention mask API under `transformers.modeling_attn_mask_utils` (`AttentionMaskConverter`) " + "is deprecated and will be removed in Transformers v5.10. Please use the new API in `transformers.masking_utils`." +) + + @dataclass class AttentionMaskConverter: """ @@ -61,6 +71,8 @@ class AttentionMaskConverter: sliding_window: int def __init__(self, is_causal: bool, sliding_window: int | None = None): + logger.warning_once(DEPRECATION_MESSAGE, FutureWarning) + self.is_causal = is_causal self.sliding_window = sliding_window @@ -160,6 +172,8 @@ def _make_causal_mask( """ Make causal mask used for bi-directional self-attention. """ + logger.warning_once(DEPRECATION_MESSAGE, FutureWarning) + bsz, tgt_len = input_ids_shape mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) mask_cond = torch.arange(mask.size(-1), device=device) @@ -188,6 +202,8 @@ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = N """ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. """ + logger.warning_once(DEPRECATION_MESSAGE, FutureWarning) + bsz, src_len = mask.size() tgt_len = tgt_len if tgt_len is not None else src_len @@ -238,6 +254,8 @@ def _unmask_unattended( [0, 1, 1]]]] ``` """ + logger.warning_once(DEPRECATION_MESSAGE, FutureWarning) + # fmt: on if expanded_mask.dtype == torch.bool: raise ValueError( @@ -263,6 +281,7 @@ def _ignore_causal_mask_sdpa( allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed). """ + logger.warning_once(DEPRECATION_MESSAGE, FutureWarning) _, query_length = inputs_embeds.shape[0], inputs_embeds.shape[1] key_value_length = query_length + past_key_values_length @@ -445,6 +464,8 @@ def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len (`int`): The target length or query length the created mask shall have. """ + logger.warning_once(DEPRECATION_MESSAGE, FutureWarning) + _, key_value_length = mask.shape tgt_len = tgt_len if tgt_len is not None else key_value_length diff --git a/src/transformers/models/align/modeling_align.py b/src/transformers/models/align/modeling_align.py index 91f96d288c70..5793ca6239bd 100644 --- a/src/transformers/models/align/modeling_align.py +++ b/src/transformers/models/align/modeling_align.py @@ -619,7 +619,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.attention_head_size) @@ -674,7 +674,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_outputs = self.self( hidden_states, @@ -732,7 +732,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_attention_outputs = self.attention( hidden_states, @@ -771,7 +771,7 @@ def forward( output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None @@ -897,7 +897,7 @@ def forward( output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" Examples: @@ -1010,7 +1010,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPoolingAndNoAttention: r""" Examples: @@ -1173,7 +1173,7 @@ def forward( output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | AlignOutput: r""" return_loss (`bool`, *optional*): diff --git a/src/transformers/models/altclip/modeling_altclip.py b/src/transformers/models/altclip/modeling_altclip.py index d6bc5965779b..769edaec72ae 100755 --- a/src/transformers/models/altclip/modeling_altclip.py +++ b/src/transformers/models/altclip/modeling_altclip.py @@ -34,7 +34,6 @@ from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int -from ...utils.generic import is_flash_attention_requested from .configuration_altclip import AltCLIPConfig, AltCLIPTextConfig, AltCLIPVisionConfig @@ -344,7 +343,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_attention_outputs = self.attention( hidden_states, @@ -384,7 +383,7 @@ def forward( output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None @@ -481,8 +480,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: """Input shape: Batch x Time x Channel""" @@ -495,15 +494,6 @@ def forward( queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - # CLIP text model uses both `causal_attention_mask` and `attention_mask` - # in case FA2 kernel is called, `is_causal` should be inferred from `causal_attention_mask` - if not is_flash_attention_requested(self.config): - if attention_mask is not None and causal_attention_mask is not None: - attention_mask = attention_mask + causal_attention_mask - elif causal_attention_mask is not None: - attention_mask = causal_attention_mask - else: - self.is_causal = causal_attention_mask is not None attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward @@ -515,9 +505,9 @@ def forward( keys, values, attention_mask, - is_causal=self.is_causal, scaling=self.scale, dropout=0.0 if not self.training else self.dropout, + **kwargs, ) attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() @@ -556,8 +546,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -575,8 +565,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -613,10 +603,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutput: r""" Args: @@ -630,13 +620,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -663,8 +646,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] @@ -896,7 +879,7 @@ def forward( output_hidden_states: bool | None = None, interpolate_pos_encoding: bool = False, return_dict: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" Examples: @@ -1070,7 +1053,7 @@ def forward( output_attentions: bool | None = None, return_dict: bool | None = None, output_hidden_states: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPoolingAndProjection: r""" Examples: @@ -1244,7 +1227,7 @@ def forward( output_hidden_states: bool | None = None, interpolate_pos_encoding: bool = False, return_dict: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | AltCLIPOutput: r""" return_loss (`bool`, *optional*): diff --git a/src/transformers/models/autoformer/modeling_autoformer.py b/src/transformers/models/autoformer/modeling_autoformer.py index e0c5e175068b..ef21969ea905 100644 --- a/src/transformers/models/autoformer/modeling_autoformer.py +++ b/src/transformers/models/autoformer/modeling_autoformer.py @@ -25,23 +25,15 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache -from ...modeling_attn_mask_utils import ( - _prepare_4d_attention_mask, - _prepare_4d_attention_mask_for_sdpa, -) +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, ModelOutput, SampleTSPredictionOutput, Seq2SeqTSPredictionOutput from ...modeling_utils import PreTrainedModel from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput -from ...utils import auto_docstring, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested +from ...utils import auto_docstring, logging from .configuration_autoformer import AutoformerConfig -if is_torch_flex_attn_available(): - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -844,27 +836,6 @@ def _init_weights(self, module: nn.Module): init.ones_(module.weight) init.zeros_(module.bias) - # copied from transformers.models.bart.modeling_bart.BartPreTrainedModel._update_full_mask - def _update_full_mask( - self, - attention_mask: torch.Tensor | None, - inputs_embeds: torch.Tensor, - ): - if attention_mask is not None: - if is_flash_attention_requested(self.config): - attention_mask = attention_mask if 0 in attention_mask else None - elif self.config._attn_implementation == "sdpa": - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype) - elif self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask, is_causal=False) - else: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) - - return attention_mask - # copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesTransformerEncoder with TimeSeriesTransformer->Autoformer,TimeSeries->Autoformer class AutoformerEncoder(AutoformerPreTrainedModel): @@ -938,9 +909,10 @@ def forward( hidden_states = self.layernorm_embedding(hidden_states + embed_pos) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - attention_mask = self._update_full_mask( - attention_mask, - inputs_embeds, + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, ) encoder_states = () if output_hidden_states else None @@ -1086,8 +1058,6 @@ def forward( ) use_cache = False - input_shape = inputs_embeds.size()[:-1] - if self.gradient_checkpointing and use_cache: logger.warning( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." @@ -1097,12 +1067,12 @@ def forward( if use_cache and past_key_values is None: past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) - # expand encoder attention mask - if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] - ) + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) hidden_states = self.value_embedding(inputs_embeds) embed_pos = self.embed_positions( diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index 1e3cd256fa1b..424ee0b721e6 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -34,15 +34,16 @@ from ... import initialization as init from ...cache_utils import Cache from ...generation import GenerationMixin -from ...integrations import lazy_load_kernel, use_kernel_forward_from_hub, use_kernelized_func -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations.hub_kernels import lazy_load_kernel +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils.generic import maybe_autocast from .configuration_bamba import BambaConfig @@ -1197,8 +1198,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) mamba_mask = self._update_mamba_mask(attention_mask, cache_position) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1250,122 +1256,6 @@ def forward( attentions=all_self_attns, ) - def _update_causal_mask( - self, - attention_mask: torch.Tensor, - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: HybridMambaAttentionDynamicCache, - output_attentions: bool, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_attention_mask = (attention_mask[:, None, None, :] == attention_mask[:, None, :, None])[ - :, :, -sequence_length:, : - ].to(dtype) - padding_mask = causal_mask[:, :, :, :mask_length] + padding_attention_mask - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - def _update_mamba_mask(self, attention_mask, cache_position): """ No need for zeroing states when diff --git a/src/transformers/models/bamba/modular_bamba.py b/src/transformers/models/bamba/modular_bamba.py index 59ef3482d3be..e396a135d1b3 100644 --- a/src/transformers/models/bamba/modular_bamba.py +++ b/src/transformers/models/bamba/modular_bamba.py @@ -42,13 +42,12 @@ ) from ... import initialization as init -from ...integrations import lazy_load_kernel -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...integrations.hub_kernels import lazy_load_kernel +from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging -from ...utils.generic import is_flash_attention_requested from .configuration_bamba import BambaConfig @@ -868,8 +867,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) mamba_mask = self._update_mamba_mask(attention_mask, cache_position) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -921,122 +925,6 @@ def forward( attentions=all_self_attns, ) - def _update_causal_mask( - self, - attention_mask: torch.Tensor, - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: HybridMambaAttentionDynamicCache, - output_attentions: bool, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_attention_mask = (attention_mask[:, None, None, :] == attention_mask[:, None, :, None])[ - :, :, -sequence_length:, : - ].to(dtype) - padding_mask = causal_mask[:, :, :, :mask_length] + padding_attention_mask - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - def _update_mamba_mask(self, attention_mask, cache_position): """ No need for zeroing states when diff --git a/src/transformers/models/bark/modeling_bark.py b/src/transformers/models/bark/modeling_bark.py index 684d5524444c..ba2f7832b9ea 100644 --- a/src/transformers/models/bark/modeling_bark.py +++ b/src/transformers/models/bark/modeling_bark.py @@ -28,7 +28,7 @@ BarkEosPrioritizerLogitsProcessor, SuppressTokensLogitsProcessor, ) -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import CausalLMOutputWithPast, MaskedLMOutput @@ -39,7 +39,6 @@ is_torch_accelerator_available, logging, ) -from ...utils.generic import is_flash_attention_requested from ..auto import AutoModel from .configuration_bark import ( BarkCoarseConfig, @@ -471,7 +470,6 @@ def forward( raise ValueError("You have to specify either input_ids or input_embeds") input_shape = input_embeds.size()[:-1] - batch_size = input_embeds.shape[0] seq_length = input_shape[-1] if self.gradient_checkpointing and self.training: @@ -499,18 +497,11 @@ def forward( position_ids = position_ids.to(self.position_embeds_layer.weight.device) position_embeds = self.position_embeds_layer(position_ids) # position embeddings of shape (1, t, n_embd) - # Attention mask. - if attention_mask is not None: - if batch_size <= 0: - raise ValueError("batch_size has to be defined and > 0") - if is_flash_attention_requested(self.config): - attention_mask = attention_mask if 0 in attention_mask else None - else: - attention_mask = attention_mask.view(batch_size, -1) - # [bsz, to_seq_length] -> [bsz, 1, 1, to_seq_length] - # from_seq_length is 1 to easily broadcast - attention_mask = _prepare_4d_attention_mask(attention_mask, input_embeds.dtype, tgt_len=1) - attention_mask = attention_mask.to(input_embeds.device) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=input_embeds, + attention_mask=attention_mask, + ) hidden_states = self.drop(input_embeds + position_embeds) output_shape = input_shape + (hidden_states.size(-1),) @@ -1090,7 +1081,6 @@ def forward( input_embeds = input_embeds[:, :, :, : codebook_idx + 1].sum(dim=-1) input_shape = input_embeds.size()[:-1] - batch_size = input_embeds.shape[0] seq_length = input_shape[1] input_embeds = input_embeds.to(self.position_embeds_layer.weight.device) @@ -1104,17 +1094,11 @@ def forward( position_ids = position_ids.to(self.position_embeds_layer.weight.device) position_embeds = self.position_embeds_layer(position_ids) # position embeddings of shape (1, t, n_embd) - # Attention mask. - if attention_mask is not None: - if batch_size <= 0: - raise ValueError("batch_size has to be defined and > 0") - if is_flash_attention_requested(self.config): - attention_mask = attention_mask if 0 in attention_mask else None - else: - # [bsz, to_seq_length] -> [bsz, 1, 1, to_seq_length] - # from_seq_length is 1 to easily broadcast - attention_mask = _prepare_4d_attention_mask(attention_mask, input_embeds.dtype, tgt_len=1) - attention_mask = attention_mask.to(input_embeds.device) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=input_embeds, + attention_mask=attention_mask, + ) hidden_states = self.drop(input_embeds + position_embeds) output_shape = input_shape + (hidden_states.size(-1),) diff --git a/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py b/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py index 8ec1a5ca230b..fed5b55149ba 100755 --- a/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py +++ b/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py @@ -26,7 +26,6 @@ from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin from ...masking_utils import create_bidirectional_mask, create_causal_mask -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -1690,7 +1689,11 @@ def forward( # expand attention_mask if self.attention_type == "original_full": # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) blocked_encoder_mask = band_mask = from_mask = to_mask = None elif self.attention_type == "block_sparse": blocked_encoder_mask, band_mask, from_mask, to_mask = self.create_masks_for_block_sparse_attn( diff --git a/src/transformers/models/bloom/modeling_bloom.py b/src/transformers/models/bloom/modeling_bloom.py index bc2ec248438d..5efb8b7e8eb7 100644 --- a/src/transformers/models/bloom/modeling_bloom.py +++ b/src/transformers/models/bloom/modeling_bloom.py @@ -14,7 +14,6 @@ """PyTorch BLOOM model.""" import math -from typing import Union import torch from torch import nn @@ -23,7 +22,7 @@ from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutputWithPastAndCrossAttentions, @@ -35,19 +34,11 @@ from ...modeling_utils import PreTrainedModel from ...utils import ( auto_docstring, - is_torch_flex_attn_available, logging, ) -from ...utils.generic import is_flash_attention_requested from .configuration_bloom import BloomConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -520,8 +511,12 @@ def forward( attention_mask = attention_mask.to(hidden_states.device) alibi = self.build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) for i, block in enumerate(self.h): @@ -560,131 +555,6 @@ def forward( attentions=all_self_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py index 495c936d5161..8ef0da7a5072 100644 --- a/src/transformers/models/blt/modular_blt.py +++ b/src/transformers/models/blt/modular_blt.py @@ -525,12 +525,6 @@ def _init_weights(self, module): init.copy_(module.inv_freq, buffer_value) init.copy_(module.original_inv_freq, buffer_value) - def _update_causal_mask(self, module): - raise AttributeError("No need to inherit it!") - - def _prepare_4d_causal_attention_mask_with_cache_position(self, module): - raise AttributeError("No need to inherit it!") - class BltLocalEncoder(BltPreTrainedModel): config: BltLocalEncoderConfig diff --git a/src/transformers/models/chinese_clip/modeling_chinese_clip.py b/src/transformers/models/chinese_clip/modeling_chinese_clip.py index 10425aa5dca8..ab19682bad3e 100644 --- a/src/transformers/models/chinese_clip/modeling_chinese_clip.py +++ b/src/transformers/models/chinese_clip/modeling_chinese_clip.py @@ -287,7 +287,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.attention_head_size) @@ -343,7 +343,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_outputs = self.self( hidden_states, @@ -473,7 +473,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_attention_outputs = self.attention( hidden_states, @@ -631,7 +631,7 @@ def forward( output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None diff --git a/src/transformers/models/clap/modeling_clap.py b/src/transformers/models/clap/modeling_clap.py index b462a1793421..4298370c8b5a 100644 --- a/src/transformers/models/clap/modeling_clap.py +++ b/src/transformers/models/clap/modeling_clap.py @@ -1104,7 +1104,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.attention_head_size) @@ -1160,7 +1160,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_outputs = self.self( hidden_states, @@ -1219,7 +1219,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_attention_outputs = self.attention( hidden_states, @@ -1259,7 +1259,7 @@ def forward( output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None diff --git a/src/transformers/models/clipseg/modeling_clipseg.py b/src/transformers/models/clipseg/modeling_clipseg.py index ed2b5cdb5b68..d041180f9a07 100644 --- a/src/transformers/models/clipseg/modeling_clipseg.py +++ b/src/transformers/models/clipseg/modeling_clipseg.py @@ -24,13 +24,12 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int -from ...utils.generic import is_flash_attention_requested from .configuration_clipseg import CLIPSegConfig, CLIPSegTextConfig, CLIPSegVisionConfig @@ -308,8 +307,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = False, + **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Input shape: Batch x Time x Channel""" @@ -322,15 +321,6 @@ def forward( queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - # CLIP text model uses both `causal_attention_mask` and `attention_mask` - # in case FA2 kernel is called, `is_causal` should be inferred from `causal_attention_mask` - if not is_flash_attention_requested(self.config): - if attention_mask is not None and causal_attention_mask is not None: - attention_mask = attention_mask + causal_attention_mask - elif causal_attention_mask is not None: - attention_mask = causal_attention_mask - else: - self.is_causal = causal_attention_mask is not None attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward @@ -342,9 +332,9 @@ def forward( keys, values, attention_mask, - is_causal=self.is_causal, scaling=self.scale, dropout=0.0 if not self.training else self.dropout, + **kwargs, ) attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() @@ -385,8 +375,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -404,8 +394,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -495,10 +485,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutput: r""" Args: @@ -512,13 +502,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -545,9 +528,10 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) + hidden_states = layer_outputs[0] if output_attentions: @@ -561,10 +545,10 @@ def forward( ) -class CLIPSegTextTransformer(nn.Module): +class CLIPSegTextTransformer(CLIPSegPreTrainedModel): def __init__(self, config: CLIPSegTextConfig): - super().__init__() - self.config = config + super().__init__(config) + embed_dim = config.hidden_size self.embeddings = CLIPSegTextEmbeddings(config) self.encoder = CLIPSegEncoder(config) @@ -573,6 +557,8 @@ def __init__(self, config: CLIPSegTextConfig): # For `pooled_output` computation self.eos_token_id = config.eos_token_id + self.post_init() + @auto_docstring def forward( self, @@ -582,6 +568,7 @@ def forward( output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutputWithPooling: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -597,23 +584,23 @@ def forward( hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) - # CLIPSeg's text model uses causal mask, prepare it here. - # https://github.com/openai/CLIPSeg/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clipseg/model.py#L324 - causal_attention_mask = _create_4d_causal_attention_mask( - input_shape, hidden_states.dtype, device=hidden_states.device + attention_mask = create_causal_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + cache_position=torch.arange(hidden_states.shape[1], device=hidden_states.device), + past_key_values=None, ) - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + kwargs.pop("is_causal", None) encoder_outputs = self.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + is_causal=True, + **kwargs, ) last_hidden_state = encoder_outputs[0] diff --git a/src/transformers/models/clvp/modeling_clvp.py b/src/transformers/models/clvp/modeling_clvp.py index 7b7ca1d12840..d1c5ee39318d 100644 --- a/src/transformers/models/clvp/modeling_clvp.py +++ b/src/transformers/models/clvp/modeling_clvp.py @@ -27,7 +27,7 @@ from ...activations import ACT2FN, get_activation from ...cache_utils import Cache, DynamicCache from ...generation import GenerationConfig, GenerationMixin -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_outputs import ( BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, @@ -923,9 +923,11 @@ def forward( raise ValueError("You have to specify either input_ids or inputs_embeds") # expand attention_mask and create position_ids if needed - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) if position_ids is None: device = input_ids.device if input_ids is not None else inputs_embeds.device @@ -1069,23 +1071,32 @@ def forward( ) use_cache = False + if inputs_embeds is None: + inputs_embeds = self.input_embeds_layer(input_ids) + if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + input_shape[-1], device=inputs_embeds.device + ) if position_ids is None: position_ids = torch.arange( past_key_values_length, input_shape[-1] + past_key_values_length, dtype=torch.long, device=device ) position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1]) - if inputs_embeds is None: - inputs_embeds = self.input_embeds_layer(input_ids) position_embeds = self.position_embeds_layer(position_ids) inputs_embeds = inputs_embeds + position_embeds - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) hidden_states = inputs_embeds diff --git a/src/transformers/models/codegen/modeling_codegen.py b/src/transformers/models/codegen/modeling_codegen.py index 8530fb4de2e1..200894a4ab7d 100644 --- a/src/transformers/models/codegen/modeling_codegen.py +++ b/src/transformers/models/codegen/modeling_codegen.py @@ -14,7 +14,6 @@ """PyTorch CodeGen model.""" import math -from typing import Union import torch from torch import nn @@ -23,25 +22,17 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel from ...utils import ( auto_docstring, - is_torch_flex_attn_available, logging, ) -from ...utils.generic import is_flash_attention_requested from .configuration_codegen import CodeGenConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -367,8 +358,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) hidden_states = inputs_embeds @@ -420,131 +416,6 @@ def forward( attentions=all_self_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/conditional_detr/modeling_conditional_detr.py b/src/transformers/models/conditional_detr/modeling_conditional_detr.py index 92cbfc628954..12fb422ccb78 100644 --- a/src/transformers/models/conditional_detr/modeling_conditional_detr.py +++ b/src/transformers/models/conditional_detr/modeling_conditional_detr.py @@ -1125,14 +1125,11 @@ def forward( hidden_states = inputs_embeds hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - # expand attention_mask - if attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - attention_mask = create_bidirectional_mask( - config=self.config, - input_embeds=inputs_embeds, - attention_mask=attention_mask, - ) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) for encoder_layer in self.layers: # we add spatial_position_embeddings as extra input to the encoder_layer diff --git a/src/transformers/models/dab_detr/modeling_dab_detr.py b/src/transformers/models/dab_detr/modeling_dab_detr.py index 745554529d96..f02418ff5c1e 100644 --- a/src/transformers/models/dab_detr/modeling_dab_detr.py +++ b/src/transformers/models/dab_detr/modeling_dab_detr.py @@ -22,7 +22,7 @@ from ... import initialization as init from ...activations import ACT2FN from ...backbone_utils import load_backbone -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions, Seq2SeqModelOutput from ...modeling_utils import PreTrainedModel @@ -919,10 +919,11 @@ def forward( hidden_states = inputs_embeds - # expand attention_mask - if attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -1048,7 +1049,6 @@ def forward( if inputs_embeds is not None: hidden_states = inputs_embeds - input_shape = inputs_embeds.size()[:-1] # decoder layers all_hidden_states = () if output_hidden_states else None @@ -1061,9 +1061,11 @@ def forward( # expand encoder attention mask if encoder_hidden_states is not None and memory_key_padding_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - memory_key_padding_mask = _prepare_4d_attention_mask( - memory_key_padding_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + memory_key_padding_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=memory_key_padding_mask, + encoder_hidden_states=encoder_hidden_states, ) for layer_id, decoder_layer in enumerate(self.layers): diff --git a/src/transformers/models/decision_transformer/modeling_decision_transformer.py b/src/transformers/models/decision_transformer/modeling_decision_transformer.py index afe70ede6667..99553520a4cc 100755 --- a/src/transformers/models/decision_transformer/modeling_decision_transformer.py +++ b/src/transformers/models/decision_transformer/modeling_decision_transformer.py @@ -23,6 +23,7 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -52,16 +53,6 @@ def eager_attention_forward(module, query, key, value, attention_mask, **kwargs) if module.scale_attn_by_inverse_layer_idx: attn_weights = attn_weights / float(module.layer_idx + 1) - if not module.is_cross_attention: - # if only "normal" attention layer implements causal mask - query_length, key_length = query.size(-2), key.size(-2) - causal_mask = module.bias[:, :, key_length - query_length : key_length, :key_length] - mask_value = torch.finfo(attn_weights.dtype).min - # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`. - # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device` - mask_value = torch.full([], mask_value, dtype=attn_weights.dtype, device=attn_weights.device) - attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value) - if attention_mask is not None: # Apply the attention mask causal_mask = attention_mask[:, :, :, : key.shape[-2]] @@ -84,15 +75,6 @@ class DecisionTransformerGPT2Attention(nn.Module): def __init__(self, config, is_cross_attention=False, layer_idx=None): super().__init__() self.config = config - max_positions = config.max_position_embeddings - self.register_buffer( - "bias", - torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view( - 1, 1, max_positions, max_positions - ), - persistent=False, - ) - self.embed_dim = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads @@ -144,16 +126,6 @@ def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None): attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor) attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len) - if not self.is_cross_attention: - # if only "normal" attention layer implements causal mask - query_length, key_length = query.size(-2), key.size(-2) - causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length] - mask_value = torch.finfo(attn_weights.dtype).min - # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`. - # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device` - mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype, device=attn_weights.device) - attn_weights = torch.where(causal_mask, attn_weights, mask_value) - if attention_mask is not None: # Apply the attention mask attn_weights = attn_weights + attention_mask @@ -366,6 +338,9 @@ class DecisionTransformerGPT2PreTrainedModel(PreTrainedModel): supports_gradient_checkpointing = True _can_compile_fullgraph = False + # No longer used as we directly use our masks instead + _keys_to_ignore_on_load_unexpected = ["attn.bias", "crossattention.bias"] + @torch.no_grad() def _init_weights(self, module): """Initialize the weights.""" @@ -382,14 +357,6 @@ def _init_weights(self, module): if "c_proj" in name and "weight" in name: # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block init.normal_(p, mean=0.0, std=self.config.initializer_range / math.sqrt(2 * self.config.n_layer)) - elif isinstance(module, DecisionTransformerGPT2Attention): - max_positions = module.config.max_position_embeddings - init.copy_( - module.bias, - torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view( - 1, 1, max_positions, max_positions - ), - ) class DecisionTransformerGPT2Model(DecisionTransformerGPT2PreTrainedModel): @@ -448,15 +415,11 @@ def forward( self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) input_shape = input_ids.size() input_ids = input_ids.view(-1, input_shape[-1]) - batch_size = input_ids.shape[0] elif inputs_embeds is not None: input_shape = inputs_embeds.size()[:-1] - batch_size = inputs_embeds.shape[0] else: raise ValueError("You have to specify either input_ids or inputs_embeds") - device = input_ids.device if input_ids is not None else inputs_embeds.device - if token_type_ids is not None: token_type_ids = token_type_ids.view(-1, input_shape[-1]) @@ -477,35 +440,23 @@ def forward( position_ids = cache_position.unsqueeze(0) # Attention mask. - if attention_mask is not None: - if batch_size <= 0: - raise ValueError("batch_size has to be defined and > 0") - attention_mask = attention_mask.view(batch_size, -1) - # We create a 3D attention mask from a 2D tensor mask. - # Sizes are [batch_size, 1, 1, to_seq_length] - # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length] - # this attention mask is more simple than the triangular masking of causal attention - # used in OpenAI GPT, we just need to prepare the broadcast dimension here. - attention_mask = attention_mask[:, None, None, :] - - # Since attention_mask is 1.0 for positions we want to attend and 0.0 for - # masked positions, this operation will create a tensor which is 0.0 for - # positions we want to attend and the dtype's smallest value for masked positions. - # Since we are adding it to the raw scores before the softmax, this is - # effectively the same as removing these entirely. - attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility - attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min - - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if self.config.add_cross_attention and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) - encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_attention_mask = None + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + encoder_attention_mask = None + if encoder_hidden_states is not None: + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) if inputs_embeds is None: inputs_embeds = self.wte(input_ids) diff --git a/src/transformers/models/detr/modeling_detr.py b/src/transformers/models/detr/modeling_detr.py index 63d75648e600..5b7dfefa3d45 100644 --- a/src/transformers/models/detr/modeling_detr.py +++ b/src/transformers/models/detr/modeling_detr.py @@ -971,14 +971,11 @@ def forward( hidden_states = inputs_embeds hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - # expand attention_mask - if attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - attention_mask = create_bidirectional_mask( - config=self.config, - input_embeds=inputs_embeds, - attention_mask=attention_mask, - ) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) for encoder_layer in self.layers: # we add spatial_position_embeddings as extra input to the encoder_layer @@ -1057,21 +1054,18 @@ def forward( if inputs_embeds is not None: hidden_states = inputs_embeds - # expand decoder attention mask (for self-attention on object queries) if attention_mask is not None: - # [batch_size, num_queries] -> [batch_size, 1, num_queries, num_queries] attention_mask = create_bidirectional_mask( config=self.config, - input_embeds=inputs_embeds, + input_embeds=hidden_states, attention_mask=attention_mask, ) # expand encoder attention mask (for cross-attention on encoder outputs) if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] encoder_attention_mask = create_bidirectional_mask( config=self.config, - input_embeds=inputs_embeds, + input_embeds=hidden_states, attention_mask=encoder_attention_mask, encoder_hidden_states=encoder_hidden_states, ) diff --git a/src/transformers/models/falcon/modeling_falcon.py b/src/transformers/models/falcon/modeling_falcon.py index b0f7fd90e029..2d01b6e98cb0 100644 --- a/src/transformers/models/falcon/modeling_falcon.py +++ b/src/transformers/models/falcon/modeling_falcon.py @@ -24,11 +24,9 @@ from ... import initialization as init from ...activations import get_activation -from ...cache_utils import Cache, DynamicCache, StaticCache +from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import ( - AttentionMaskConverter, -) +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -47,7 +45,7 @@ auto_docstring, logging, ) -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils.generic import maybe_autocast from .configuration_falcon import FalconConfig @@ -357,8 +355,8 @@ def forward( if self.config._attn_implementation == "sdpa" and not output_attentions: # We dispatch to SDPA's Flash Attention or Efficient kernels via this if statement instead of an # inline conditional assignment to support both torch.compile's `dynamic=True` and `fullgraph=True` - # The query_length > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not - # create a causal mask in case query_length == 1. + # The query_length > 1 is necessary to match with a bidirectional attention mask we do not have + # a causal pattern in those cases. is_causal = self.is_causal and attention_mask is None and query_length > 1 attn_output = torch.nn.functional.scaled_dot_product_attention( query_layer, @@ -782,9 +780,31 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions, alibi + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + # Force mask creation for alibi + and_mask_function=lambda *args: torch.tensor(True, dtype=torch.bool), ) + if alibi is not None and causal_mask is not None and causal_mask.ndim == 4: + min_dtype = torch.finfo(inputs_embeds.dtype).min + + # Only using non-bool mask for alibi + if causal_mask.dtype == torch.bool: + causal_mask = torch.where( + causal_mask, torch.tensor(0.0, device=causal_mask.device, dtype=inputs_embeds.dtype), min_dtype + ) + + # We take care to integrate alibi bias in the causal_mask here + alibi = alibi.reshape(batch_size, -1, *alibi.shape[1:]) + causal_mask = torch.masked_fill( + alibi / math.sqrt(self.config.hidden_size // self.num_heads), + causal_mask < -1, + min_dtype, + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -830,147 +850,6 @@ def forward( attentions=all_self_attentions, ) - def _update_causal_mask( - self, - attention_mask: torch.Tensor, - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool, - alibi: torch.Tensor, - ): - # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static - # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. - # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using - # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 - - if is_flash_attention_requested(self.config): - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_static_cache = isinstance(past_key_values, StaticCache) - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if ( - self.config._attn_implementation == "sdpa" - and not using_static_cache - and not output_attentions - and alibi is None - ): - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype, device = input_tensor.dtype, input_tensor.device - min_dtype = torch.finfo(dtype).min - batch_size, sequence_length, _ = input_tensor.shape - if using_static_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - device=device, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - # We take care to integrate alibi bias in the causal_mask here - if alibi is not None: - alibi = alibi.reshape(batch_size, -1, *alibi.shape[1:]) - causal_mask = torch.masked_fill( - alibi / math.sqrt(self.config.hidden_size // self.num_heads), - causal_mask < -1, - min_dtype, - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 498697a321f1..96f10b6d51a8 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -36,7 +36,7 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast @@ -44,7 +44,7 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils.generic import maybe_autocast from ...utils.import_utils import is_causal_conv1d_available, is_mamba_2_ssm_available from .configuration_falcon_h1 import FalconH1Config @@ -1319,8 +1319,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) mamba_mask = self._update_mamba_mask(attention_mask, cache_position) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1380,122 +1385,6 @@ def _update_mamba_mask(self, attention_mask, cache_position): mamba_mask = None return mamba_mask - def _update_causal_mask( - self, - attention_mask: torch.Tensor, - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: FalconHybridMambaAttentionDynamicCache, - output_attentions: bool, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_attention_mask = (attention_mask[:, None, None, :] == attention_mask[:, None, :, None])[ - :, :, -sequence_length:, : - ].to(dtype) - padding_mask = causal_mask[:, :, :, :mask_length] + padding_attention_mask - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin): diff --git a/src/transformers/models/falcon_h1/modular_falcon_h1.py b/src/transformers/models/falcon_h1/modular_falcon_h1.py index ac0b2462d48d..b50e20ac235d 100644 --- a/src/transformers/models/falcon_h1/modular_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modular_falcon_h1.py @@ -46,14 +46,13 @@ from ... import initialization as init from ...cache_utils import Cache -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging -from ...utils.generic import is_flash_attention_requested from ...utils.import_utils import is_causal_conv1d_available, is_mamba_2_ssm_available from .configuration_falcon_h1 import FalconH1Config @@ -1045,8 +1044,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) mamba_mask = self._update_mamba_mask(attention_mask, cache_position) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1106,122 +1110,6 @@ def _update_mamba_mask(self, attention_mask, cache_position): mamba_mask = None return mamba_mask - def _update_causal_mask( - self, - attention_mask: torch.Tensor, - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: FalconHybridMambaAttentionDynamicCache, - output_attentions: bool, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_attention_mask = (attention_mask[:, None, None, :] == attention_mask[:, None, :, None])[ - :, :, -sequence_length:, : - ].to(dtype) - padding_mask = causal_mask[:, :, :, :mask_length] + padding_attention_mask - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - class FalconH1ForCausalLM(LlamaForCausalLM): def forward( diff --git a/src/transformers/models/git/modeling_git.py b/src/transformers/models/git/modeling_git.py index 1baea7a449a8..b0febbd1a446 100644 --- a/src/transformers/models/git/modeling_git.py +++ b/src/transformers/models/git/modeling_git.py @@ -35,15 +35,16 @@ CausalLMOutputWithPast, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward from ...utils import ( ModelOutput, + TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int, ) -from ...utils.generic import is_flash_attention_requested from .configuration_git import GitConfig, GitVisionConfig @@ -655,7 +656,6 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Input shape: Batch x Time x Channel""" @@ -669,15 +669,6 @@ def forward( queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - # CLIP text model uses both `causal_attention_mask` and `attention_mask` - # in case FA2 kernel is called, `is_causal` should be inferred from `causal_attention_mask` - if not is_flash_attention_requested(self.config): - if attention_mask is not None and causal_attention_mask is not None: - attention_mask = attention_mask + causal_attention_mask - elif causal_attention_mask is not None: - attention_mask = causal_attention_mask - else: - self.is_causal = causal_attention_mask is not None attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward @@ -715,8 +706,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -734,8 +725,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -773,10 +764,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutput: r""" Args: @@ -790,13 +781,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -823,8 +807,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] diff --git a/src/transformers/models/gpt2/modeling_gpt2.py b/src/transformers/models/gpt2/modeling_gpt2.py index a55437c0ae5a..ced92b31b9d2 100644 --- a/src/transformers/models/gpt2/modeling_gpt2.py +++ b/src/transformers/models/gpt2/modeling_gpt2.py @@ -26,8 +26,7 @@ from ...activations import ACT2FN, get_activation from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...masking_utils import create_causal_mask -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutputWithPastAndCrossAttentions, @@ -43,7 +42,7 @@ auto_docstring, logging, ) -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils.generic import maybe_autocast from .configuration_gpt2 import GPT2Config @@ -62,16 +61,6 @@ def eager_attention_forward(module, query, key, value, attention_mask, **kwargs) if module.scale_attn_by_inverse_layer_idx: attn_weights = attn_weights / float(module.layer_idx + 1) - if not module.is_cross_attention: - # if only "normal" attention layer implements causal mask - query_length, key_length = query.size(-2), key.size(-2) - causal_mask = module.bias[:, :, key_length - query_length : key_length, :key_length] - mask_value = torch.finfo(attn_weights.dtype).min - # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`. - # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device` - mask_value = torch.full([], mask_value, dtype=attn_weights.dtype, device=attn_weights.device) - attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value) - if attention_mask is not None: # Apply the attention mask causal_mask = attention_mask[:, :, :, : key.shape[-2]] @@ -93,15 +82,6 @@ class GPT2Attention(nn.Module): def __init__(self, config, is_cross_attention=False, layer_idx=None): super().__init__() self.config = config - max_positions = config.max_position_embeddings - self.register_buffer( - "bias", - torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view( - 1, 1, max_positions, max_positions - ), - persistent=False, - ) - self.embed_dim = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads @@ -153,16 +133,6 @@ def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None): attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor) attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len) - if not self.is_cross_attention: - # if only "normal" attention layer implements causal mask - query_length, key_length = query.size(-2), key.size(-2) - causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length] - mask_value = torch.finfo(attn_weights.dtype).min - # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`. - # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device` - mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype, device=attn_weights.device) - attn_weights = torch.where(causal_mask, attn_weights, mask_value) - if attention_mask is not None: # Apply the attention mask attn_weights = attn_weights + attention_mask @@ -475,6 +445,9 @@ class GPT2PreTrainedModel(PreTrainedModel): _supports_attention_backend = True _can_compile_fullgraph = True + # No longer used as we directly use our masks instead + _keys_to_ignore_on_load_unexpected = ["attn.bias", "crossattention.bias"] + @torch.no_grad() def _init_weights(self, module): """Initialize the weights.""" @@ -490,14 +463,6 @@ def _init_weights(self, module): elif isinstance(module, nn.LayerNorm): init.zeros_(module.bias) init.ones_(module.weight) - elif isinstance(module, GPT2Attention): - max_positions = module.config.max_position_embeddings - init.copy_( - module.bias, - torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view( - 1, 1, max_positions, max_positions - ), - ) # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale @@ -622,8 +587,6 @@ def forward( else: raise ValueError("You have to specify either input_ids or inputs_embeds") - device = input_ids.device if input_ids is not None else inputs_embeds.device - if token_type_ids is not None: token_type_ids = token_type_ids.view(-1, input_shape[-1]) @@ -657,7 +620,6 @@ def forward( hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device) # Attention mask. - # ._update_causal_mask() and ._prepare_4d_causal_attention_mask_with_cache_position() copied from LlamaModel if attention_mask is not None and attention_mask.ndim < 4: attention_mask = attention_mask.view(batch_size, -1) @@ -670,22 +632,14 @@ def forward( position_ids=position_ids, ) - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - _use_sdpa = self._attn_implementation == "sdpa" and output_attentions is False - if self.config.add_cross_attention and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) - if _use_sdpa: - encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa( - mask=encoder_attention_mask, dtype=inputs_embeds.dtype, tgt_len=input_shape[-1] - ) - elif not is_flash_attention_requested(requested_attention_implementation=self._attn_implementation): - encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_attention_mask = None + encoder_attention_mask = None + if encoder_hidden_states is not None: + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) if token_type_ids is not None: token_type_embeds = self.wte(token_type_ids) @@ -711,6 +665,7 @@ def forward( encoder_attention_mask=encoder_attention_mask, use_cache=use_cache, output_attentions=output_attentions, + position_ids=position_ids, **kwargs, ) diff --git a/src/transformers/models/gpt_neo/modeling_gpt_neo.py b/src/transformers/models/gpt_neo/modeling_gpt_neo.py index e756bc8d8ea9..57140b9ad47b 100755 --- a/src/transformers/models/gpt_neo/modeling_gpt_neo.py +++ b/src/transformers/models/gpt_neo/modeling_gpt_neo.py @@ -13,8 +13,6 @@ # limitations under the License. """PyTorch GPT Neo model.""" -from typing import Union - import torch from torch import nn from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss @@ -23,7 +21,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -38,19 +36,11 @@ from ...modeling_utils import PreTrainedModel from ...utils import ( auto_docstring, - is_torch_flex_attn_available, logging, ) -from ...utils.generic import is_flash_attention_requested from .configuration_gpt_neo import GPTNeoConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - if is_flash_attn_available(): from ...modeling_flash_attention_utils import _flash_attention_forward @@ -473,8 +463,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) position_embeds = self.wpe(position_ids) @@ -526,131 +521,6 @@ def forward( attentions=all_self_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py index a3e486e2c39f..396396f9fff0 100755 --- a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py +++ b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py @@ -15,7 +15,7 @@ import math from collections.abc import Callable -from typing import Optional, Union +from typing import Optional import torch from torch import Tensor, nn @@ -24,21 +24,15 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils import auto_docstring, logging +from ...utils.generic import maybe_autocast from .configuration_gpt_neox_japanese import GPTNeoXJapaneseConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -471,8 +465,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) hidden_states = inputs_embeds @@ -515,131 +514,6 @@ def forward( attentions=all_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/gptj/modeling_gptj.py b/src/transformers/models/gptj/modeling_gptj.py index afb8a660da68..caa940a4e865 100644 --- a/src/transformers/models/gptj/modeling_gptj.py +++ b/src/transformers/models/gptj/modeling_gptj.py @@ -14,7 +14,6 @@ """PyTorch GPT-J model.""" import math -from typing import Union import torch from torch import nn @@ -24,7 +23,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -34,17 +33,10 @@ SequenceClassifierOutputWithPast, ) from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested +from ...utils import auto_docstring, logging from .configuration_gptj import GPTJConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - if is_flash_attn_available(): from ...modeling_flash_attention_utils import _flash_attention_forward @@ -529,8 +521,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) hidden_states = inputs_embeds @@ -582,129 +579,6 @@ def forward( attentions=all_self_attentions, ) - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index 3edf31765bea..3144048eae69 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -30,12 +30,8 @@ from ... import initialization as init from ...cache_utils import Cache from ...generation import GenerationMixin -from ...integrations import ( - lazy_load_kernel, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from ...integrations.hub_kernels import lazy_load_kernel from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, MoeCausalLMOutputWithPast, MoeModelOutputWithPast diff --git a/src/transformers/models/groupvit/modeling_groupvit.py b/src/transformers/models/groupvit/modeling_groupvit.py index a0e83e5fbfb8..e04a1c1ba2b9 100644 --- a/src/transformers/models/groupvit/modeling_groupvit.py +++ b/src/transformers/models/groupvit/modeling_groupvit.py @@ -23,7 +23,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel @@ -535,7 +535,7 @@ def forward( cat_x = self.concat_x(x, group_token) for layer in self.layers: - layer_out = layer(cat_x, attention_mask=None, causal_attention_mask=None) + layer_out = layer(cat_x, attention_mask=None) cat_x = layer_out[0] x, group_token = self.split_x(cat_x) @@ -610,9 +610,9 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, encoder_hidden_states: torch.FloatTensor | None = None, output_attentions: bool | None = False, + **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: """Input shape: Batch x Time x Channel""" @@ -642,16 +642,6 @@ def forward( f" {attn_weights.size()}" ) - # apply the causal_attention_mask first - if causal_attention_mask is not None: - if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is" - f" {causal_attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - if attention_mask is not None: if attention_mask.size() != (bsz, 1, tgt_len, src_len): raise ValueError( @@ -705,8 +695,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -724,8 +714,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -864,10 +854,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutput: r""" Args: @@ -881,13 +871,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -914,8 +897,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] @@ -933,10 +916,9 @@ def forward( ) -class GroupViTTextTransformer(nn.Module): +class GroupViTTextTransformer(GroupViTPreTrainedModel): def __init__(self, config: GroupViTTextConfig): - super().__init__() - self.config = config + super().__init__(config) embed_dim = config.hidden_size self.embeddings = GroupViTTextEmbeddings(config) self.encoder = GroupViTTextEncoder(config) @@ -945,6 +927,8 @@ def __init__(self, config: GroupViTTextConfig): # For `pooled_output` computation self.eos_token_id = config.eos_token_id + self.post_init() + @auto_docstring def forward( self, @@ -954,6 +938,7 @@ def forward( output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutputWithPooling: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -969,24 +954,23 @@ def forward( hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) - # CLIP's text model uses causal mask, prepare it here. - # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 - causal_attention_mask = _create_4d_causal_attention_mask( - input_shape, hidden_states.dtype, device=hidden_states.device + attention_mask = create_causal_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + cache_position=torch.arange(hidden_states.shape[1], device=hidden_states.device), + past_key_values=None, ) - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) - + kwargs.pop("is_causal", None) encoder_outputs = self.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + is_causal=True, + **kwargs, ) last_hidden_state = encoder_outputs[0] diff --git a/src/transformers/models/idefics/vision.py b/src/transformers/models/idefics/vision.py index adcf7f2167da..2cc51ac34993 100644 --- a/src/transformers/models/idefics/vision.py +++ b/src/transformers/models/idefics/vision.py @@ -24,8 +24,10 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack from ...utils import ( ModelOutput, + TransformersKwargs, can_return_tuple, logging, ) @@ -292,8 +294,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -311,8 +313,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -350,10 +352,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutput: r""" Args: @@ -367,13 +369,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -400,8 +395,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] diff --git a/src/transformers/models/idefics2/modeling_idefics2.py b/src/transformers/models/idefics2/modeling_idefics2.py index 455f1be57dd9..5047e46c9673 100644 --- a/src/transformers/models/idefics2/modeling_idefics2.py +++ b/src/transformers/models/idefics2/modeling_idefics2.py @@ -23,14 +23,14 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging -from ...utils.generic import check_model_inputs, is_flash_attention_requested +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_idefics2 import Idefics2Config, Idefics2PerceiverConfig, Idefics2VisionConfig @@ -489,13 +489,11 @@ def forward( hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) patch_attention_mask = patch_attention_mask.view(batch_size, -1) - # The call to `_upad_input` in `_flash_attention_forward` is expensive - # So when the `patch_attention_mask` is full of 1s (i.e. attending to the whole sequence), - # avoiding passing the attention_mask, which is equivalent to attending to the full sequence - if not is_torchdynamo_compiling() and not torch.any(~patch_attention_mask): - patch_attention_mask = None - elif not is_flash_attention_requested(self.config): - patch_attention_mask = _prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype) + patch_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=patch_attention_mask, + ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, @@ -735,10 +733,10 @@ def forward( (attention_mask.size(0), latents.size(1)), dtype=attention_mask.dtype, device=attention_mask.device ) attention_mask = torch.cat([attention_mask, latent_attention_mask], dim=-1) - attention_mask = ( - _prepare_4d_attention_mask(attention_mask, latents.dtype, tgt_len=self.n_latents) - if not is_flash_attention_requested(self.config) - else attention_mask + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=latents, + attention_mask=attention_mask, ) compressed_context = latents diff --git a/src/transformers/models/informer/modeling_informer.py b/src/transformers/models/informer/modeling_informer.py index 34b7fe5ee668..255d6460eb4c 100644 --- a/src/transformers/models/informer/modeling_informer.py +++ b/src/transformers/models/informer/modeling_informer.py @@ -28,7 +28,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...masking_utils import create_bidirectional_mask, create_causal_mask -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -913,10 +912,11 @@ def forward( hidden_states = self.layernorm_embedding(hidden_states + embed_pos) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None diff --git a/src/transformers/models/informer/modular_informer.py b/src/transformers/models/informer/modular_informer.py index e0d01538301c..790bd874c4dc 100644 --- a/src/transformers/models/informer/modular_informer.py +++ b/src/transformers/models/informer/modular_informer.py @@ -19,7 +19,7 @@ from ... import initialization as init from ...cache_utils import Cache, EncoderDecoderCache -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -448,10 +448,11 @@ def forward( hidden_states = self.layernorm_embedding(hidden_states + embed_pos) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None diff --git a/src/transformers/models/kosmos2/modeling_kosmos2.py b/src/transformers/models/kosmos2/modeling_kosmos2.py index 95a5618d935c..93f5a2117a10 100644 --- a/src/transformers/models/kosmos2/modeling_kosmos2.py +++ b/src/transformers/models/kosmos2/modeling_kosmos2.py @@ -387,8 +387,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -406,8 +406,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -444,10 +444,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutput: r""" Args: @@ -461,13 +461,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -494,8 +487,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] @@ -1133,7 +1126,7 @@ class Kosmos2PreTrainedModel(PreTrainedModel): supports_gradient_checkpointing = True _no_split_modules = ["Kosmos2VisionEncoderLayer", "Kosmos2TextBlock"] _supports_attention_backend = True - _supports_flash_attn = True + _supports_flash_attn = False # cuda device errors _supports_sdpa = True @torch.no_grad() diff --git a/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py b/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py index ffc43a557d3d..c61e602384b3 100644 --- a/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py +++ b/src/transformers/models/kosmos2_5/modeling_kosmos2_5.py @@ -16,7 +16,7 @@ import math from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Union +from typing import Any import torch from torch import nn @@ -25,7 +25,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -42,7 +42,6 @@ add_start_docstrings, add_start_docstrings_to_model_forward, can_return_tuple, - is_torch_flex_attn_available, logging, replace_return_docstrings, ) @@ -54,12 +53,6 @@ ) -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = Kosmos2_5Config @@ -916,131 +909,6 @@ def __init__(self, config: Kosmos2_5TextConfig): self.layer_norm = nn.LayerNorm(config.embed_dim, config.layer_norm_eps) self.gradient_checkpointing = False - # TODO (ydshieh): Remove this (to match Llama's code) - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # TODO (ydshieh): Remove this (to match Llama's code) - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - def forward( self, input_ids: torch.Tensor | None = None, @@ -1122,8 +990,12 @@ def forward( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py index d9f930b00b8f..4fe1c4f6482d 100644 --- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py +++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py @@ -21,7 +21,7 @@ import math import types from collections.abc import Callable -from typing import Optional, Union +from typing import Optional import torch import torch.nn as nn @@ -30,15 +30,15 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationConfig, GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast from ..auto import AutoModel from .configuration_kyutai_speech_to_text import KyutaiSpeechToTextConfig @@ -46,11 +46,6 @@ if is_flash_attn_available(): from ...modeling_flash_attention_utils import _flash_attention_forward -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - logger = logging.get_logger(__name__) @@ -863,8 +858,13 @@ def forward( causal_mask = None if attention_mask is not None: - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) # embed positions @@ -913,157 +913,6 @@ def forward( attentions=all_self_attns, ) - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and past_key_values is not None: - is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] - if is_padding_right: - raise ValueError( - "You are attempting to perform batched generation with padding_side='right'" - " this may lead to unexpected behaviour for Flash Attention version of KyutaiSpeechToText. Make sure to " - " call `tokenizer.padding_side = 'left'` before tokenizing the input. " - ) - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_static_cache = isinstance(past_key_values, StaticCache) - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - sliding_window=self.config.sliding_window, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - min_dtype = torch.finfo(dtype).min - sequence_length = input_tensor.shape[1] - # StaticCache - if using_static_cache: - target_length = past_key_values.get_max_cache_shape() - # DynamicCache or no cache - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - config=self.config, - past_key_values=past_key_values, - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - config: KyutaiSpeechToTextConfig, - past_key_values: Cache, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - config (`KyutaiSpeechToTextConfig`): - The model's configuration class - past_key_values (`Cache`): - The cache class that is being used currently to generate - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape( - -1, 1 - ) - text_config = config.get_text_config() - if getattr(text_config, "use_sliding_window", True) and text_config.sliding_window is not None: - # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also - # the check is needed to verify is current checkpoint was trained with sliding window or not - is_static_sliding_cache = isinstance(past_key_values, StaticCache) and all(past_key_values.is_sliding) - if not is_static_sliding_cache or sequence_length > target_length: - sliding_attend_mask = torch.arange(target_length, device=cache_position.device) <= ( - cache_position.reshape(-1, 1) - text_config.sliding_window - ) - diagonal_attend_mask.bitwise_or_(sliding_attend_mask) - causal_mask *= diagonal_attend_mask - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - if attention_mask.shape[-1] > target_length: - attention_mask = attention_mask[:, :target_length] - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - return causal_mask - @auto_docstring class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedModel, GenerationMixin): diff --git a/src/transformers/models/layoutlm/modeling_layoutlm.py b/src/transformers/models/layoutlm/modeling_layoutlm.py index d206be4301e7..70b058d33bd4 100644 --- a/src/transformers/models/layoutlm/modeling_layoutlm.py +++ b/src/transformers/models/layoutlm/modeling_layoutlm.py @@ -31,8 +31,9 @@ TokenClassifierOutput, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward -from ...utils import auto_docstring, can_return_tuple, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from .configuration_layoutlm import LayoutLMConfig @@ -170,7 +171,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.attention_head_size) @@ -226,7 +227,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_outputs = self.self( hidden_states, @@ -285,7 +286,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_attention_outputs = self.attention( hidden_states, @@ -325,7 +326,7 @@ def forward( output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None diff --git a/src/transformers/models/led/modeling_led.py b/src/transformers/models/led/modeling_led.py index 417c07ce7052..6c8cee33eeea 100755 --- a/src/transformers/models/led/modeling_led.py +++ b/src/transformers/models/led/modeling_led.py @@ -24,7 +24,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions from ...modeling_utils import PreTrainedModel @@ -1671,26 +1671,27 @@ def forward( past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - # create causal mask - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - combined_attention_mask = None - if input_shape[-1] > 1: - combined_attention_mask = _create_4d_causal_attention_mask( - input_shape, inputs_embeds.dtype, inputs_embeds.device, past_key_values_length=past_key_values_length + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + inputs_embeds.shape[1], device=inputs_embeds.device ) - if attention_mask is not None and combined_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - combined_attention_mask = combined_attention_mask + _prepare_4d_attention_mask_inverted( - attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + combined_attention_mask = None + if input_shape[-1] > 1: # only create a causal mask when we go over a single token + combined_attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) - # expand encoder attention mask - if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask_inverted( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] - ) + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) # embed positions positions = self.embed_positions(input_shape, past_key_values_length) diff --git a/src/transformers/models/llava_next/modeling_llava_next.py b/src/transformers/models/llava_next/modeling_llava_next.py index 39d11bce4965..25d06e21a568 100644 --- a/src/transformers/models/llava_next/modeling_llava_next.py +++ b/src/transformers/models/llava_next/modeling_llava_next.py @@ -718,61 +718,5 @@ def prepare_inputs_for_generation( return model_inputs - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - __all__ = ["LlavaNextForConditionalGeneration", "LlavaNextPreTrainedModel", "LlavaNextModel"] diff --git a/src/transformers/models/llava_next_video/modeling_llava_next_video.py b/src/transformers/models/llava_next_video/modeling_llava_next_video.py index 7d667e2c3abd..cc90624c7d0b 100644 --- a/src/transformers/models/llava_next_video/modeling_llava_next_video.py +++ b/src/transformers/models/llava_next_video/modeling_llava_next_video.py @@ -896,61 +896,6 @@ def prepare_inputs_for_generation( return model_inputs - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring def get_video_features( self, diff --git a/src/transformers/models/llava_onevision/modeling_llava_onevision.py b/src/transformers/models/llava_onevision/modeling_llava_onevision.py index ad9a648402bd..a849b7b9cd2f 100644 --- a/src/transformers/models/llava_onevision/modeling_llava_onevision.py +++ b/src/transformers/models/llava_onevision/modeling_llava_onevision.py @@ -866,61 +866,6 @@ def prepare_inputs_for_generation( return model_inputs - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring def get_video_features( self, diff --git a/src/transformers/models/longt5/modeling_longt5.py b/src/transformers/models/longt5/modeling_longt5.py index 2b6e914699a4..51679f564c17 100644 --- a/src/transformers/models/longt5/modeling_longt5.py +++ b/src/transformers/models/longt5/modeling_longt5.py @@ -15,7 +15,7 @@ import copy import math -from typing import Any, Union +from typing import Any import torch from torch import nn @@ -25,7 +25,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -38,20 +38,12 @@ DUMMY_INPUTS, DUMMY_MASK, auto_docstring, - is_torch_flex_attn_available, is_torchdynamo_compiling, logging, ) -from ...utils.generic import is_flash_attention_requested from .configuration_longt5 import LongT5Config -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -1344,14 +1336,12 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.is_decoder: - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache - if isinstance(past_key_values, EncoderDecoderCache) - else past_key_values, - output_attentions, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) # We use local attention in encoder self-attention, otherwise standard self & cross attentions are used elif self.config.encoder_attention_type == "local": @@ -1440,131 +1430,6 @@ def forward( cross_attentions=all_cross_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring class LongT5Model(LongT5PreTrainedModel): diff --git a/src/transformers/models/markuplm/modeling_markuplm.py b/src/transformers/models/markuplm/modeling_markuplm.py index 978b186f9544..b276e55c1cfc 100755 --- a/src/transformers/models/markuplm/modeling_markuplm.py +++ b/src/transformers/models/markuplm/modeling_markuplm.py @@ -31,8 +31,9 @@ TokenClassifierOutput, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward -from ...utils import auto_docstring, can_return_tuple, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from .configuration_markuplm import MarkupLMConfig @@ -364,7 +365,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.attention_head_size) @@ -405,7 +406,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_outputs = self.self( hidden_states, @@ -433,7 +434,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_attention_outputs = self.attention( hidden_states, @@ -473,7 +474,7 @@ def forward( output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None diff --git a/src/transformers/models/maskformer/modeling_maskformer.py b/src/transformers/models/maskformer/modeling_maskformer.py index 2aed70cdd1b6..24470bbc0cb2 100644 --- a/src/transformers/models/maskformer/modeling_maskformer.py +++ b/src/transformers/models/maskformer/modeling_maskformer.py @@ -24,13 +24,15 @@ from ... import initialization as init from ...activations import ACT2FN from ...backbone_utils import load_backbone -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithCrossAttentions from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache from ...utils import ( ModelOutput, + TransformersKwargs, auto_docstring, is_accelerate_available, is_scipy_available, @@ -554,6 +556,7 @@ def forward( encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ): """ Args: @@ -625,7 +628,10 @@ def forward( return outputs -class DetrDecoder(nn.Module): +class DetrDecoder(PreTrainedModel): + config: DetrConfig + base_model_prefix = "model" + """ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DetrDecoderLayer`]. @@ -641,8 +647,8 @@ class DetrDecoder(nn.Module): """ def __init__(self, config: DetrConfig): - super().__init__() - self.config = config + super().__init__(config) + self.dropout = config.dropout self.layerdrop = config.decoder_layerdrop @@ -652,6 +658,8 @@ def __init__(self, config: DetrConfig): self.gradient_checkpointing = False + self.post_init() + def forward( self, inputs_embeds=None, @@ -663,6 +671,7 @@ def forward( output_attentions=None, output_hidden_states=None, return_dict=None, + **kwargs: Unpack[TransformersKwargs], ): r""" Args: @@ -707,14 +716,13 @@ def forward( if inputs_embeds is not None: hidden_states = inputs_embeds - input_shape = inputs_embeds.size()[:-1] - # expand encoder attention mask - if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] - ) + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) # optional intermediate hidden states intermediate = () if self.config.auxiliary_loss else None @@ -741,6 +749,7 @@ def forward( encoder_hidden_states, # as a positional argument for gradient checkpointing encoder_attention_mask=encoder_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] diff --git a/src/transformers/models/mllama/modeling_mllama.py b/src/transformers/models/mllama/modeling_mllama.py index f5fbd8ed0b83..bd0edfa08f4a 100644 --- a/src/transformers/models/mllama/modeling_mllama.py +++ b/src/transformers/models/mllama/modeling_mllama.py @@ -15,7 +15,7 @@ import math from collections.abc import Callable -from typing import Optional, Union +from typing import Optional import torch import torch.nn.functional as F @@ -25,7 +25,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast @@ -35,17 +35,12 @@ ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging -from ...utils.generic import check_model_inputs, is_flash_attention_requested, maybe_autocast +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import check_model_inputs, maybe_autocast from ...utils.output_capturing import OutputRecorder from .configuration_mllama import MllamaConfig, MllamaTextConfig, MllamaVisionConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - logger = logging.get_logger(__name__) @@ -852,131 +847,6 @@ def _init_weights(self, module): init.copy_(module.inv_freq, buffer_value) init.copy_(module.original_inv_freq, buffer_value) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" @@ -1286,7 +1156,14 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values) + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) # decoder layers diff --git a/src/transformers/models/moshi/modeling_moshi.py b/src/transformers/models/moshi/modeling_moshi.py index aab096682f68..59d444f51590 100644 --- a/src/transformers/models/moshi/modeling_moshi.py +++ b/src/transformers/models/moshi/modeling_moshi.py @@ -16,7 +16,7 @@ import math from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Optional, Union +from typing import Any, Optional import torch import torch.nn as nn @@ -26,14 +26,14 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationConfig, GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput, Seq2SeqLMOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils import auto_docstring, logging +from ...utils.generic import maybe_autocast from ..auto.modeling_auto import AutoModel from .configuration_moshi import MoshiConfig, MoshiDepthConfig @@ -41,11 +41,6 @@ if is_flash_attn_available(): from ...modeling_flash_attention_utils import _flash_attention_forward -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - logger = logging.get_logger(__name__) @@ -976,8 +971,13 @@ def forward( causal_mask = None if attention_mask is not None: - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) # decoder layers @@ -1032,157 +1032,6 @@ def forward( attentions=all_self_attns, ) - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and past_key_values is not None: - is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] - if is_padding_right: - raise ValueError( - "You are attempting to perform batched generation with padding_side='right'" - " this may lead to unexpected behaviour for Flash Attention version of Moshi. Make sure to " - " call `tokenizer.padding_side = 'left'` before tokenizing the input. " - ) - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_static_cache = isinstance(past_key_values, StaticCache) - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - sliding_window=self.config.sliding_window, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - min_dtype = torch.finfo(dtype).min - sequence_length = input_tensor.shape[1] - # StaticCache - if using_static_cache: - target_length = past_key_values.get_max_cache_shape() - # DynamicCache or no cache - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - config=self.config, - past_key_values=past_key_values, - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - config: MoshiDepthConfig, - past_key_values: Cache, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - config (`MoshiDepthConfig`): - The model's configuration class - past_key_values (`Cache`): - The cache class that is being used currently to generate - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape( - -1, 1 - ) - text_config = config.get_text_config() - if getattr(text_config, "use_sliding_window", True) and text_config.sliding_window is not None: - # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also - # the check is needed to verify is current checkpoint was trained with sliding window or not - is_static_sliding_cache = isinstance(past_key_values, StaticCache) and all(past_key_values.is_sliding) - if not is_static_sliding_cache or sequence_length > target_length: - sliding_attend_mask = torch.arange(target_length, device=cache_position.device) <= ( - cache_position.reshape(-1, 1) - text_config.sliding_window - ) - diagonal_attend_mask.bitwise_or_(sliding_attend_mask) - causal_mask *= diagonal_attend_mask - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - if attention_mask.shape[-1] > target_length: - attention_mask = attention_mask[:, :target_length] - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - return causal_mask - @auto_docstring class MoshiModel(MoshiPreTrainedModel): @@ -1246,8 +1095,13 @@ def forward( causal_mask = None if attention_mask is not None: - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) # embed positions @@ -1296,157 +1150,6 @@ def forward( attentions=all_self_attns, ) - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and past_key_values is not None: - is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] - if is_padding_right: - raise ValueError( - "You are attempting to perform batched generation with padding_side='right'" - " this may lead to unexpected behaviour for Flash Attention version of Moshi. Make sure to " - " call `tokenizer.padding_side = 'left'` before tokenizing the input. " - ) - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_static_cache = isinstance(past_key_values, StaticCache) - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - sliding_window=self.config.sliding_window, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - min_dtype = torch.finfo(dtype).min - sequence_length = input_tensor.shape[1] - # StaticCache - if using_static_cache: - target_length = past_key_values.get_max_cache_shape() - # DynamicCache or no cache - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - config=self.config, - past_key_values=past_key_values, - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - config: MoshiConfig, - past_key_values: Cache, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - config (`MoshiConfig`): - The model's configuration class - past_key_values (`Cache`): - The cache class that is being used currently to generate - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape( - -1, 1 - ) - text_config = config.get_text_config() - if getattr(text_config, "use_sliding_window", True) and text_config.sliding_window is not None: - # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also - # the check is needed to verify is current checkpoint was trained with sliding window or not - is_static_sliding_cache = isinstance(past_key_values, StaticCache) and all(past_key_values.is_sliding) - if not is_static_sliding_cache or sequence_length > target_length: - sliding_attend_mask = torch.arange(target_length, device=cache_position.device) <= ( - cache_position.reshape(-1, 1) - text_config.sliding_window - ) - diagonal_attend_mask.bitwise_or_(sliding_attend_mask) - causal_mask *= diagonal_attend_mask - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - if attention_mask.shape[-1] > target_length: - attention_mask = attention_mask[:, :target_length] - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/mpt/modeling_mpt.py b/src/transformers/models/mpt/modeling_mpt.py index d80b2815f4ab..553243db0437 100644 --- a/src/transformers/models/mpt/modeling_mpt.py +++ b/src/transformers/models/mpt/modeling_mpt.py @@ -22,7 +22,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import _prepare_4d_causal_attention_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutputWithPastAndCrossAttentions, @@ -317,19 +317,21 @@ def forward( all_hidden_states = () if output_hidden_states else None # Compute alibi tensor: check build_alibi_tensor documentation - past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - seq_length_with_past = seq_length + past_key_values_length - if attention_mask is None: - attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device) - else: - attention_mask = attention_mask.to(hidden_states.device) - alibi = self.build_mpt_alibi_tensor(self.num_heads, self.config.max_seq_len, device=hidden_states.device) - causal_mask = _prepare_4d_causal_attention_mask( - attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length - ) - causal_mask = causal_mask.bool() + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device + ) + + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + ).to(torch.bool) for block in self.blocks: if output_hidden_states: diff --git a/src/transformers/models/musicgen/modeling_musicgen.py b/src/transformers/models/musicgen/modeling_musicgen.py index 9bfe09749fa8..6cb6adeafa89 100644 --- a/src/transformers/models/musicgen/modeling_musicgen.py +++ b/src/transformers/models/musicgen/modeling_musicgen.py @@ -36,12 +36,7 @@ LogitsProcessorList, StoppingCriteriaList, ) -from ...modeling_attn_mask_utils import ( - _prepare_4d_attention_mask, - _prepare_4d_attention_mask_for_sdpa, - _prepare_4d_causal_attention_mask, - _prepare_4d_causal_attention_mask_for_sdpa, -) +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_flash_attention_utils import ( FlashAttentionKwargs, ) @@ -55,17 +50,12 @@ ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested +from ...utils import TransformersKwargs, auto_docstring, logging from ..auto.configuration_auto import AutoConfig from ..auto.modeling_auto import AutoModel from .configuration_musicgen import MusicgenConfig, MusicgenDecoderConfig -if is_torch_flex_attn_available(): - from ...integrations.flex_attention import make_flex_block_causal_mask - - if TYPE_CHECKING: from ...generation.streamers import BaseStreamer @@ -533,9 +523,7 @@ def forward( # (bsz * codebooks, seq_len) -> (bsz, codebooks, seq_len) input = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1]) bsz, num_codebooks, seq_len = input.shape - input_shape = (bsz, seq_len) elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] input = inputs_embeds[:, :, -1:] else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") @@ -555,17 +543,24 @@ def forward( if inputs_embeds is None: inputs_embeds = sum(self.embed_tokens[codebook](input[:, codebook]) for codebook in range(num_codebooks)) - attention_mask = self._update_causal_mask( - attention_mask, - input_shape, - inputs_embeds, - past_key_values_length, + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position: torch.Tensor = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + ) + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) - encoder_attention_mask = self._update_cross_attn_mask( - encoder_hidden_states, - encoder_attention_mask, - input_shape, - inputs_embeds, + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) # embed positions @@ -623,75 +618,6 @@ def forward( cross_attentions=all_cross_attentions, ) - def _update_causal_mask( - self, - attention_mask: torch.Tensor | None, - input_shape: torch.Size, - inputs_embeds: torch.Tensor, - past_key_values_length: int, - ): - if is_flash_attention_requested(self.config): - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - elif self.config._attn_implementation == "sdpa": - attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - input_shape, - inputs_embeds, - past_key_values_length, - ) - elif self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - # Other attention flavors support in-built causal (when `mask is None`) - # while we need to create our specific block mask regardless - elif attention_mask is None: - attention_mask = make_flex_block_causal_mask( - torch.ones( - size=(input_shape), - device=inputs_embeds.device, - ) - ) - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length - ) - - return attention_mask - - def _update_cross_attn_mask( - self, - encoder_hidden_states: torch.Tensor | None, - encoder_attention_mask: torch.Tensor | None, - input_shape: torch.Size, - inputs_embeds: torch.Tensor, - ): - # expand encoder attention mask - if encoder_hidden_states is not None and encoder_attention_mask is not None: - if is_flash_attention_requested(self.config): - encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None - elif self.config._attn_implementation == "sdpa": - encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa( - encoder_attention_mask, - inputs_embeds.dtype, - tgt_len=input_shape[-1], - ) - elif self.config._attn_implementation == "flex_attention": - if isinstance(encoder_attention_mask, torch.Tensor): - encoder_attention_mask = make_flex_block_causal_mask( - encoder_attention_mask, - query_length=input_shape[-1], - is_causal=False, - ) - else: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] - ) - - return encoder_attention_mask - @auto_docstring class MusicgenModel(MusicgenPreTrainedModel): diff --git a/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py b/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py index ba68b800f517..7efbdb019019 100644 --- a/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py +++ b/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py @@ -35,10 +35,7 @@ LogitsProcessorList, StoppingCriteriaList, ) -from ...modeling_attn_mask_utils import ( - _prepare_4d_causal_attention_mask, - _prepare_4d_causal_attention_mask_for_sdpa, -) +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import ( FlashAttentionKwargs, ) @@ -46,17 +43,12 @@ from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested +from ...utils import TransformersKwargs, auto_docstring, logging from ..auto.configuration_auto import AutoConfig from ..auto.modeling_auto import AutoModel, AutoModelForTextEncoding from .configuration_musicgen_melody import MusicgenMelodyConfig, MusicgenMelodyDecoderConfig -if is_torch_flex_attn_available(): - from ...integrations.flex_attention import make_flex_block_causal_mask - - if TYPE_CHECKING: from ...generation.streamers import BaseStreamer @@ -505,9 +497,7 @@ def forward( # (bsz * codebooks, seq_len) -> (bsz, codebooks, seq_len) input = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1]) bsz, num_codebooks, seq_len = input.shape - input_shape = (bsz, seq_len) elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] input = inputs_embeds[:, :, -1:] else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") @@ -539,13 +529,18 @@ def forward( # fuse encoder_hidden_states and inputs_embeds inputs_embeds = torch.cat([encoder_hidden_states, inputs_embeds], dim=1) - input_shape = inputs_embeds.size()[:-1] + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position: torch.Tensor = ( + torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + ) - attention_mask = self._update_causal_mask( - attention_mask, - input_shape, - inputs_embeds, - past_key_values_length, + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) # embed positions @@ -594,55 +589,6 @@ def forward( attentions=all_attentions, ) - def _update_causal_mask( - self, - attention_mask: torch.Tensor | None, - input_shape: torch.Size, - inputs_embeds: torch.Tensor, - past_key_values_length: int, - ): - if is_flash_attention_requested(self.config): - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - elif self.config._attn_implementation == "sdpa": - attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - input_shape, - inputs_embeds, - past_key_values_length, - ) - elif self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - # Other attention flavors support in-built causal (when `mask is None`) - # while we need to create our specific block mask regardless - elif attention_mask is None: - attention_mask = make_flex_block_causal_mask( - torch.ones( - size=(input_shape), - device=inputs_embeds.device, - ) - ) - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length - ) - - return attention_mask - - # Ignore copy - def _update_cross_attn_mask( - self, - encoder_hidden_states: torch.Tensor | None, - encoder_attention_mask: torch.Tensor | None, - input_shape: torch.Size, - inputs_embeds: torch.Tensor, - ): - # MusicgenMelody doesn't apply cross attention, hence it's ignored here - # and only exists to not confuse any copy checks - pass - @auto_docstring # Copied from transformers.models.musicgen.modeling_musicgen.MusicgenModel with MUSICGEN->MUSICGEN_MELODY,Musicgen->MusicgenMelody diff --git a/src/transformers/models/mvp/modeling_mvp.py b/src/transformers/models/mvp/modeling_mvp.py index 15d802941e35..4d3c844e76d1 100644 --- a/src/transformers/models/mvp/modeling_mvp.py +++ b/src/transformers/models/mvp/modeling_mvp.py @@ -23,10 +23,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import ( - _prepare_4d_attention_mask, - _prepare_4d_causal_attention_mask, -) +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -603,8 +600,11 @@ def forward( # expand attention_mask if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -791,15 +791,26 @@ def forward( ) past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) # embed positions diff --git a/src/transformers/models/nemotron/modeling_nemotron.py b/src/transformers/models/nemotron/modeling_nemotron.py index 8b4f5d28e427..7b25e2166719 100644 --- a/src/transformers/models/nemotron/modeling_nemotron.py +++ b/src/transformers/models/nemotron/modeling_nemotron.py @@ -16,7 +16,7 @@ import math from collections.abc import Callable -from typing import Optional, Union +from typing import Optional import torch import torch.nn.functional as F @@ -26,7 +26,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, StaticCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask from ...modeling_layers import ( GenericForQuestionAnswering, @@ -43,17 +43,11 @@ dynamic_rope_update, ) from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils import auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast from .configuration_nemotron import NemotronConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -672,8 +666,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) # embed positions @@ -717,131 +716,6 @@ def forward( attentions=all_self_attns, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - # TODO: re-enable check: Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->NEMOTRON,Llama->Nemotron,llama->nemotron class NemotronForCausalLM(NemotronPreTrainedModel, GenerationMixin): diff --git a/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py b/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py index c0191b58dae2..a6f979bd7aa0 100644 --- a/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py +++ b/src/transformers/models/omdet_turbo/modeling_omdet_turbo.py @@ -28,7 +28,7 @@ from ...activations import ACT2CLS, ACT2FN from ...backbone_utils import load_backbone from ...integrations import use_kernel_forward_from_hub -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack @@ -1361,12 +1361,16 @@ def forward( fusion_size = attn_mask_len + task_features.shape[0] key_padding_mask = torch.zeros([batch_size, fusion_size], dtype=torch.bool).to(task_features.device) key_padding_mask[:, attn_mask_len:] = src_key_mask - attention_mask = _prepare_4d_attention_mask(~key_padding_mask, dtype=vision_features.dtype) decoder_embeddings, reference_points, encoder_bboxes, encoder_class_similarity, init_reference_points = ( self._get_decoder_input( vision_features, tuple(vision_shapes_list), class_features, denoise_embeddings, denoise_bboxes ) ) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=torch.ones_like(key_padding_mask, dtype=decoder_embeddings.dtype)[..., None], + attention_mask=~key_padding_mask, + ) all_hidden_states = () if output_hidden_states else None all_attns = () if output_attentions else None diff --git a/src/transformers/models/opt/modeling_opt.py b/src/transformers/models/opt/modeling_opt.py index c8636b012e33..35b31b6cf162 100644 --- a/src/transformers/models/opt/modeling_opt.py +++ b/src/transformers/models/opt/modeling_opt.py @@ -14,7 +14,6 @@ """PyTorch OPT model.""" from collections.abc import Callable -from typing import Union import torch from torch import nn @@ -23,7 +22,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -34,17 +33,10 @@ ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from .configuration_opt import OPTConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -348,131 +340,6 @@ def __init__(self, config: OPTConfig): # Initialize weights and apply final processing self.post_init() - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @can_return_tuple def forward( self, @@ -573,11 +440,6 @@ def forward( seq_length = past_seen_tokens + inputs_embeds.shape[1] attention_mask = torch.ones(inputs_embeds.shape[0], seq_length, device=inputs_embeds.device) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions - ) - - # embed positions if position_ids is None: # position_ids = cache_position.unsqueeze(0) position_ids = torch.cumsum(attention_mask, dim=1) @@ -585,6 +447,14 @@ def forward( # cut positions if `past_seen_tokens` is > 0 position_ids = position_ids[:, past_seen_tokens:] + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + ) + pos_embeds = self.embed_positions(attention_mask, past_seen_tokens, position_ids=position_ids) if self.project_in is not None: diff --git a/src/transformers/models/owlv2/modeling_owlv2.py b/src/transformers/models/owlv2/modeling_owlv2.py index 3079bb35054f..4ca64b240e57 100644 --- a/src/transformers/models/owlv2/modeling_owlv2.py +++ b/src/transformers/models/owlv2/modeling_owlv2.py @@ -21,7 +21,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel @@ -414,8 +414,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = False, + **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: """Input shape: Batch x Time x Channel""" @@ -440,16 +440,6 @@ def forward( f" {attn_weights.size()}" ) - # apply the causal_attention_mask first - if causal_attention_mask is not None: - if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is" - f" {causal_attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - if attention_mask is not None: if attention_mask.size() != (bsz, 1, tgt_len, src_len): raise ValueError( @@ -522,8 +512,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -541,8 +531,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -633,10 +623,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutput: r""" Args: @@ -646,11 +636,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. @@ -676,8 +661,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] @@ -696,15 +681,18 @@ def forward( # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTTextTransformer with OWLVIT->OWLV2,OwlViT->Owlv2 -class Owlv2TextTransformer(nn.Module): +class Owlv2TextTransformer(Owlv2PreTrainedModel): def __init__(self, config: Owlv2TextConfig): - super().__init__() - self.config = config + super().__init__(config) + embed_dim = config.hidden_size self.embeddings = Owlv2TextEmbeddings(config) self.encoder = Owlv2Encoder(config) self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + # Initialize weights and apply final processing + self.post_init() + @auto_docstring def forward( self, @@ -714,6 +702,7 @@ def forward( output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutputWithPooling: r""" input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`): @@ -731,24 +720,23 @@ def forward( input_ids = input_ids.view(-1, input_shape[-1]) hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) - # num_samples, seq_len = input_shape where num_samples = batch_size * num_max_text_queries - # OWLV2's text model uses causal mask, prepare it here. - # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 - causal_attention_mask = _create_4d_causal_attention_mask( - input_shape, hidden_states.dtype, device=hidden_states.device + attention_mask = create_causal_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + cache_position=torch.arange(hidden_states.shape[1], device=hidden_states.device), + past_key_values=None, ) - # expand attention_mask - if attention_mask is not None: - # [num_samples, seq_len] -> [num_samples, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + kwargs.pop("is_causal", None) encoder_outputs = self.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + is_causal=True, + **kwargs, ) last_hidden_state = encoder_outputs[0] @@ -830,16 +818,18 @@ def forward( # Copied from transformers.models.owlvit.modeling_owlvit.OwlViTVisionTransformer with OWLVIT->OWLV2,OwlViT->Owlv2 -class Owlv2VisionTransformer(nn.Module): +class Owlv2VisionTransformer(Owlv2PreTrainedModel): def __init__(self, config: Owlv2VisionConfig): - super().__init__() - self.config = config + super().__init__(config) self.embeddings = Owlv2VisionEmbeddings(config) self.pre_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.encoder = Owlv2Encoder(config) self.post_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + # Initialize weights and apply final processing + self.post_init() + @auto_docstring def forward( self, @@ -848,6 +838,7 @@ def forward( output_hidden_states: bool | None = None, interpolate_pos_encoding: bool | None = False, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutputWithPooling: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -867,6 +858,7 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + **kwargs, ) last_hidden_state = encoder_outputs[0] diff --git a/src/transformers/models/owlvit/modeling_owlvit.py b/src/transformers/models/owlvit/modeling_owlvit.py index 4d2c96f793be..28d4a83bffee 100644 --- a/src/transformers/models/owlvit/modeling_owlvit.py +++ b/src/transformers/models/owlvit/modeling_owlvit.py @@ -21,7 +21,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel @@ -402,8 +402,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = False, + **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: """Input shape: Batch x Time x Channel""" @@ -428,16 +428,6 @@ def forward( f" {attn_weights.size()}" ) - # apply the causal_attention_mask first - if causal_attention_mask is not None: - if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is" - f" {causal_attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - if attention_mask is not None: if attention_mask.size() != (bsz, 1, tgt_len, src_len): raise ValueError( @@ -510,8 +500,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -529,8 +519,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -619,10 +609,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutput: r""" Args: @@ -632,11 +622,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. @@ -662,8 +647,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] @@ -681,15 +666,18 @@ def forward( ) -class OwlViTTextTransformer(nn.Module): +class OwlViTTextTransformer(OwlViTPreTrainedModel): def __init__(self, config: OwlViTTextConfig): - super().__init__() - self.config = config + super().__init__(config) + embed_dim = config.hidden_size self.embeddings = OwlViTTextEmbeddings(config) self.encoder = OwlViTEncoder(config) self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + # Initialize weights and apply final processing + self.post_init() + @auto_docstring def forward( self, @@ -699,6 +687,7 @@ def forward( output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutputWithPooling: r""" input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`): @@ -716,24 +705,23 @@ def forward( input_ids = input_ids.view(-1, input_shape[-1]) hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) - # num_samples, seq_len = input_shape where num_samples = batch_size * num_max_text_queries - # OWLVIT's text model uses causal mask, prepare it here. - # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 - causal_attention_mask = _create_4d_causal_attention_mask( - input_shape, hidden_states.dtype, device=hidden_states.device + attention_mask = create_causal_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + cache_position=torch.arange(hidden_states.shape[1], device=hidden_states.device), + past_key_values=None, ) - # expand attention_mask - if attention_mask is not None: - # [num_samples, seq_len] -> [num_samples, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + kwargs.pop("is_causal", None) encoder_outputs = self.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + is_causal=True, + **kwargs, ) last_hidden_state = encoder_outputs[0] @@ -813,16 +801,18 @@ def forward( ) -class OwlViTVisionTransformer(nn.Module): +class OwlViTVisionTransformer(OwlViTPreTrainedModel): def __init__(self, config: OwlViTVisionConfig): - super().__init__() - self.config = config + super().__init__(config) self.embeddings = OwlViTVisionEmbeddings(config) self.pre_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.encoder = OwlViTEncoder(config) self.post_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + # Initialize weights and apply final processing + self.post_init() + @auto_docstring def forward( self, @@ -831,6 +821,7 @@ def forward( output_hidden_states: bool | None = None, interpolate_pos_encoding: bool | None = False, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutputWithPooling: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -850,6 +841,7 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + **kwargs, ) last_hidden_state = encoder_outputs[0] diff --git a/src/transformers/models/pe_audio/modeling_pe_audio.py b/src/transformers/models/pe_audio/modeling_pe_audio.py index 9c14a046abfa..92fa76547e5b 100644 --- a/src/transformers/models/pe_audio/modeling_pe_audio.py +++ b/src/transformers/models/pe_audio/modeling_pe_audio.py @@ -31,7 +31,7 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -647,7 +647,11 @@ def forward( inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) if attention_mask is not None: - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) position_embeddings = self.rotary_emb(inputs_embeds, position_ids) diff --git a/src/transformers/models/pe_audio/modular_pe_audio.py b/src/transformers/models/pe_audio/modular_pe_audio.py index 7455fed105bb..0ee1d76706a1 100644 --- a/src/transformers/models/pe_audio/modular_pe_audio.py +++ b/src/transformers/models/pe_audio/modular_pe_audio.py @@ -20,7 +20,7 @@ from ... import initialization as init from ...configuration_utils import PreTrainedConfig -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput from ...utils import ModelOutput, auto_docstring, can_return_tuple from ...utils.generic import check_model_inputs @@ -122,7 +122,11 @@ def forward( inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) if attention_mask is not None: - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) position_embeddings = self.rotary_emb(inputs_embeds, position_ids) diff --git a/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py b/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py index b64d9ec45bdc..18edabdd91a3 100644 --- a/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py +++ b/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py @@ -28,7 +28,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -601,7 +601,11 @@ def forward( inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) if attention_mask is not None: - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) position_embeddings = self.rotary_emb(inputs_embeds, position_ids) diff --git a/src/transformers/models/pe_audio_video/modular_pe_audio_video.py b/src/transformers/models/pe_audio_video/modular_pe_audio_video.py index c5584e9aa617..23c4bdb53048 100644 --- a/src/transformers/models/pe_audio_video/modular_pe_audio_video.py +++ b/src/transformers/models/pe_audio_video/modular_pe_audio_video.py @@ -19,7 +19,7 @@ import torch.nn as nn from ... import initialization as init -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel, eager_attention_forward from ...processing_utils import Unpack @@ -390,7 +390,11 @@ def forward( inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) if attention_mask is not None: - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) position_embeddings = self.rotary_emb(inputs_embeds, position_ids) diff --git a/src/transformers/models/pe_video/modeling_pe_video.py b/src/transformers/models/pe_video/modeling_pe_video.py index 4c32deecb300..5fb68e85de21 100644 --- a/src/transformers/models/pe_video/modeling_pe_video.py +++ b/src/transformers/models/pe_video/modeling_pe_video.py @@ -29,7 +29,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...integrations import use_kernel_forward_from_hub, use_kernelized_func -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -531,7 +531,11 @@ def forward( inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) if attention_mask is not None: - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) position_embeddings = self.rotary_emb(inputs_embeds, position_ids) diff --git a/src/transformers/models/pe_video/modular_pe_video.py b/src/transformers/models/pe_video/modular_pe_video.py index 38576598cf44..c79d698ee1de 100644 --- a/src/transformers/models/pe_video/modular_pe_video.py +++ b/src/transformers/models/pe_video/modular_pe_video.py @@ -18,7 +18,7 @@ import torch.nn as nn import torch.nn.functional as F -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple @@ -112,7 +112,11 @@ def forward( inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) if attention_mask is not None: - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) position_embeddings = self.rotary_emb(inputs_embeds, position_ids) diff --git a/src/transformers/models/persimmon/modeling_persimmon.py b/src/transformers/models/persimmon/modeling_persimmon.py index b2d2718bb25a..655080a416ef 100644 --- a/src/transformers/models/persimmon/modeling_persimmon.py +++ b/src/transformers/models/persimmon/modeling_persimmon.py @@ -19,7 +19,7 @@ """PyTorch Persimmon model.""" from collections.abc import Callable -from typing import Optional, Union +from typing import Optional import torch from torch import nn @@ -27,7 +27,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import ( GenericForSequenceClassification, @@ -44,17 +44,11 @@ ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils import auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast from .configuration_persimmon import PersimmonConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -496,8 +490,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) hidden_states = inputs_embeds @@ -541,131 +540,6 @@ def forward( attentions=all_self_attns, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - class PersimmonForCausalLM(PersimmonPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index 2a51c5fa7c69..d2aec9613faf 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -32,8 +32,7 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub -from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( @@ -46,7 +45,7 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, can_return_tuple, torch_int -from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested, maybe_autocast +from ...utils.generic import TransformersKwargs, check_model_inputs, maybe_autocast from .configuration_phi4_multimodal import Phi4MultimodalAudioConfig, Phi4MultimodalConfig, Phi4MultimodalVisionConfig @@ -432,17 +431,11 @@ def forward( hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) patch_attention_mask = patch_attention_mask.view(batch_size, -1) - # The call to `_upad_input` in `_flash_attention_forward` is expensive - # So when the `patch_attention_mask` is full of 1s (i.e. attending to the whole sequence), - # avoiding passing the attention_mask, which is equivalent to attending to the full sequence - if not torch.any(~patch_attention_mask): - attention_mask = None - else: - attention_mask = ( - _prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype) - if not is_flash_attention_requested(self.config) - else patch_attention_mask - ) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=patch_attention_mask, + ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, diff --git a/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py index 27e0e88348c9..bedd780d45ae 100644 --- a/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py @@ -24,8 +24,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig -from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import ( BaseModelOutput, BaseModelOutputWithPast, @@ -36,7 +35,7 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, logging -from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested, maybe_autocast +from ...utils.generic import TransformersKwargs, check_model_inputs, maybe_autocast from ..phi3.configuration_phi3 import Phi3Config from ..phi3.modeling_phi3 import ( Phi3DecoderLayer, @@ -702,17 +701,11 @@ def forward( hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) patch_attention_mask = patch_attention_mask.view(batch_size, -1) - # The call to `_upad_input` in `_flash_attention_forward` is expensive - # So when the `patch_attention_mask` is full of 1s (i.e. attending to the whole sequence), - # avoiding passing the attention_mask, which is equivalent to attending to the full sequence - if not torch.any(~patch_attention_mask): - attention_mask = None - else: - attention_mask = ( - _prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype) - if not is_flash_attention_requested(self.config) - else patch_attention_mask - ) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=patch_attention_mask, + ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, diff --git a/src/transformers/models/pix2struct/modeling_pix2struct.py b/src/transformers/models/pix2struct/modeling_pix2struct.py index 2d73fadd90ae..f7e7305c974b 100644 --- a/src/transformers/models/pix2struct/modeling_pix2struct.py +++ b/src/transformers/models/pix2struct/modeling_pix2struct.py @@ -14,7 +14,6 @@ """Pix2Struct modeling file""" import math -from typing import Union import torch from torch import nn @@ -23,7 +22,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -37,20 +36,12 @@ DUMMY_INPUTS, DUMMY_MASK, auto_docstring, - is_torch_flex_attn_available, is_torchdynamo_compiling, logging, ) -from ...utils.generic import is_flash_attention_requested from .configuration_pix2struct import Pix2StructConfig, Pix2StructTextConfig, Pix2StructVisionConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) # General docstring @@ -1083,14 +1074,12 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache - if isinstance(past_key_values, EncoderDecoderCache) - else past_key_values, - output_attentions, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) else: causal_mask = attention_mask[:, None, None, :] @@ -1186,131 +1175,6 @@ def forward( cross_attentions=all_cross_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring( custom_intro=""" diff --git a/src/transformers/models/pop2piano/modeling_pop2piano.py b/src/transformers/models/pop2piano/modeling_pop2piano.py index 5f1e65eca8c8..a98681af2411 100644 --- a/src/transformers/models/pop2piano/modeling_pop2piano.py +++ b/src/transformers/models/pop2piano/modeling_pop2piano.py @@ -15,7 +15,6 @@ import copy import math -from typing import Union import torch from torch import nn @@ -27,21 +26,14 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import 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_torch_flex_attn_available, is_torchdynamo_compiling, logging -from ...utils.generic import is_flash_attention_requested +from ...utils import auto_docstring, is_torchdynamo_compiling, logging from .configuration_pop2piano import Pop2PianoConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) _load_pop2piano_layer_norm = True @@ -710,14 +702,12 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache - if isinstance(past_key_values, EncoderDecoderCache) - else past_key_values, - output_attentions, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) else: causal_mask = attention_mask[:, None, None, :] @@ -801,131 +791,6 @@ def forward( cross_attentions=all_cross_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - class Pop2PianoConcatEmbeddingToMel(nn.Module): """Embedding Matrix for `composer` tokens.""" diff --git a/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py index 7818b37ae10f..733aad22a813 100644 --- a/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py @@ -149,61 +149,6 @@ def _init_weights(self, module): class Qwen2_5OmniPreTrainedModelForConditionalGeneration(Qwen2_5OmniPreTrainedModel): input_modalities = ("image", "video", "audio", "text") - def _prepare_4d_causal_attention_mask_with_cache_position( - self, - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - device: torch.device, - min_dtype: float, - cache_position: torch.Tensor, - batch_size: int, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - device (`torch.device`): - The device to place the 4D attention mask on. - min_dtype (`float`): - The minimum value representable with the dtype `dtype`. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - def get_llm_pos_ids_for_vision( self, start_idx: int, diff --git a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py index dba41dde7b0d..c8b70ed0807e 100644 --- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py @@ -1090,61 +1090,6 @@ def _init_weights(self, module): class Qwen2_5OmniPreTrainedModelForConditionalGeneration(Qwen2_5OmniPreTrainedModel): input_modalities = ("image", "video", "audio", "text") - def _prepare_4d_causal_attention_mask_with_cache_position( - self, - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - device: torch.device, - min_dtype: float, - cache_position: torch.Tensor, - batch_size: int, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - device (`torch.device`): - The device to place the 4D attention mask on. - min_dtype (`float`): - The minimum value representable with the dtype `dtype`. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - def get_llm_pos_ids_for_vision( self, start_idx: int, diff --git a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py index 4e07f1b3863a..e125819a9562 100644 --- a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py @@ -151,61 +151,6 @@ def _get_feat_extract_output_lengths(input_lengths): class Qwen3OmniMoePreTrainedModelForConditionalGeneration(Qwen3OmniMoePreTrainedModel): input_modalities = ("image", "video", "audio", "text") - def _prepare_4d_causal_attention_mask_with_cache_position( - self, - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - device: torch.device, - min_dtype: float, - cache_position: torch.Tensor, - batch_size: int, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - device (`torch.device`): - The device to place the 4D attention mask on. - min_dtype (`float`): - The minimum value representable with the dtype `dtype`. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - def get_llm_pos_ids_for_vision( self, start_idx: int, diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py index 424287e47656..c1840b3dec82 100644 --- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -24,7 +24,6 @@ from ... import initialization as init from ...activations import ACT2FN from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithNoAttention, CausalLMOutput from ...modeling_rope_utils import dynamic_rope_update @@ -707,7 +706,6 @@ def forward( hidden_states=all_hidden_states, ) - # Ignore copy def _update_causal_mask(self, attention_mask, input_tensor, cache_position): dtype, device = input_tensor.dtype, input_tensor.device min_dtype = torch.finfo(dtype).min @@ -734,7 +732,7 @@ def _update_causal_mask(self, attention_mask, input_tensor, cache_position): # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + causal_mask = causal_mask.mul(~torch.all(causal_mask == min_dtype, dim=-1, keepdim=True)) return causal_mask diff --git a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py index be82aca79b78..50c9f430d1f7 100755 --- a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py +++ b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py @@ -27,10 +27,7 @@ from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled from ...integrations.fsdp import is_fsdp_managed_module -from ...modeling_attn_mask_utils import ( - _prepare_4d_attention_mask, - _prepare_4d_causal_attention_mask, -) +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -754,6 +751,8 @@ def forward( class SeamlessM4TConformerAdapterLayer(nn.Module): def __init__(self, config): super().__init__() + self.config = config + embed_dim = config.hidden_size dropout = config.adaptor_dropout @@ -826,9 +825,10 @@ def forward( hidden_states.device ) attention_mask = _compute_new_attention_mask(hidden_states=hidden_states, seq_lens=sub_sampled_lengths) - attention_mask = _prepare_4d_attention_mask( - attention_mask, - hidden_states.dtype, + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, ) # The rest of the computation is identical to a vanilla Transformer @@ -1671,8 +1671,11 @@ def forward( # expand attention_mask if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -1821,15 +1824,26 @@ def forward( past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) # embed positions diff --git a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py index c8dc99af7ccf..8af4bad722eb 100644 --- a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py +++ b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py @@ -27,7 +27,7 @@ from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled from ...integrations.fsdp import is_fsdp_managed_module -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -623,6 +623,8 @@ def forward( class SeamlessM4Tv2ConformerAdapterLayer(nn.Module): def __init__(self, config): super().__init__() + self.config = config + embed_dim = config.hidden_size dropout = config.adaptor_dropout @@ -695,9 +697,10 @@ def forward( hidden_states.device ) attention_mask = _compute_new_attention_mask(hidden_states=hidden_states, seq_lens=sub_sampled_lengths) - attention_mask = _prepare_4d_attention_mask( - attention_mask, - hidden_states.dtype, + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, ) # The rest of the computation is identical to a vanilla Transformer @@ -1702,8 +1705,11 @@ def forward( # expand attention_mask if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -1853,15 +1859,26 @@ def forward( past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) # embed positions @@ -1896,7 +1913,6 @@ def forward( cache_position=cache_position, ) hidden_states = layer_outputs[0] - if output_attentions: all_self_attns += (layer_outputs[1],) @@ -2049,7 +2065,11 @@ def forward( hidden_states = char_hidden_states + positions padding_mask = _compute_new_attention_mask(hidden_states, dur_out.sum(1)) - attention_mask = _prepare_4d_attention_mask(padding_mask, hidden_states.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=padding_mask, + ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) diff --git a/src/transformers/models/siglip/modeling_siglip.py b/src/transformers/models/siglip/modeling_siglip.py index 97aaa620c93d..f68e8189443c 100644 --- a/src/transformers/models/siglip/modeling_siglip.py +++ b/src/transformers/models/siglip/modeling_siglip.py @@ -23,7 +23,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -35,7 +35,7 @@ can_return_tuple, torch_int, ) -from ...utils.generic import check_model_inputs, is_flash_attention_requested +from ...utils.generic import check_model_inputs from .configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig @@ -504,13 +504,11 @@ def forward( hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) # note: SigLIP's text model does not use a causal mask, unlike the original CLIP model. - # expand attention_mask - uses_flash_attention = is_flash_attention_requested(self.config) - if uses_flash_attention: - attention_mask = None - elif attention_mask is not None and not uses_flash_attention: - # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, diff --git a/src/transformers/models/siglip2/modeling_siglip2.py b/src/transformers/models/siglip2/modeling_siglip2.py index f73e432e32b7..6852cb264053 100644 --- a/src/transformers/models/siglip2/modeling_siglip2.py +++ b/src/transformers/models/siglip2/modeling_siglip2.py @@ -29,13 +29,13 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check -from ...utils.generic import check_model_inputs, is_flash_attention_requested +from ...utils.generic import check_model_inputs from .configuration_siglip2 import Siglip2Config, Siglip2TextConfig, Siglip2VisionConfig @@ -403,9 +403,10 @@ class Siglip2PreTrainedModel(PreTrainedModel): "Siglip2EncoderLayer", "Siglip2MultiheadAttentionPoolingHead", ] - _supports_flash_attn = True + _supports_flash_attn = False _supports_sdpa = True - _supports_flex_attn = True + # nn.MultiHeadAttention mask doesn't allow for non 4d mask + _supports_flex_attn = False _supports_attention_backend = True _can_record_outputs = { @@ -536,11 +537,11 @@ def forward( hidden_states = self.embeddings(pixel_values, spatial_shapes) - if attention_mask is not None and not is_flash_attention_requested(self.config): - # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) - else: - encoder_attention_mask = attention_mask + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, @@ -594,13 +595,11 @@ def forward( hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) # note: Siglip2's text model does not use a causal mask, unlike the original CLIP model. - # expand attention_mask - uses_flash_attention = is_flash_attention_requested(self.config) - if uses_flash_attention: - attention_mask = None - elif attention_mask is not None and not uses_flash_attention: - # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, @@ -686,6 +685,8 @@ def __init__(self, config: Siglip2VisionConfig): self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.mlp = Siglip2MLP(config) + + self.config = config self.num_heads = config.num_attention_heads def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: @@ -694,9 +695,23 @@ def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | Non if attention_mask is not None: target_len, source_len = probe.shape[1], hidden_state.shape[1] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len) - attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1) - attention_mask = attention_mask.reshape(-1, target_len, source_len) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=probe, + attention_mask=attention_mask, + encoder_hidden_states=hidden_state, + ) + if attention_mask is not None: + attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1) + attention_mask = attention_mask.reshape(-1, target_len, source_len) + + # `nn.MultiheadAttention` cannot handle boolean masks (which SDPA can) + if attention_mask.dtype == torch.bool: + attention_mask = torch.where( + attention_mask, + torch.tensor(0.0, device=attention_mask.device, dtype=probe.dtype), + torch.finfo(probe.dtype).min, + ) hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0] diff --git a/src/transformers/models/siglip2/modular_siglip2.py b/src/transformers/models/siglip2/modular_siglip2.py index 5bf68682ffee..02bffac690b2 100644 --- a/src/transformers/models/siglip2/modular_siglip2.py +++ b/src/transformers/models/siglip2/modular_siglip2.py @@ -35,7 +35,7 @@ SiglipVisionTransformer, ) -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...processing_utils import Unpack from ...utils import ( TransformersKwargs, @@ -43,7 +43,7 @@ can_return_tuple, torch_compilable_check, ) -from ...utils.generic import check_model_inputs, is_flash_attention_requested +from ...utils.generic import check_model_inputs class Siglip2Tokenizer(GemmaTokenizer): @@ -280,7 +280,9 @@ def forward(self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTen class Siglip2PreTrainedModel(SiglipPreTrainedModel): - pass + # nn.MultiHeadAttention mask doesn't allow for non 4d mask + _supports_flex_attn = False + _supports_flash_attn = False class Siglip2VisionTransformer(SiglipVisionTransformer): @@ -308,11 +310,11 @@ def forward( hidden_states = self.embeddings(pixel_values, spatial_shapes) - if attention_mask is not None and not is_flash_attention_requested(self.config): - # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) - else: - encoder_attention_mask = attention_mask + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, @@ -341,6 +343,8 @@ class Siglip2TextModel(SiglipTextModel): class Siglip2MultiheadAttentionPoolingHead(SiglipMultiheadAttentionPoolingHead): def __init__(self, config: Siglip2VisionConfig): super().__init__(config) + + self.config = config self.num_heads = config.num_attention_heads def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: @@ -349,9 +353,23 @@ def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | Non if attention_mask is not None: target_len, source_len = probe.shape[1], hidden_state.shape[1] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len) - attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1) - attention_mask = attention_mask.reshape(-1, target_len, source_len) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=probe, + attention_mask=attention_mask, + encoder_hidden_states=hidden_state, + ) + if attention_mask is not None: + attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1) + attention_mask = attention_mask.reshape(-1, target_len, source_len) + + # `nn.MultiheadAttention` cannot handle boolean masks (which SDPA can) + if attention_mask.dtype == torch.bool: + attention_mask = torch.where( + attention_mask, + torch.tensor(0.0, device=attention_mask.device, dtype=probe.dtype), + torch.finfo(probe.dtype).min, + ) hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0] diff --git a/src/transformers/models/speecht5/modeling_speecht5.py b/src/transformers/models/speecht5/modeling_speecht5.py index 1f734303de50..56ad3aa3ee93 100644 --- a/src/transformers/models/speecht5/modeling_speecht5.py +++ b/src/transformers/models/speecht5/modeling_speecht5.py @@ -26,7 +26,7 @@ from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled from ...integrations.fsdp import is_fsdp_managed_module -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -1288,10 +1288,11 @@ def forward( ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) hidden_states = self.layer_norm(hidden_states) hidden_states = self.dropout(hidden_states) @@ -1531,8 +1532,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.use_return_dict - input_shape = hidden_states.size()[:-1] - if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( @@ -1544,16 +1543,26 @@ def forward( past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + hidden_states.shape[1], device=hidden_states.device + ) - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, hidden_states, past_key_values_length + attention_mask = create_causal_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, hidden_states.dtype, tgt_len=input_shape[-1] + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self) diff --git a/src/transformers/models/splinter/modeling_splinter.py b/src/transformers/models/splinter/modeling_splinter.py index 48f39c614dc7..baa3c34f9161 100755 --- a/src/transformers/models/splinter/modeling_splinter.py +++ b/src/transformers/models/splinter/modeling_splinter.py @@ -25,8 +25,10 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, ModelOutput, QuestionAnsweringModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward from ...utils import ( + TransformersKwargs, auto_docstring, can_return_tuple, logging, @@ -140,7 +142,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.attention_head_size) @@ -196,7 +198,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_outputs = self.self( hidden_states, @@ -255,7 +257,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.FloatTensor | None = None, output_attentions: bool | None = False, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor]: self_attention_outputs = self.attention( hidden_states, @@ -295,7 +297,7 @@ def forward( output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | BaseModelOutput: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None diff --git a/src/transformers/models/stablelm/modeling_stablelm.py b/src/transformers/models/stablelm/modeling_stablelm.py index dcdf29e1b2c0..dd804b1cbe28 100755 --- a/src/transformers/models/stablelm/modeling_stablelm.py +++ b/src/transformers/models/stablelm/modeling_stablelm.py @@ -19,7 +19,7 @@ """PyTorch StableLM model.""" from collections.abc import Callable -from typing import Optional, Union +from typing import Optional import torch from torch import nn @@ -27,7 +27,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import ( GenericForSequenceClassification, GenericForTokenClassification, @@ -43,17 +43,11 @@ ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import maybe_autocast from .configuration_stablelm import StableLmConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -522,8 +516,13 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) hidden_states = inputs_embeds @@ -566,131 +565,6 @@ def forward( attentions=all_self_attns, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - # Copied from transformers.models.persimmon.modeling_persimmon.PersimmonForCausalLM with PERSIMMON->STABLELM,Persimmon->StableLm class StableLmForCausalLM(StableLmPreTrainedModel, GenerationMixin): diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index 32add26f8d4c..9955074ae3cb 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -20,7 +20,6 @@ import copy import math -from typing import Union import torch import torch.nn as nn @@ -30,7 +29,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( MoEModelOutput, @@ -40,24 +39,12 @@ ) from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import ( - TransformersKwargs, - auto_docstring, - is_torch_flex_attn_available, - is_torchdynamo_compiling, - logging, -) -from ...utils.generic import can_return_tuple, check_model_inputs, is_flash_attention_requested +from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling, logging +from ...utils.generic import can_return_tuple, check_model_inputs from ...utils.output_capturing import OutputRecorder from .configuration_switch_transformers import SwitchTransformersConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -732,13 +719,12 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache - if isinstance(past_key_values, EncoderDecoderCache) - else past_key_values, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) else: causal_mask = attention_mask[:, None, None, :] @@ -783,129 +769,6 @@ def forward( past_key_values=past_key_values, ) - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring class SwitchTransformersModel(SwitchTransformersPreTrainedModel): diff --git a/src/transformers/models/switch_transformers/modular_switch_transformers.py b/src/transformers/models/switch_transformers/modular_switch_transformers.py index 45b4e576f715..e5a7be13473e 100644 --- a/src/transformers/models/switch_transformers/modular_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modular_switch_transformers.py @@ -14,7 +14,6 @@ """PyTorch SwitchTransformers model.""" import copy -from typing import Union import torch import torch.nn as nn @@ -23,7 +22,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( MoEModelOutput, @@ -36,22 +35,15 @@ from ...utils import ( TransformersKwargs, auto_docstring, - is_torch_flex_attn_available, is_torchdynamo_compiling, logging, ) -from ...utils.generic import can_return_tuple, check_model_inputs, is_flash_attention_requested +from ...utils.generic import can_return_tuple, check_model_inputs from ...utils.output_capturing import OutputRecorder from ..t5.modeling_t5 import T5Attention, T5DenseActDense, T5LayerCrossAttention, T5LayerNorm, T5LayerSelfAttention from .configuration_switch_transformers import SwitchTransformersConfig -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - - logger = logging.get_logger(__name__) @@ -488,13 +480,12 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache - if isinstance(past_key_values, EncoderDecoderCache) - else past_key_values, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) else: causal_mask = attention_mask[:, None, None, :] @@ -539,129 +530,6 @@ def forward( past_key_values=past_key_values, ) - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring class SwitchTransformersModel(SwitchTransformersPreTrainedModel): diff --git a/src/transformers/models/table_transformer/modeling_table_transformer.py b/src/transformers/models/table_transformer/modeling_table_transformer.py index 725f109555b5..33213a2443b4 100644 --- a/src/transformers/models/table_transformer/modeling_table_transformer.py +++ b/src/transformers/models/table_transformer/modeling_table_transformer.py @@ -22,7 +22,7 @@ from ... import initialization as init from ...activations import ACT2FN from ...backbone_utils import load_backbone -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions, Seq2SeqModelOutput from ...modeling_utils import PreTrainedModel @@ -759,8 +759,11 @@ def forward( # expand attention_mask if attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -888,21 +891,21 @@ def forward( if inputs_embeds is not None: hidden_states = inputs_embeds - input_shape = inputs_embeds.size()[:-1] - - combined_attention_mask = None - if attention_mask is not None and combined_attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - combined_attention_mask = combined_attention_mask + _prepare_4d_attention_mask( - attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + if attention_mask is not None: + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, ) # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) # optional intermediate hidden states @@ -924,7 +927,7 @@ def forward( layer_outputs = decoder_layer( hidden_states, - combined_attention_mask, + attention_mask, object_queries, query_position_embeddings, encoder_hidden_states, # as a positional argument for gradient checkpointing diff --git a/src/transformers/models/trocr/modeling_trocr.py b/src/transformers/models/trocr/modeling_trocr.py index 832655984da5..db81457ee594 100644 --- a/src/transformers/models/trocr/modeling_trocr.py +++ b/src/transformers/models/trocr/modeling_trocr.py @@ -22,10 +22,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import ( - _prepare_4d_attention_mask, - _prepare_4d_causal_attention_mask, -) +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions from ...modeling_utils import PreTrainedModel @@ -521,7 +518,6 @@ def forward( input = input_ids input_ids = input_ids.view(-1, input.shape[-1]) elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] input = inputs_embeds[:, :, -1] else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") @@ -557,17 +553,26 @@ def forward( hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - input_shape = input.shape + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + inputs_embeds.shape[1], device=inputs_embeds.device + ) - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) # decoder layers diff --git a/src/transformers/models/udop/modeling_udop.py b/src/transformers/models/udop/modeling_udop.py index ff3b591947a1..5285793c3084 100644 --- a/src/transformers/models/udop/modeling_udop.py +++ b/src/transformers/models/udop/modeling_udop.py @@ -21,7 +21,7 @@ from collections.abc import Sequence from copy import deepcopy from dataclasses import dataclass -from typing import Any, Union +from typing import Any import torch from torch import Tensor, nn @@ -37,22 +37,14 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_utils import PreTrainedModel from ...utils import ( ModelOutput, auto_docstring, - is_torch_flex_attn_available, is_torchdynamo_compiling, ) -from ...utils.generic import is_flash_attention_requested - - -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask logger = logging.getLogger(__name__) @@ -1196,14 +1188,12 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache - if isinstance(past_key_values, EncoderDecoderCache) - else past_key_values, - output_attentions, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) else: causal_mask = attention_mask[:, None, None, :] @@ -1292,131 +1282,6 @@ def forward( cross_attentions=all_cross_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring class UdopModel(UdopPreTrainedModel): diff --git a/src/transformers/models/umt5/modeling_umt5.py b/src/transformers/models/umt5/modeling_umt5.py index a50ba757db39..4a721829ef83 100644 --- a/src/transformers/models/umt5/modeling_umt5.py +++ b/src/transformers/models/umt5/modeling_umt5.py @@ -15,7 +15,6 @@ import copy import math -from typing import Union import torch from torch import nn @@ -25,7 +24,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -41,20 +40,13 @@ DUMMY_INPUTS, DUMMY_MASK, auto_docstring, - is_torch_flex_attn_available, is_torchdynamo_compiling, logging, torch_compilable_check, ) -from ...utils.generic import is_flash_attention_requested from .configuration_umt5 import UMT5Config -if is_torch_flex_attn_available(): - from torch.nn.attention.flex_attention import BlockMask - - from ...integrations.flex_attention import make_flex_block_causal_mask - logger = logging.get_logger(__name__) @@ -689,14 +681,12 @@ def forward( attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.is_decoder: - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache - if isinstance(past_key_values, EncoderDecoderCache) - else past_key_values, - output_attentions, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) elif attention_mask is not None: causal_mask = attention_mask[:, None, None, :] @@ -771,131 +761,6 @@ def forward( cross_attentions=all_cross_attentions, ) - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask - def _update_causal_mask( - self, - attention_mask: Union[torch.Tensor, "BlockMask"], - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool = False, - ): - if is_flash_attention_requested(self.config): - if attention_mask is not None and (attention_mask == 0.0).any(): - return attention_mask - return None - if self.config._attn_implementation == "flex_attention": - if isinstance(attention_mask, torch.Tensor): - attention_mask = make_flex_block_causal_mask(attention_mask) - return attention_mask - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype = input_tensor.dtype - sequence_length = input_tensor.shape[1] - if using_compilable_cache: - target_length = past_key_values.get_max_cache_shape() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - min_dtype = torch.finfo(dtype).min - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - @auto_docstring class UMT5Model(UMT5PreTrainedModel): diff --git a/src/transformers/models/video_llava/modeling_video_llava.py b/src/transformers/models/video_llava/modeling_video_llava.py index 629a0fd1ab3a..d543ba97ac95 100644 --- a/src/transformers/models/video_llava/modeling_video_llava.py +++ b/src/transformers/models/video_llava/modeling_video_llava.py @@ -623,61 +623,5 @@ def prepare_inputs_for_generation( return model_inputs - @staticmethod - # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position - def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - cache_position: torch.Tensor, - batch_size: int, - **kwargs, - ): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape - `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, - to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - min_dtype = torch.finfo(dtype).min - causal_mask = torch.full( - (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - __all__ = ["VideoLlavaPreTrainedModel", "VideoLlavaModel", "VideoLlavaForConditionalGeneration"] diff --git a/src/transformers/models/vits/modeling_vits.py b/src/transformers/models/vits/modeling_vits.py index f40409b39d8e..6626ed94b3b6 100644 --- a/src/transformers/models/vits/modeling_vits.py +++ b/src/transformers/models/vits/modeling_vits.py @@ -25,7 +25,7 @@ from ...activations import ACT2FN from ...integrations.deepspeed import is_deepspeed_zero3_enabled from ...integrations.fsdp import is_fsdp_managed_module -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, ModelOutput from ...modeling_utils import PreTrainedModel @@ -1098,10 +1098,11 @@ def forward( all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + ) hidden_states = hidden_states * padding_mask diff --git a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py index 0c65e8880623..518a8c18431f 100644 --- a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py +++ b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py @@ -16,7 +16,7 @@ from ...activations import ACT2FN from ...integrations.deepspeed import is_deepspeed_zero3_enabled from ...integrations.fsdp import is_fsdp_managed_module -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -617,6 +617,8 @@ def _compute_new_attention_mask(hidden_states: torch.Tensor, seq_lens: torch.Ten class Wav2Vec2BertAdapterLayer(nn.Module): def __init__(self, config): super().__init__() + self.config = config + embed_dim = config.output_hidden_size dropout = config.conformer_conv_dropout @@ -679,9 +681,10 @@ def forward( if attention_mask is not None: attention_mask = _compute_new_attention_mask(hidden_states=hidden_states, seq_lens=sub_sampled_lengths) - attention_mask = _prepare_4d_attention_mask( - attention_mask, - hidden_states.dtype, + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, ) # The rest of the computation is identical to a vanilla Transformer diff --git a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py index c4b856f70c67..005ddedf2594 100644 --- a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py +++ b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py @@ -8,7 +8,7 @@ from ...activations import ACT2FN from ...integrations.deepspeed import is_deepspeed_zero3_enabled from ...integrations.fsdp import is_fsdp_managed_module -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask +from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, @@ -490,6 +490,8 @@ def forward(self, hidden_states, attention_mask=None): class Wav2Vec2BertAdapterLayer(nn.Module): def __init__(self, config): super().__init__() + self.config = config + embed_dim = config.output_hidden_size dropout = config.conformer_conv_dropout @@ -552,9 +554,10 @@ def forward( if attention_mask is not None: attention_mask = _compute_new_attention_mask(hidden_states=hidden_states, seq_lens=sub_sampled_lengths) - attention_mask = _prepare_4d_attention_mask( - attention_mask, - hidden_states.dtype, + attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, ) # The rest of the computation is identical to a vanilla Transformer diff --git a/src/transformers/models/x_clip/modeling_x_clip.py b/src/transformers/models/x_clip/modeling_x_clip.py index 72bc8700c10a..6a62299f188b 100644 --- a/src/transformers/models/x_clip/modeling_x_clip.py +++ b/src/transformers/models/x_clip/modeling_x_clip.py @@ -23,7 +23,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -36,7 +36,6 @@ logging, torch_int, ) -from ...utils.generic import is_flash_attention_requested from .configuration_x_clip import XCLIPConfig, XCLIPTextConfig, XCLIPVisionConfig @@ -275,8 +274,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = False, + **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Input shape: Batch x Time x Channel""" @@ -289,15 +288,6 @@ def forward( queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - # CLIP text model uses both `causal_attention_mask` and `attention_mask` - # in case FA2 kernel is called, `is_causal` should be inferred from `causal_attention_mask` - if not is_flash_attention_requested(self.config): - if attention_mask is not None and causal_attention_mask is not None: - attention_mask = attention_mask + causal_attention_mask - elif causal_attention_mask is not None: - attention_mask = causal_attention_mask - else: - self.is_causal = causal_attention_mask is not None attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward @@ -351,8 +341,8 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: @@ -370,8 +360,8 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = residual + hidden_states @@ -444,7 +434,6 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - causal_attention_mask: torch.Tensor, output_attentions: bool | None = False, ) -> tuple[torch.FloatTensor]: """ @@ -453,11 +442,6 @@ def forward( attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. `(config.encoder_attention_heads,)`. - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. @@ -479,7 +463,6 @@ def forward( hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, ) hidden_states = residual + hidden_states @@ -578,10 +561,10 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutput: r""" Args: @@ -595,13 +578,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -628,8 +604,8 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, + **kwargs, ) hidden_states = layer_outputs[0] @@ -645,15 +621,17 @@ def forward( ) -class XCLIPTextTransformer(nn.Module): +class XCLIPTextTransformer(XCLIPPreTrainedModel): def __init__(self, config: XCLIPTextConfig): - super().__init__() - self.config = config + super().__init__(config) + embed_dim = config.hidden_size self.embeddings = XCLIPTextEmbeddings(config) self.encoder = XCLIPEncoder(config) self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.post_init() + @auto_docstring def forward( self, @@ -663,6 +641,7 @@ def forward( output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, + **kwargs, ) -> tuple | BaseModelOutputWithPooling: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -678,23 +657,23 @@ def forward( hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) - # X_CLIP's text model uses causal mask, prepare it here. - # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 - causal_attention_mask = _create_4d_causal_attention_mask( - input_shape, hidden_states.dtype, device=hidden_states.device + attention_mask = create_causal_mask( + config=self.config, + input_embeds=hidden_states, + attention_mask=attention_mask, + cache_position=torch.arange(hidden_states.shape[1], device=hidden_states.device), + past_key_values=None, ) - # expand attention_mask - if attention_mask is not None: - # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + kwargs.pop("is_causal", None) encoder_outputs = self.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + is_causal=True, + **kwargs, ) last_hidden_state = encoder_outputs[0] @@ -786,7 +765,6 @@ def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, - causal_attention_mask: torch.Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, @@ -803,13 +781,6 @@ def forward( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) - causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Causal mask for the text model. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under @@ -836,7 +807,6 @@ def forward( layer_outputs = encoder_layer( hidden_states, attention_mask, - causal_attention_mask, output_attentions=output_attentions, ) diff --git a/src/transformers/models/xglm/modeling_xglm.py b/src/transformers/models/xglm/modeling_xglm.py index f84e92a207de..183930968d51 100755 --- a/src/transformers/models/xglm/modeling_xglm.py +++ b/src/transformers/models/xglm/modeling_xglm.py @@ -22,7 +22,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions from ...modeling_utils import PreTrainedModel @@ -464,8 +464,17 @@ def forward( ) past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length + if cache_position is None: + cache_position = torch.arange( + past_key_values_length, past_key_values_length + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, ) if position_ids is None: @@ -479,9 +488,11 @@ def forward( # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _prepare_4d_attention_mask( - encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, ) hidden_states = inputs_embeds + self.embed_positions(position_ids, past_key_values_length).to( diff --git a/src/transformers/models/zamba/modeling_zamba.py b/src/transformers/models/zamba/modeling_zamba.py index 0adc5fe505df..80d9ae9b71fe 100644 --- a/src/transformers/models/zamba/modeling_zamba.py +++ b/src/transformers/models/zamba/modeling_zamba.py @@ -30,14 +30,13 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, logging -from ...utils.generic import is_flash_attention_requested from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_zamba import ZambaConfig @@ -188,6 +187,14 @@ def get_seq_length(self, layer_idx: int | None = 0) -> int: return 0 return self.key_cache[layer_idx].shape[-2] + # Copied from transformers.models.jamba.modeling_jamba.HybridMambaAttentionDynamicCache.get_mask_sizes + def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]: + """Return the length and offset of the cache, used to generate the mask""" + kv_offset = 0 + query_length = cache_position.shape[0] + kv_length = self.get_seq_length(layer_idx) + query_length + return kv_length, kv_offset + def eager_attention_forward( module: nn.Module, @@ -910,7 +917,14 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position) + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None @@ -955,45 +969,6 @@ def forward( ) return output if return_dict else output.to_tuple() - def _update_causal_mask(self, attention_mask, input_tensor, cache_position): - if is_flash_attention_requested(self.config): - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - dtype, device = input_tensor.dtype, input_tensor.device - min_dtype = torch.finfo(dtype).min - sequence_length = input_tensor.shape[1] - if cache_position is None: - target_length = sequence_length - else: - target_length = cache_position[-1] + 1 - - causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - if cache_position is not None: - causal_mask *= (torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)).to(torch.bool) - causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - if attention_mask.dim() == 2: - expanded_attn_mask = attention_mask[:, None, None, :].expand(-1, 1, sequence_length, -1) - padding_mask = causal_mask.eq(0.0) * expanded_attn_mask.eq(0.0) - causal_mask.masked_fill_(padding_mask, min_dtype) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - # Adapted from transformers.models.jamba.modeling_jamba.JambaForCausalLM with Jamba->Zamba, JAMBA->ZAMBA class ZambaForCausalLM(ZambaPreTrainedModel, GenerationMixin): diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py index d6dc70a406b5..d5241cafcbef 100644 --- a/src/transformers/models/zamba2/modeling_zamba2.py +++ b/src/transformers/models/zamba2/modeling_zamba2.py @@ -32,7 +32,7 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...integrations import use_kernel_func_from_hub -from ...modeling_attn_mask_utils import AttentionMaskConverter +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast @@ -40,7 +40,7 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, is_torchdynamo_compiling, logging -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils.generic import maybe_autocast from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_zamba2 import Zamba2Config @@ -190,6 +190,13 @@ def get_seq_length(self, layer_idx: int | None = 0) -> int: return 0 return self.key_cache[layer_idx].shape[-2] + def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]: + """Return the length and offset of the cache, used to generate the mask""" + kv_offset = 0 + query_length = cache_position.shape[0] + kv_length = self.get_seq_length(layer_idx) + query_length + return kv_length, kv_offset + def update_conv_state( self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor ) -> torch.Tensor: @@ -1323,7 +1330,14 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position) + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) all_hidden_states = () if output_hidden_states else None @@ -1370,45 +1384,6 @@ def forward( ) return output if return_dict else output.to_tuple() - def _update_causal_mask(self, attention_mask, input_tensor, cache_position): - if is_flash_attention_requested(self.config): - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - dtype, device = input_tensor.dtype, input_tensor.device - min_dtype = torch.finfo(dtype).min - sequence_length = input_tensor.shape[1] - if cache_position is None: - target_length = sequence_length - else: - target_length = cache_position[-1] + 1 - - causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - if cache_position is not None: - causal_mask *= (torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)).to(torch.bool) - causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - if attention_mask.dim() == 2: - expanded_attn_mask = attention_mask[:, None, None, :].expand(-1, 1, sequence_length, -1) - padding_mask = causal_mask.eq(0.0) * expanded_attn_mask.eq(0.0) - causal_mask.masked_fill_(padding_mask, min_dtype) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu", "npu"] - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - def get_layers(self): layers = [] self._tied_weights_keys = {} diff --git a/src/transformers/models/zamba2/modular_zamba2.py b/src/transformers/models/zamba2/modular_zamba2.py index 4dd9a0753a81..7ac5558276c3 100644 --- a/src/transformers/models/zamba2/modular_zamba2.py +++ b/src/transformers/models/zamba2/modular_zamba2.py @@ -21,6 +21,7 @@ from ... import initialization as init from ...activations import ACT2FN +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -161,6 +162,13 @@ def get_seq_length(self, layer_idx: int | None = 0) -> int: return 0 return self.key_cache[layer_idx].shape[-2] + def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]: + """Return the length and offset of the cache, used to generate the mask""" + kv_offset = 0 + query_length = cache_position.shape[0] + kv_length = self.get_seq_length(layer_idx) + query_length + return kv_length, kv_offset + class Zamba2RotaryEmbedding(LlamaRotaryEmbedding): pass @@ -1042,7 +1050,14 @@ def forward( if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position) + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) all_hidden_states = () if output_hidden_states else None diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 50a95e796e4f..3feeb3c8277a 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -2023,12 +2023,12 @@ def attention_mask_padding_matches_padding_free_with_position_ids( if config.is_encoder_decoder: self.skipTest("Model is an encoder-decoder") - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: self.skipTest("Model dummy inputs should contain text input ids") + if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: + self.skipTest("Model dummy inputs should contain padding in their attention mask") + # make sure that all models have enough positions for generation dummy_input_ids = inputs_dict["input_ids"] if hasattr(config, "max_position_embeddings"): @@ -2061,6 +2061,13 @@ def attention_mask_padding_matches_padding_free_with_position_ids( dummy_attention_mask = inputs_dict["attention_mask"] dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id + # We need to prepare position ids according to the attention mask as we use it to extract embeddings that + # rely on the correct position - naively increasing sequences do not suffice anymore atp. The solution here + # calculates an increasing sequences for all 1s and puts 0s else. + inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( + inputs_dict["attention_mask"] == 1 + ).long() + model = ( model_class.from_pretrained( tmpdirname, @@ -2083,6 +2090,7 @@ def attention_mask_padding_matches_padding_free_with_position_ids( padfree_inputs_dict = { k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() } + padfree_inputs_dict.pop("labels", None) # can lead to silent upcasts on logits else: # create packed position_ids position_ids = ( diff --git a/tests/models/autoformer/test_modeling_autoformer.py b/tests/models/autoformer/test_modeling_autoformer.py index 0a96be5f73cb..35b8a79ad66e 100644 --- a/tests/models/autoformer/test_modeling_autoformer.py +++ b/tests/models/autoformer/test_modeling_autoformer.py @@ -210,6 +210,7 @@ class AutoformerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCa test_missing_keys = False test_inputs_embeds = False + test_torch_exportable = False # not worth to fix def setUp(self): self.model_tester = AutoformerModelTester(self) diff --git a/tests/models/bert/test_modeling_bert.py b/tests/models/bert/test_modeling_bert.py index 70f272e43b4d..7b30c46990b7 100644 --- a/tests/models/bert/test_modeling_bert.py +++ b/tests/models/bert/test_modeling_bert.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import tempfile import unittest import pytest @@ -45,7 +43,6 @@ BertForTokenClassification, BertLMHeadModel, BertModel, - DataCollatorWithFlattening, ) @@ -571,120 +568,6 @@ def test_model_from_pretrained(self): model = BertModel.from_pretrained(model_name) self.assertIsNotNone(model) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - @require_torch class BertModelIntegrationTest(unittest.TestCase): diff --git a/tests/models/bert_generation/test_modeling_bert_generation.py b/tests/models/bert_generation/test_modeling_bert_generation.py index 7b48c614e892..fe55e374b164 100644 --- a/tests/models/bert_generation/test_modeling_bert_generation.py +++ b/tests/models/bert_generation/test_modeling_bert_generation.py @@ -13,8 +13,6 @@ # limitations under the License. -import inspect -import tempfile import unittest from transformers import BertGenerationConfig, is_torch_available @@ -29,7 +27,7 @@ if is_torch_available(): import torch - from transformers import BertGenerationDecoder, BertGenerationEncoder, DataCollatorWithFlattening + from transformers import BertGenerationDecoder, BertGenerationEncoder class BertGenerationEncoderTester: @@ -310,120 +308,6 @@ def test_model_from_pretrained(self): model = BertGenerationEncoder.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder") self.assertIsNotNone(model) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - @require_torch class BertGenerationEncoderIntegrationTest(unittest.TestCase): diff --git a/tests/models/bloom/test_modeling_bloom.py b/tests/models/bloom/test_modeling_bloom.py index a6b3f7ed662d..d1ad38d43109 100644 --- a/tests/models/bloom/test_modeling_bloom.py +++ b/tests/models/bloom/test_modeling_bloom.py @@ -224,7 +224,7 @@ def test_embeddings(self): # TODO change the script (or just add skip) when building the env with tokenizers 0.12.0 """ # The config in this checkpoint has `bfloat16` as `dtype` -> model in `bfloat16` - model = BloomForCausalLM.from_pretrained(self.path_bigscience_model, dtype="auto") + model = BloomForCausalLM.from_pretrained(self.path_bigscience_model, use_safetensors=False, dtype="auto") model.eval() EMBEDDINGS_DS_BEFORE_LN_BF_16_MEAN = { @@ -452,7 +452,9 @@ def test_embeddings(self): @require_torch def test_hidden_states_transformers(self): - model = BloomModel.from_pretrained(self.path_bigscience_model, use_cache=False, dtype="auto").to(torch_device) + model = BloomModel.from_pretrained( + self.path_bigscience_model, use_safetensors=False, use_cache=False, dtype="auto" + ).to(torch_device) model.eval() EXAMPLE_IDS = [3478, 368, 109586, 35433, 2, 77, 132619, 3478, 368, 109586, 35433, 2, 2175, 23714, 73173, 144252, 2, 77, 132619, 3478] # fmt: skip @@ -477,9 +479,9 @@ def test_hidden_states_transformers(self): @require_torch def test_logits(self): - model = BloomForCausalLM.from_pretrained(self.path_bigscience_model, use_cache=False, dtype="auto").to( - torch_device - ) # load in bf16 + model = BloomForCausalLM.from_pretrained( + self.path_bigscience_model, use_safetensors=False, use_cache=False, dtype="auto" + ).to(torch_device) # load in bf16 model.eval() EXAMPLE_IDS = [3478, 368, 109586, 35433, 2, 77, 132619, 3478, 368, 109586, 35433, 2, 2175, 23714, 73173, 144252, 2, 77, 132619, 3478] # fmt: skip @@ -518,7 +520,9 @@ def test_simple_generation(self): # >=1b1 + allow_fp16_reduced_precision_reduction = False + torch.bmm ==> PASS path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) + model = BloomForCausalLM.from_pretrained( + path_560m, use_safetensors=False, use_cache=True, revision="gs555750" + ).to(torch_device) model = model.eval() tokenizer = AutoTokenizer.from_pretrained(path_560m) @@ -538,7 +542,9 @@ def test_simple_generation(self): @require_torch_accelerator def test_batch_generation(self): path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) + model = BloomForCausalLM.from_pretrained( + path_560m, use_safetensors=False, use_cache=True, revision="gs555750" + ).to(torch_device) model = model.eval() tokenizer = AutoTokenizer.from_pretrained(path_560m, padding_side="left") @@ -558,7 +564,9 @@ def test_batch_generation(self): @require_torch_accelerator def test_batch_generation_padding(self): path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) + model = BloomForCausalLM.from_pretrained( + path_560m, use_safetensors=False, use_cache=True, revision="gs555750" + ).to(torch_device) model = model.eval() tokenizer = AutoTokenizer.from_pretrained(path_560m, padding_side="left") @@ -588,7 +596,9 @@ def test_batch_generation_padding(self): def test_batch_generated_text(self): path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) + model = BloomForCausalLM.from_pretrained( + path_560m, use_safetensors=False, use_cache=True, revision="gs555750" + ).to(torch_device) model = model.eval() tokenizer = AutoTokenizer.from_pretrained(path_560m, padding_side="left") diff --git a/tests/models/data2vec/test_modeling_data2vec_text.py b/tests/models/data2vec/test_modeling_data2vec_text.py index 4709920e1817..dbdc51622d6a 100644 --- a/tests/models/data2vec/test_modeling_data2vec_text.py +++ b/tests/models/data2vec/test_modeling_data2vec_text.py @@ -13,8 +13,6 @@ # limitations under the License. """Testing suite for the PyTorch Data2VecAudio model.""" -import inspect -import tempfile import unittest from tests.test_modeling_common import floats_tensor, ids_tensor, random_attention_mask @@ -38,7 +36,6 @@ Data2VecTextForSequenceClassification, Data2VecTextForTokenClassification, Data2VecTextModel, - DataCollatorWithFlattening, ) from transformers.models.data2vec.modeling_data2vec_text import Data2VecTextEmbeddings @@ -504,120 +501,6 @@ def test_create_position_ids_from_inputs_embeds(self): self.assertEqual(position_ids.shape, expected_positions.shape) self.assertTrue(torch.all(torch.eq(position_ids, expected_positions))) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - @require_torch class Data2VecTextModelIntegrationTest(TestCasePlus): diff --git a/tests/models/electra/test_modeling_electra.py b/tests/models/electra/test_modeling_electra.py index 7bf778aa4d9e..65326b324e00 100644 --- a/tests/models/electra/test_modeling_electra.py +++ b/tests/models/electra/test_modeling_electra.py @@ -13,8 +13,6 @@ # limitations under the License. -import inspect -import tempfile import unittest from transformers import ElectraConfig, is_torch_available @@ -32,7 +30,6 @@ from transformers import ( MODEL_FOR_PRETRAINING_MAPPING, - DataCollatorWithFlattening, ElectraForCausalLM, ElectraForMaskedLM, ElectraForMultipleChoice, @@ -472,120 +469,6 @@ def test_for_causal_lm(self): config_and_inputs = self.model_tester.prepare_config_and_inputs_for_decoder() self.model_tester.create_and_check_electra_for_causal_lm(*config_and_inputs) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - @require_torch class ElectraModelIntegrationTest(unittest.TestCase): diff --git a/tests/models/ernie/test_modeling_ernie.py b/tests/models/ernie/test_modeling_ernie.py index 6ca4d53eeea1..2784a27bd6f6 100644 --- a/tests/models/ernie/test_modeling_ernie.py +++ b/tests/models/ernie/test_modeling_ernie.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import tempfile import unittest from transformers import ErnieConfig, is_torch_available @@ -30,7 +28,6 @@ from transformers import ( MODEL_FOR_PRETRAINING_MAPPING, - DataCollatorWithFlattening, ErnieForCausalLM, ErnieForMaskedLM, ErnieForMultipleChoice, @@ -563,117 +560,3 @@ def test_model_from_pretrained(self): model_name = "nghuyong/ernie-1.0-base-zh" model = ErnieModel.from_pretrained(model_name) self.assertIsNotNone(model) - - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) diff --git a/tests/models/falcon/test_modeling_falcon.py b/tests/models/falcon/test_modeling_falcon.py index 6c6cc8c12b42..c1e2935d30ed 100644 --- a/tests/models/falcon/test_modeling_falcon.py +++ b/tests/models/falcon/test_modeling_falcon.py @@ -67,6 +67,10 @@ def is_pipeline_test_to_skip( ): return True + @unittest.skip(reason="To support alibi, we are forced to create a mask in all cases") + def test_sdpa_can_dispatch_on_flash(self): + pass + @require_torch class FalconLanguageGenerationTest(unittest.TestCase): diff --git a/tests/models/gpt2/test_modeling_gpt2.py b/tests/models/gpt2/test_modeling_gpt2.py index 9fb8e3f29593..f0e829af0085 100644 --- a/tests/models/gpt2/test_modeling_gpt2.py +++ b/tests/models/gpt2/test_modeling_gpt2.py @@ -108,8 +108,8 @@ def get_config(self, scale_attn_by_inverse_layer_idx=False, reorder_and_upcast_a def prepare_config_and_inputs_for_common(self): # Overwritten: we want `token_type_ids` as part of the common inputs config_and_inputs = self.prepare_config_and_inputs(extra_inputs=True) - config, input_ids, _, token_type_ids, _, _, _, _ = config_and_inputs - inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids} + config, input_ids, attention_mask, token_type_ids, _, _, _, _ = config_and_inputs + inputs_dict = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} return config, inputs_dict def prepare_config_and_inputs_for_decoder(self): @@ -185,6 +185,7 @@ def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): device=torch_device, ) inputs_dict["input_ids"] = inputs_dict["labels"] + inputs_dict["attention_mask"] = torch.tril(torch.ones_like(inputs_dict["input_ids"]).to(torch_device)) inputs_dict["token_type_ids"] = inputs_dict["labels"] inputs_dict["mc_token_ids"] = torch.zeros( (self.model_tester.batch_size, self.model_tester.num_choices), diff --git a/tests/models/gptj/test_modeling_gptj.py b/tests/models/gptj/test_modeling_gptj.py index 3f6bab0a827f..f7974f97f49c 100644 --- a/tests/models/gptj/test_modeling_gptj.py +++ b/tests/models/gptj/test_modeling_gptj.py @@ -464,7 +464,7 @@ def test_batch_generation(self): @slow def test_model_from_pretrained(self): model_name = "EleutherAI/gpt-j-6B" - model = GPTJModel.from_pretrained(model_name, revision="float16", dtype=torch.float16) + model = GPTJModel.from_pretrained(model_name, use_safetensors=False, revision="float16", dtype=torch.float16) self.assertIsNotNone(model) diff --git a/tests/models/kosmos2_5/test_modeling_kosmos2_5.py b/tests/models/kosmos2_5/test_modeling_kosmos2_5.py index ebaa95553cc2..e4ab26d4721f 100644 --- a/tests/models/kosmos2_5/test_modeling_kosmos2_5.py +++ b/tests/models/kosmos2_5/test_modeling_kosmos2_5.py @@ -471,9 +471,9 @@ def test_model_parallelism(self): def test_sdpa_can_dispatch_on_flash(self): pass - # TODO: ydshieh - @unittest.skip(reason="doesn't support padding yet") - def test_eager_matches_sdpa_inference_1_bfloat16(self): + # TODO: vasqu + @unittest.skip(reason="why the heck does this have bigger tols") + def test_eager_matches_sdpa_inference_24_fp32_pad_left_output_attentions(self): pass # TODO: ydshieh diff --git a/tests/models/opt/test_modeling_opt.py b/tests/models/opt/test_modeling_opt.py index d12fbd0812f9..c13bd7e9e8b9 100644 --- a/tests/models/opt/test_modeling_opt.py +++ b/tests/models/opt/test_modeling_opt.py @@ -332,6 +332,12 @@ def test_opt_sequence_classification_model_for_multi_label(self): def test_model_parallelism(self): super().test_model_parallelism() + @unittest.skip(reason="Position ids follow custom logic") + def attention_mask_padding_matches_padding_free_with_position_ids( + self, attn_implementation: str, fa_kwargs: bool = False + ): + pass + def assert_tensors_close(a, b, atol=1e-12, prefix=""): """If tensors have different shapes, different values or a and b are not both tensors, raise a nice Assertion error.""" diff --git a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py index 6bcb46323892..a73f83c6c261 100644 --- a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py +++ b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py @@ -21,6 +21,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, is_torch_available, set_seed from transformers.testing_utils import ( Expectations, + cleanup, require_bitsandbytes, require_torch, require_torch_accelerator, @@ -129,6 +130,12 @@ class RecurrentGemmaIntegrationTest(unittest.TestCase): input_long_text = ['Marseille, France (CNN)The French prosecutor leading an investigation into the crash of Germanwings Flight 9525 insisted Wednesday that he was not aware of any video footage from on board the plane. Marseille prosecutor Brice Robin told CNN that "so far no videos were used in the crash investigation." He added, "A person who has such a video needs to immediately give it to the investigators." Robin\'s comments follow claims by two magazines, German daily Bild and French Paris Match, of a cell phone video showing the harrowing final seconds from on board Germanwings Flight 9525 as it crashed into the French Alps. All 150 on board were killed. Paris Match and Bild reported that the video was recovered from a phone at the wreckage site. The two publications described the supposed video, but did not post it on their websites. The publications said that they watched the video, which was found by a source close to the investigation. "One can hear cries of \'My God\' in several languages," Paris Match reported. "Metallic banging can also be heard more than three times, perhaps of the pilot trying to open the cockpit door with a heavy object. Towards the end, after a heavy shake, stronger than the others, the screaming intensifies. Then nothing." "It is a very disturbing scene," said Julian Reichelt, editor-in-chief of Bild online. An official with France\'s accident investigation agency, the BEA, said the agency is not aware of any such video. Lt. Col.'] # fmt: skip model_id = "google/recurrentgemma-2b" + def setup(self): + cleanup(torch_device, gc_collect=True) + + def tearDown(self): + cleanup(torch_device, gc_collect=True) + def test_2b_generate(self): EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', 'Hi today is a new app that allows you to make money by watching videos.\n\nThe app is very simple to use and you can earn money by watching videos.\n\nThe app is available for both Android and iOS devices and you can download it from the Google Play Store or the App Store.\n\nOnce you have downloaded the app'] # fmt: skip model = AutoModelForCausalLM.from_pretrained( diff --git a/tests/models/roberta/test_modeling_roberta.py b/tests/models/roberta/test_modeling_roberta.py index 4220c9ea0c59..4af3f445e175 100644 --- a/tests/models/roberta/test_modeling_roberta.py +++ b/tests/models/roberta/test_modeling_roberta.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import tempfile import unittest import pytest @@ -31,7 +29,6 @@ import torch from transformers import ( - DataCollatorWithFlattening, RobertaForCausalLM, RobertaForMaskedLM, RobertaForMultipleChoice, @@ -511,120 +508,6 @@ def test_create_position_ids_from_inputs_embeds(self): self.assertEqual(position_ids.shape, expected_positions.shape) self.assertTrue(torch.all(torch.eq(position_ids, expected_positions))) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - @require_torch class RobertaModelIntegrationTest(TestCasePlus): diff --git a/tests/models/roberta_prelayernorm/test_modeling_roberta_prelayernorm.py b/tests/models/roberta_prelayernorm/test_modeling_roberta_prelayernorm.py index 9394d3a8dcfd..415d0dbb8da3 100644 --- a/tests/models/roberta_prelayernorm/test_modeling_roberta_prelayernorm.py +++ b/tests/models/roberta_prelayernorm/test_modeling_roberta_prelayernorm.py @@ -13,8 +13,6 @@ # limitations under the License. -import inspect -import tempfile import unittest from transformers import RobertaPreLayerNormConfig, is_torch_available @@ -30,7 +28,6 @@ import torch from transformers import ( - DataCollatorWithFlattening, RobertaPreLayerNormForCausalLM, RobertaPreLayerNormForMaskedLM, RobertaPreLayerNormForMultipleChoice, @@ -522,120 +519,6 @@ def test_create_position_ids_from_inputs_embeds(self): self.assertEqual(position_ids.shape, expected_positions.shape) self.assertTrue(torch.all(torch.eq(position_ids, expected_positions))) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - @require_torch class RobertaPreLayerNormModelIntegrationTest(TestCasePlus): diff --git a/tests/models/roc_bert/test_modeling_roc_bert.py b/tests/models/roc_bert/test_modeling_roc_bert.py index dc92ffbefd22..ca801e265c0e 100644 --- a/tests/models/roc_bert/test_modeling_roc_bert.py +++ b/tests/models/roc_bert/test_modeling_roc_bert.py @@ -13,8 +13,6 @@ # limitations under the License. """Testing suite for the PyTorch RoCBert model.""" -import inspect -import tempfile import unittest from transformers import RoCBertConfig, is_torch_available @@ -31,7 +29,6 @@ from transformers import ( MODEL_FOR_PRETRAINING_MAPPING, - DataCollatorWithFlattening, RoCBertForCausalLM, RoCBertForMaskedLM, RoCBertForMultipleChoice, @@ -698,120 +695,6 @@ def test_model_from_pretrained(self): model = RoCBertModel.from_pretrained(model_name) self.assertIsNotNone(model) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - def flash_attn_inference_equivalence( self, attn_implementation: str, padding_side: str, atol: float = 4e-2, rtol: float = 4e-2 ): diff --git a/tests/models/switch_transformers/test_modeling_switch_transformers.py b/tests/models/switch_transformers/test_modeling_switch_transformers.py index 687c260a9992..87d177d91d4f 100644 --- a/tests/models/switch_transformers/test_modeling_switch_transformers.py +++ b/tests/models/switch_transformers/test_modeling_switch_transformers.py @@ -620,7 +620,7 @@ def test_v1_1_resize_embeddings(self): @slow def test_model_from_pretrained(self): model_name = "google/switch-base-8" - model = SwitchTransformersModel.from_pretrained(model_name) + model = SwitchTransformersModel.from_pretrained(model_name, use_safetensors=False) self.assertIsNotNone(model) @unittest.skip( @@ -945,7 +945,9 @@ def test_small_logits(self): and `transformers` implementation of Switch-C transformers. We only check the logits of the first batch. """ - model = SwitchTransformersModel.from_pretrained("google/switch-base-8", dtype=torch.bfloat16).to(torch_device) + model = SwitchTransformersModel.from_pretrained( + "google/switch-base-8", use_safetensors=False, dtype=torch.bfloat16 + ).to(torch_device) input_ids = torch.ones((32, 64), dtype=torch.long).to(torch_device) decoder_input_ids = torch.ones((32, 64), dtype=torch.long).to(torch_device) diff --git a/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py b/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py index e556dbf76a1d..7371a354b5ce 100644 --- a/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py +++ b/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py @@ -13,7 +13,6 @@ # limitations under the License. -import inspect import tempfile import unittest @@ -30,7 +29,6 @@ import torch from transformers import ( - DataCollatorWithFlattening, XLMRobertaXLForCausalLM, XLMRobertaXLForMaskedLM, XLMRobertaXLForMultipleChoice, @@ -514,120 +512,6 @@ def test_create_position_ids_from_inputs_embeds(self): self.assertEqual(position_ids.shape, expected_positions.shape) self.assertTrue(torch.all(torch.eq(position_ids, expected_positions))) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - def flash_attn_inference_equivalence( self, attn_implementation: str, padding_side: str, atol: float = 4e-2, rtol: float = 4e-2 ): diff --git a/tests/models/xmod/test_modeling_xmod.py b/tests/models/xmod/test_modeling_xmod.py index 583bd51e6375..1cfbf23f64fc 100644 --- a/tests/models/xmod/test_modeling_xmod.py +++ b/tests/models/xmod/test_modeling_xmod.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import tempfile import unittest from transformers import XLMRobertaTokenizer, is_torch_available @@ -28,7 +26,6 @@ import torch from transformers import ( - DataCollatorWithFlattening, XmodConfig, XmodForCausalLM, XmodForMaskedLM, @@ -528,120 +525,6 @@ def test_freeze_embeddings_and_language_adapters(self): num_trainable_params_after = sum(p.numel() for p in model.parameters() if p.requires_grad) self.assertLess(num_trainable_params_after, num_trainable_params_before) - def attention_mask_padding_matches_padding_free_with_position_ids( - self, attn_implementation: str, fa_kwargs: bool = False - ): - """ - Overwritten to account for the embeddings that rely on position ids. - """ - if not self.has_attentions: - self.skipTest(reason="Model architecture does not support attentions") - - max_new_tokens = 30 - support_flag = { - "sdpa": "_supports_sdpa", - "flash_attention_2": "_supports_flash_attn", - "flash_attention_3": "_supports_flash_attn", - } - - for model_class in self.all_generative_model_classes: - if attn_implementation != "eager" and not getattr(model_class, support_flag[attn_implementation]): - self.skipTest(f"{model_class.__name__} does not support {attn_implementation}") - - # can't infer if new attn mask API is supported by assume that only model with attention backend support it - if not model_class._supports_attention_backend: - self.skipTest(f"{model_class.__name__} does not support new attention mask API") - - if model_class._is_stateful: # non-transformer models most probably have no packing support - self.skipTest(f"{model_class.__name__} doesn't support packing!") - - config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - if config.is_encoder_decoder: - self.skipTest("Model is an encoder-decoder") - - if 0 not in inputs_dict.get("attention_mask", []) or "attention_mask" not in inputs_dict: - self.skipTest("Model dummy inputs should contain padding in their attention mask") - - if "input_ids" not in inputs_dict or inputs_dict["input_ids"].ndim != 2: - self.skipTest("Model dummy inputs should contain text input ids") - - # make sure that all models have enough positions for generation - dummy_input_ids = inputs_dict["input_ids"] - if hasattr(config, "max_position_embeddings"): - config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1 - - model = model_class(config) - if "position_ids" not in inspect.signature(model.forward).parameters: - self.skipTest("Model does not support position_ids") - - if (not fa_kwargs) and "position_ids" not in inspect.signature(model.forward).parameters: - continue # this model doesn't accept position ids as input - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname) - - # Drop all keys except for the minimal set. Hard to manipulate with multimodals etc - inputs_dict = {k: v for k, v in inputs_dict.items() if k in ["input_ids", "attention_mask"]} - - # Ensure left padding, to adapt for some models - if 0 in inputs_dict["attention_mask"][:, -1]: - inputs_dict["attention_mask"] = inputs_dict["attention_mask"].flip(1) - dummy_attention_mask = inputs_dict["attention_mask"] - dummy_input_ids[~dummy_attention_mask.bool()] = config.get_text_config().pad_token_id - - # Main difference to other models, we need to prepare position ids according to the attention mask - # as we use it to extract embeddings that rely on the correct position - naively increasing sequences do - # not suffice anymore atp. The solution here calculates an increasing sequences for all 1s and puts 0s else. - inputs_dict["position_ids"] = ((inputs_dict["attention_mask"] == 1).long().cumsum(dim=1) - 1) * ( - inputs_dict["attention_mask"] == 1 - ).long() - - model = ( - model_class.from_pretrained( - tmpdirname, - dtype=torch.bfloat16, - attn_implementation=attn_implementation, - ) - .to(torch_device) - .eval() - ) - - if fa_kwargs: - # flatten - features = [ - {"input_ids": i[a.bool()].tolist()} for i, a in zip(dummy_input_ids, dummy_attention_mask) - ] - - # add position_ids + fa_kwargs - data_collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) - batch = data_collator(features) - padfree_inputs_dict = { - k: t.to(torch_device) if torch.is_tensor(t) else t for k, t in batch.items() - } - else: - # create packed position_ids - position_ids = ( - torch.cat([torch.arange(length) for length in dummy_attention_mask.sum(1).tolist()]) - .long() - .unsqueeze(0) - .to(torch_device) - ) - padfree_inputs_dict = { - "input_ids": dummy_input_ids[dummy_attention_mask.bool()].unsqueeze(0), - "position_ids": position_ids, - } - - # We need to do simple forward without cache in order to trigger packed SDPA/flex/eager attention path - res_padded = model(**inputs_dict, use_cache=False) - res_padfree = model(**padfree_inputs_dict, use_cache=False) - - logits_padded = res_padded.logits[dummy_attention_mask.bool()] - logits_padfree = res_padfree.logits[0] - - # acceptable numerical instability - tol = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol) - @require_sentencepiece @require_tokenizers diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 2228861ea3e5..3ef9ce00cce2 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -125,12 +125,6 @@ T5ForConditionalGeneration, ) from transformers.conversion_mapping import MergeModulelist, WeightConverter, get_model_conversion_mapping - from transformers.modeling_attn_mask_utils import ( - AttentionMaskConverter, - _create_4d_causal_attention_mask, - _prepare_4d_attention_mask, - _prepare_4d_causal_attention_mask, - ) from transformers.modeling_utils import ( FLASH_ATTN_KERNEL_FALLBACK, _find_disjoint, @@ -337,32 +331,6 @@ def __init__(self, config): def forward(self, x): return self.head(self.language_model(self.vision_model(x))) - class Prepare4dCausalAttentionMaskModel(nn.Module): - def forward(self, inputs_embeds): - batch_size, seq_length, _ = inputs_embeds.shape - past_key_values_length = 4 - attention_mask = _prepare_4d_causal_attention_mask( - None, (batch_size, seq_length), inputs_embeds, past_key_values_length - ) - return attention_mask - - class Create4dCausalAttentionMaskModel(nn.Module): - def forward(self, inputs_embeds): - batch_size, seq_length, _ = inputs_embeds.shape - past_key_values_length = 4 - attention_mask = _create_4d_causal_attention_mask( - (batch_size, seq_length), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - past_key_values_length=past_key_values_length, - ) - return attention_mask - - class Prepare4dAttentionMaskModel(nn.Module): - def forward(self, mask, inputs_embeds): - attention_mask = _prepare_4d_attention_mask(mask, dtype=inputs_embeds.dtype) - return attention_mask - class TestOffline(unittest.TestCase): def test_offline(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -2782,295 +2750,6 @@ def test_push_to_hub_with_tags(self): self.assertEqual(loaded_model_card.data.tags, new_tags) -@require_torch -class AttentionMaskTester(unittest.TestCase): - def check_non_causal(self, bsz, q_len, kv_len, mask_2d, mask_4d): - mask_indices = (mask_2d != 1)[:, None].broadcast_to((bsz, q_len, kv_len)) - mask_4d_values = mask_4d[:, 0][mask_indices] - is_inf = mask_4d_values == -float("inf") - is_min = mask_4d_values == torch.finfo(mask_4d.dtype).min - assert torch.logical_or(is_inf, is_min).all() - - def check_to_4d(self, mask_converter, q_len, kv_len, additional_mask=None, bsz=3): - mask_2d = torch.ones((bsz, kv_len), device=torch_device, dtype=torch.long) - - if additional_mask is not None: - for bsz_idx, seq_idx in additional_mask: - mask_2d[bsz_idx, seq_idx] = 0 - - mask_4d = mask_converter.to_4d(mask_2d, query_length=q_len, key_value_length=kv_len, dtype=torch.float32) - - assert mask_4d.shape == (bsz, 1, q_len, kv_len) - - # make sure there are no overflows - assert mask_4d.min() != float("-inf") - - context = mask_converter.sliding_window - if mask_converter.is_causal and context is None: - # k * (k+1) / 2 tokens are masked in triangualar masks - num_tokens_masked = bsz * (q_len * (q_len - 1) // 2) - - if 0 not in mask_2d: - assert (mask_4d != 0).sum().item() == num_tokens_masked - if 0 in mask_2d: - # at least causal mask + maybe more - assert (mask_4d != 0).sum().item() >= num_tokens_masked - self.check_non_causal(bsz, q_len, kv_len, mask_2d, mask_4d) - elif not mask_converter.is_causal and context is None: - if 0 not in mask_2d: - assert (mask_4d != 0).sum().item() == 0 - if 0 in mask_2d: - self.check_non_causal(bsz, q_len, kv_len, mask_2d, mask_4d) - elif mask_converter.is_causal and context is not None: - # k * (k+1) / 2 tokens are masked in triangualar masks - num_tokens_masked = (q_len * (q_len - 1) // 2) + self.compute_num_context_mask(kv_len, context, q_len) - num_tokens_masked = bsz * num_tokens_masked - - if 0 not in mask_2d: - assert (mask_4d != 0).sum().item() == num_tokens_masked - if 0 in mask_2d: - # at least causal mask + maybe more - assert (mask_4d != 0).sum().item() >= num_tokens_masked - self.check_non_causal(bsz, q_len, kv_len, mask_2d, mask_4d) - - def check_to_causal(self, mask_converter, q_len, kv_len, bsz=3): - mask_4d = mask_converter.to_causal_4d( - bsz, query_length=q_len, key_value_length=kv_len, device=torch_device, dtype=torch.float32 - ) - - if q_len == 1 and mask_converter.sliding_window is None: - # no causal mask if q_len is 1 - assert mask_4d is None - return - - context = mask_converter.sliding_window - if mask_converter.is_causal and context is None: - # k * (k+1) / 2 tokens are masked in triangualar masks - num_tokens_masked = bsz * (q_len * (q_len - 1) // 2) - - assert (mask_4d != 0).sum().item() == num_tokens_masked - elif not mask_converter.is_causal and context is None: - assert (mask_4d != 0).sum().item() == 0 - elif mask_converter.is_causal and context is not None: - # k * (k+1) / 2 tokens are masked in triangualar masks - num_tokens_masked = (q_len * (q_len - 1) // 2) + self.compute_num_context_mask(kv_len, context, q_len) - num_tokens_masked = bsz * num_tokens_masked - - assert (mask_4d != 0).sum().item() == num_tokens_masked - - def compute_num_context_mask(self, kv_len, context, q_len): - # This function computes the # of attention tokens that are added for - # the sliding window - c_mask_len = kv_len - context - 1 - num_mask_triangle = c_mask_len * (c_mask_len + 1) // 2 - cut_mask_len = max(c_mask_len - q_len, 0) - num_cut_mask = cut_mask_len * (cut_mask_len + 1) // 2 - return num_mask_triangle - num_cut_mask - - def test_2d_to_4d_causal(self): - mask_converter = AttentionMaskConverter(is_causal=True) - - # auto-regressive use case - self.check_to_4d(mask_converter, q_len=1, kv_len=7) - # special auto-regressive case - self.check_to_4d(mask_converter, q_len=3, kv_len=7) - # non auto-regressive case - self.check_to_4d(mask_converter, q_len=7, kv_len=7) - - # same with extra attention masks - self.check_to_4d(mask_converter, q_len=1, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)]) - self.check_to_4d(mask_converter, q_len=3, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)]) - self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)]) - - # check that the mask does not overflow on causal masked tokens - self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 0), (1, 0), (1, 1)]) - - def test_2d_to_4d(self): - mask_converter = AttentionMaskConverter(is_causal=False) - - # non auto-regressive case - self.check_to_4d(mask_converter, q_len=7, kv_len=7) - - # same with extra attention masks - self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)]) - - def test_2d_to_4d_causal_sliding(self): - mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=5) - - # auto-regressive use case - self.check_to_4d(mask_converter, q_len=1, kv_len=7) - # special auto-regressive case - self.check_to_4d(mask_converter, q_len=3, kv_len=7) - # non auto-regressive case - self.check_to_4d(mask_converter, q_len=7, kv_len=7) - - # same with extra attention masks - self.check_to_4d(mask_converter, q_len=1, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)]) - self.check_to_4d(mask_converter, q_len=3, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)]) - self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)]) - - def test_causal_mask(self): - mask_converter = AttentionMaskConverter(is_causal=True) - - # auto-regressive use case - self.check_to_causal(mask_converter, q_len=1, kv_len=7) - # special auto-regressive case - self.check_to_causal(mask_converter, q_len=3, kv_len=7) - # non auto-regressive case - self.check_to_causal(mask_converter, q_len=7, kv_len=7) - - def test_causal_mask_sliding(self): - mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=3) - - # auto-regressive use case - self.check_to_causal(mask_converter, q_len=1, kv_len=7) - # special auto-regressive case - self.check_to_causal(mask_converter, q_len=3, kv_len=7) - # non auto-regressive case - self.check_to_causal(mask_converter, q_len=7, kv_len=7) - - @pytest.mark.torch_compile_test - def test_torch_compile_fullgraph(self): - model = Prepare4dCausalAttentionMaskModel() - - inputs_embeds = torch.rand([1, 3, 32]) - res_non_compiled = model(inputs_embeds) - - compiled_model = torch.compile(model, fullgraph=True) - - res_compiled = compiled_model(inputs_embeds) - - self.assertTrue(torch.equal(res_non_compiled, res_compiled)) - - model = Create4dCausalAttentionMaskModel() - - inputs_embeds = torch.rand(2, 4, 16) - res_non_compiled = model(inputs_embeds) - - compiled_model = torch.compile(model, fullgraph=True) - res_compiled = compiled_model(inputs_embeds) - - self.assertTrue(torch.equal(res_non_compiled, res_compiled)) - - model = Prepare4dAttentionMaskModel() - - mask = torch.ones(2, 4) - mask[0, :2] = 0 - inputs_embeds = torch.rand(2, 4, 16) - - res_non_compiled = model(mask, inputs_embeds) - - compiled_model = torch.compile(model, fullgraph=True) - res_compiled = compiled_model(mask, inputs_embeds) - - self.assertTrue(torch.equal(res_non_compiled, res_compiled)) - - @require_torch - @slow - def test_unmask_unattended_left_padding(self): - attention_mask = torch.Tensor([[0, 0, 1], [1, 1, 1], [0, 1, 1]]).to(torch.int64) - - expanded_mask = torch.Tensor( - [ - [[[0, 0, 0], [0, 0, 0], [0, 0, 1]]], - [[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], - [[[0, 0, 0], [0, 1, 0], [0, 1, 1]]], - ] - ).to(torch.int64) - - reference_output = torch.Tensor( - [ - [[[1, 1, 1], [1, 1, 1], [0, 0, 1]]], - [[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], - [[[1, 1, 1], [0, 1, 0], [0, 1, 1]]], - ] - ).to(torch.int64) - - result = AttentionMaskConverter._unmask_unattended(expanded_mask, attention_mask, unmasked_value=1) - - self.assertTrue(torch.equal(result, reference_output)) - - attention_mask = torch.Tensor([[0, 0, 1, 1, 1], [1, 1, 1, 1, 1], [0, 1, 1, 1, 1]]).to(torch.int64) - - attn_mask_converter = AttentionMaskConverter(is_causal=True) - past_key_values_length = 0 - key_value_length = attention_mask.shape[-1] + past_key_values_length - - expanded_mask = attn_mask_converter.to_4d( - attention_mask, attention_mask.shape[-1], key_value_length=key_value_length, dtype=torch.float32 - ) - - result = AttentionMaskConverter._unmask_unattended(expanded_mask, attention_mask, unmasked_value=0) - min_inf = torch.finfo(torch.float32).min - reference_output = torch.Tensor( - [ - [ - [ - [0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - [min_inf, min_inf, 0, min_inf, min_inf], - [min_inf, min_inf, 0, 0, min_inf], - [min_inf, min_inf, 0, 0, 0], - ] - ], - [ - [ - [0, min_inf, min_inf, min_inf, min_inf], - [0, 0, min_inf, min_inf, min_inf], - [0, 0, 0, min_inf, min_inf], - [0, 0, 0, 0, min_inf], - [0, 0, 0, 0, 0], - ] - ], - [ - [ - [0, 0, 0, 0, 0], - [min_inf, 0, min_inf, min_inf, min_inf], - [min_inf, 0, 0, min_inf, min_inf], - [min_inf, 0, 0, 0, min_inf], - [min_inf, 0, 0, 0, 0], - ] - ], - ] - ) - - self.assertTrue(torch.equal(reference_output, result)) - - @require_torch - @slow - def test_unmask_unattended_right_padding(self): - attention_mask = torch.Tensor([[1, 1, 1, 0], [1, 1, 1, 1], [1, 1, 0, 0]]).to(torch.int64) - - attn_mask_converter = AttentionMaskConverter(is_causal=True) - past_key_values_length = 0 - key_value_length = attention_mask.shape[-1] + past_key_values_length - - expanded_mask = attn_mask_converter.to_4d( - attention_mask, attention_mask.shape[-1], key_value_length=key_value_length, dtype=torch.float32 - ) - - result = AttentionMaskConverter._unmask_unattended(expanded_mask, attention_mask, unmasked_value=0) - - self.assertTrue(torch.equal(expanded_mask, result)) - - @require_torch - @slow - def test_unmask_unattended_random_mask(self): - attention_mask = torch.Tensor([[1, 0, 1, 0], [1, 0, 1, 1], [1, 1, 0, 1]]).to(torch.int64) - - attn_mask_converter = AttentionMaskConverter(is_causal=True) - past_key_values_length = 0 - key_value_length = attention_mask.shape[-1] + past_key_values_length - - expanded_mask = attn_mask_converter.to_4d( - attention_mask, attention_mask.shape[-1], key_value_length=key_value_length, dtype=torch.float32 - ) - - result = AttentionMaskConverter._unmask_unattended(expanded_mask, attention_mask, unmasked_value=0) - - self.assertTrue(torch.equal(expanded_mask, result)) - - @require_torch class TestAttentionImplementation(unittest.TestCase): @unittest.skip("Just a bit annoying") diff --git a/utils/check_inits.py b/utils/check_inits.py index 3d0525f8015e..b07e5d50d964 100644 --- a/utils/check_inits.py +++ b/utils/check_inits.py @@ -307,7 +307,6 @@ def get_transformers_submodules() -> list[str]: IGNORE_SUBMODULES = [ "convert_pytorch_checkpoint_to_tf2", "models.esm.openfold_utils", - "modeling_attn_mask_utils", "safetensors_conversion", "modeling_gguf_pytorch_utils", "kernels.falcon_mamba", diff --git a/utils/check_repo.py b/utils/check_repo.py index 598ae9e90a36..b554645e3069 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -112,11 +112,21 @@ "PeAudioPreTrainedModel", "PeAudioVideoPreTrainedModel", "PeVideoPreTrainedModel", - "CLIPTextTransformer", # was not a PreTrainedModel originally but needs to be - "CLIPVisionTransformer", # was not a PreTrainedModel originally but needs to be - "MetaClip2TextTransformer", # was not a PreTrainedModel originally but needs to be - "MetaClip2VisionTransformer", # was not a PreTrainedModel originally but needs to be - "MLCDVisionTransformer", # was not a PreTrainedModel originally but needs to be + # the following models should have been PreTrainedModels + "Owlv2TextTransformer", + "Owlv2VisionTransformer", + "OwlViTTextTransformer", + "OwlViTVisionTransformer", + "XCLIPTextTransformer", + "CLIPSegTextTransformer", + "DetrDecoder", + "GroupViTTextTransformer", + "CLIPTextTransformer", + "CLIPVisionTransformer", + "MetaClip2TextTransformer", + "MetaClip2VisionTransformer", + "MLCDVisionTransformer", + # end of should have beens ] # Update this list for models that are not tested with a comment explaining the reason it should not be.