cyril review

Author: manueldeprada

src/transformers/models/minimax/modular_minimax.py

@@ -378,10 +378,7 @@ def forward(
 
 
 class MiniMaxAttention(MixtralAttention):
-    def __init__(self, config: MiniMaxConfig, layer_idx: int):
-        super().__init__(config, layer_idx)
-        del is_sliding  # noqa: F821
-        self.sliding_window = getattr(config, "sliding_window", None)
+    pass
 
 
 class MiniMaxSparseMoeBlock(MixtralSparseMoeBlock):

src/transformers/models/mistral/modeling_mistral.py

@@ -15,7 +15,7 @@
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
 from ...integrations import use_kernel_forward_from_hub
-from ...masking_utils import create_masks_for_generate
+from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import (
     GenericForQuestionAnswering,
@@ -135,13 +135,7 @@ def __init__(self, config: MistralConfig, layer_idx: int):
         self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
         self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
         self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
-        # This check is necessary to support models that inherit via modular (e.g. Mixtral) and do not use layer_types
-        is_sliding = (
-            config.layer_types[layer_idx] == "sliding_attention"
-            if getattr(config, "layer_types", None) is not None
-            else getattr(config, "sliding_window", None) is not None
-        )
-        self.sliding_window = config.sliding_window if is_sliding else None
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
 
     def forward(
         self,
@@ -217,9 +211,7 @@ def __init__(self, config: MistralConfig, layer_idx: int):
         self.mlp = MistralMLP(config)
         self.input_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.attention_type = (
-            config.layer_types[layer_idx] if getattr(config, "layer_types", None) is not None else None
-        )
+        self.attention_type = config.layer_types[layer_idx]
 
     def forward(
         self,
@@ -358,31 +350,35 @@ def forward(
             position_ids = cache_position.unsqueeze(0)
 
         # It may already have been prepared by e.g. `generate`
-        mask_already_prepared = isinstance(attention_mask, dict) or (
-            isinstance(attention_mask, torch.Tensor) and len(attention_mask.shape) > 2
-        )
-        if not mask_already_prepared:
-            attention_mask = create_masks_for_generate(
-                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,
-            )
-
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "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,
+            }
+            full_mask_already_prepared = isinstance(attention_mask, torch.Tensor) and len(attention_mask.shape) > 2
+            causal_mask_mapping = {}
+            if "sliding_attention" in self.config.layer_types:
+                sliding_attention_mask = (
+                    create_sliding_window_causal_mask(**mask_kwargs)
+                    if not full_mask_already_prepared
+                    else attention_mask
+                )
+                causal_mask_mapping["sliding_attention"] = sliding_attention_mask
+            if "full_attention" in self.config.layer_types:
+                causal_mask = create_causal_mask(**mask_kwargs) if not full_mask_already_prepared else attention_mask
+                causal_mask_mapping["full_attention"] = causal_mask
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
-            causal_mask = (
-                attention_mask[decoder_layer.attention_type]
-                if decoder_layer.attention_type is not None and isinstance(attention_mask, dict)
-                else attention_mask
-            )
             hidden_states = decoder_layer(
                 hidden_states,
-                attention_mask=causal_mask,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
                 position_ids=position_ids,
                 past_key_value=past_key_values,
                 use_cache=use_cache,

src/transformers/models/mistral/modular_mistral.py

@@ -6,7 +6,7 @@
 from transformers.utils.generic import check_model_inputs
 
 from ...cache_utils import Cache, DynamicCache
-from ...masking_utils import create_masks_for_generate
+from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import (
     GenericForQuestionAnswering,
@@ -49,13 +49,7 @@ def __init__(self, config: MistralConfig, layer_idx: int):
         self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
         self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
         self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
-        # This check is necessary to support models that inherit via modular (e.g. Mixtral) and do not use layer_types
-        is_sliding = (
-            config.layer_types[layer_idx] == "sliding_attention"
-            if getattr(config, "layer_types", None) is not None
-            else getattr(config, "sliding_window", None) is not None
-        )
-        self.sliding_window = config.sliding_window if is_sliding else None
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
 
     def forward(
         self,
@@ -107,9 +101,7 @@ def __init__(self, config: MistralConfig, layer_idx: int):
         super().__init__(config, layer_idx)
         self.self_attn = MistralAttention(config=config, layer_idx=layer_idx)
         self.mlp = MistralMLP(config)
-        self.attention_type = (
-            config.layer_types[layer_idx] if getattr(config, "layer_types", None) is not None else None
-        )
+        self.attention_type = config.layer_types[layer_idx]
 
 
 class MistralPreTrainedModel(LlamaPreTrainedModel):
@@ -152,31 +144,35 @@ def forward(
             position_ids = cache_position.unsqueeze(0)
 
         # It may already have been prepared by e.g. `generate`
-        mask_already_prepared = isinstance(attention_mask, dict) or (
-            isinstance(attention_mask, torch.Tensor) and len(attention_mask.shape) > 2
-        )
-        if not mask_already_prepared:
-            attention_mask = create_masks_for_generate(
-                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,
-            )
-
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "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,
+            }
+            full_mask_already_prepared = isinstance(attention_mask, torch.Tensor) and len(attention_mask.shape) > 2
+            causal_mask_mapping = {}
+            if "sliding_attention" in self.config.layer_types:
+                sliding_attention_mask = (
+                    create_sliding_window_causal_mask(**mask_kwargs)
+                    if not full_mask_already_prepared
+                    else attention_mask
+                )
+                causal_mask_mapping["sliding_attention"] = sliding_attention_mask
+            if "full_attention" in self.config.layer_types:
+                causal_mask = create_causal_mask(**mask_kwargs) if not full_mask_already_prepared else attention_mask
+                causal_mask_mapping["full_attention"] = causal_mask
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
-            causal_mask = (
-                attention_mask[decoder_layer.attention_type]
-                if decoder_layer.attention_type is not None and isinstance(attention_mask, dict)
-                else attention_mask
-            )
             hidden_states = decoder_layer(
                 hidden_states,
-                attention_mask=causal_mask,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
                 position_ids=position_ids,
                 past_key_value=past_key_values,
                 use_cache=use_cache,

src/transformers/models/mixtral/modular_mixtral.py

@@ -225,7 +225,6 @@ class MixtralRMSNorm(MistralRMSNorm):
 class MixtralAttention(MistralAttention):
     def __init__(self, config: MixtralConfig, layer_idx: int):
         super().__init__(config, layer_idx)
-        del is_sliding  # noqa: F821
         self.sliding_window = getattr(config, "sliding_window", None)
 
 

src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py

@@ -35,6 +35,7 @@
 from ...processing_utils import Unpack
 from ...utils import auto_docstring, logging
 from ...utils.generic import TransformersKwargs, check_model_inputs
+from ..phi3.configuration_phi3 import Phi3Config
 from ..phi3.modeling_phi3 import (
     Phi3DecoderLayer,
     Phi3ForCausalLM,
@@ -277,7 +278,7 @@ def __init__(
         self.nemo_final_size = length
 
 
-class Phi4MultimodalConfig(PretrainedConfig):
+class Phi4MultimodalConfig(Phi3Config):
     r"""
     This is the configuration class to store the configuration of a [`Phi4MultimodalModel`]. It is used to instantiate a
     Phi4Multimodal model according to the specified arguments, defining the model architecture. Instantiating a configuration
@@ -370,20 +371,6 @@ class Phi4MultimodalConfig(PretrainedConfig):
     >>> configuration = model.config
     ```"""
 
-    model_type = "phi4_multimodal"
-    keys_to_ignore_at_inference = ["past_key_values"]
-    base_model_tp_plan = {
-        "layers.*.self_attn.qkv_proj": "colwise_rep",  # we need to replicate here due to the slicing of qkv
-        "layers.*.self_attn.o_proj": "rowwise_rep",  # we need to replicate here due to the slicing of qkv
-        "layers.*.mlp.gate_up_proj": "colwise_rep",  # we need to replicate here due to the `chunk` operation
-        "layers.*.mlp.down_proj": "rowwise_rep",  # we need to replicate here due to the `chunk` operation
-    }
-    base_model_pp_plan = {
-        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
-        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
-        "norm": (["hidden_states"], ["hidden_states"]),
-    }
-
     sub_configs = {"audio_config": Phi4MultimodalAudioConfig, "vision_config": Phi4MultimodalVisionConfig}
 
     def __init__(
@@ -416,37 +403,31 @@ def __init__(
         **kwargs,
     ):
         super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            resid_pdrop=resid_pdrop,
+            embd_pdrop=embd_pdrop,
+            attention_dropout=attention_dropout,
+            hidden_act=hidden_act,
+            max_position_embeddings=max_position_embeddings,
+            initializer_range=initializer_range,
+            rms_norm_eps=rms_norm_eps,
+            use_cache=use_cache,
+            tie_word_embeddings=tie_word_embeddings,
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            partial_rotary_factor=partial_rotary_factor,
             bos_token_id=bos_token_id,
             eos_token_id=eos_token_id,
             pad_token_id=pad_token_id,
-            tie_word_embeddings=tie_word_embeddings,
+            original_max_position_embeddings=original_max_position_embeddings,
+            sliding_window=sliding_window,
             **kwargs,
         )
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.resid_pdrop = resid_pdrop
-        self.embd_pdrop = embd_pdrop
-        self.attention_dropout = attention_dropout
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.original_max_position_embeddings = original_max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.partial_rotary_factor = partial_rotary_factor
-        self._rope_scaling_adjustment()
-        self._rope_scaling_validation()
-        self.sliding_window = sliding_window
 
         if isinstance(vision_config, dict):
             vision_config = Phi4MultimodalVisionConfig(**vision_config)
@@ -460,60 +441,6 @@ def __init__(
             audio_config = Phi4MultimodalAudioConfig()
         self.audio_config = audio_config
 
-    def _rope_scaling_adjustment(self):
-        """
-        Adjust the `type` of the `rope_scaling` configuration for backward compatibility.
-        """
-        if self.rope_scaling is None:
-            return
-
-        rope_scaling_type = self.rope_scaling.get("type", None)
-
-        # For backward compatibility if previous version used "su" or "yarn"
-        if rope_scaling_type is not None and rope_scaling_type in ["su", "yarn"]:
-            self.rope_scaling["type"] = "longrope"
-
-    def _rope_scaling_validation(self):
-        """
-        Validate the `rope_scaling` configuration.
-        """
-        if self.rope_scaling is None:
-            return
-
-        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 3:
-            raise ValueError(
-                "`rope_scaling` must be a dictionary with three fields, `type`, `short_factor` and `long_factor`, "
-                f"got {self.rope_scaling}"
-            )
-        rope_scaling_type = self.rope_scaling.get("type", None)
-        rope_scaling_short_factor = self.rope_scaling.get("short_factor", None)
-        rope_scaling_long_factor = self.rope_scaling.get("long_factor", None)
-        if rope_scaling_type is None or rope_scaling_type not in ["longrope"]:
-            raise ValueError(f"`rope_scaling`'s type field must be one of ['longrope'], got {rope_scaling_type}")
-        if not (
-            isinstance(rope_scaling_short_factor, list)
-            and all(isinstance(x, (int, float)) for x in rope_scaling_short_factor)
-        ):
-            raise ValueError(
-                f"`rope_scaling`'s short_factor field must be a list of numbers, got {rope_scaling_short_factor}"
-            )
-        rotary_ndims = int(self.hidden_size // self.num_attention_heads * self.partial_rotary_factor)
-        if not len(rope_scaling_short_factor) == rotary_ndims // 2:
-            raise ValueError(
-                f"`rope_scaling`'s short_factor field must have length {rotary_ndims // 2}, got {len(rope_scaling_short_factor)}"
-            )
-        if not (
-            isinstance(rope_scaling_long_factor, list)
-            and all(isinstance(x, (int, float)) for x in rope_scaling_long_factor)
-        ):
-            raise ValueError(
-                f"`rope_scaling`'s long_factor field must be a list of numbers, got {rope_scaling_long_factor}"
-            )
-        if not len(rope_scaling_long_factor) == rotary_ndims // 2:
-            raise ValueError(
-                f"`rope_scaling`'s long_factor field must have length {rotary_ndims // 2}, got {len(rope_scaling_long_factor)}"
-            )
-
 
 class Phi4MultimodalVisionMLP(SiglipMLP):
     pass

src/transformers/models/starcoder2/modular_starcoder2.py

@@ -77,7 +77,6 @@ def __init__(self, config: Starcoder2Config, layer_idx: Optional[int] = None):
         self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias)
         self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.use_bias)
         self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.use_bias)
-        del is_sliding  # noqa: F821
         del self.sliding_window
 
     def forward(