Merge branch 'main' into fix_biogpt_test1

Author: Sai-Suraj-27

src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py

@@ -1705,6 +1705,8 @@ def prepare_inputs_for_generation(
         past_key_values=None,
         image_grid_thw=None,
         video_grid_thw=None,
+        use_cache=True,
+        is_first_iteration=False,
         # Intentionally ignore position ids to force custom cache logic
         position_ids=None,
         **kwargs,
@@ -1717,6 +1719,8 @@ def prepare_inputs_for_generation(
             past_key_values=past_key_values,
             image_grid_thw=image_grid_thw,
             video_grid_thw=video_grid_thw,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
@@ -1732,7 +1736,7 @@ def prepare_inputs_for_generation(
             mm_token_type_ids=model_inputs.get("mm_token_type_ids"),
         )
 
-        if model_inputs["cache_position"][0] != 0:
+        if not is_first_iteration and use_cache:
             model_inputs["pixel_values"] = None
             model_inputs["pixel_values_videos"] = None
             model_inputs["mm_token_type_ids"] = None

src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py

@@ -1393,6 +1393,8 @@ def prepare_inputs_for_generation(
         past_key_values=None,
         image_grid_thw=None,
         video_grid_thw=None,
+        use_cache=True,
+        is_first_iteration=False,
         # Intentionally ignore position ids to force custom cache logic
         position_ids=None,
         **kwargs,
@@ -1405,6 +1407,8 @@ def prepare_inputs_for_generation(
             past_key_values=past_key_values,
             image_grid_thw=image_grid_thw,
             video_grid_thw=video_grid_thw,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
@@ -1420,7 +1424,7 @@ def prepare_inputs_for_generation(
             mm_token_type_ids=model_inputs.get("mm_token_type_ids"),
         )
 
-        if model_inputs["cache_position"][0] != 0:
+        if not is_first_iteration and use_cache:
             model_inputs["pixel_values"] = None
             model_inputs["pixel_values_videos"] = None
             model_inputs["mm_token_type_ids"] = None

src/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -88,8 +88,8 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
 
         Args:
             hidden_states (torch.Tensor): (batch_size, seq_len, hidden_size)
-            selected_experts (torch.Tensor): (batch_size * token_num, top_k)
-            routing_weights (torch.Tensor): (batch_size * token_num, num_experts)
+            selected_experts (torch.Tensor): (batch_size * seq_len, top_k)
+            routing_weights (torch.Tensor): (batch_size * seq_len, top_k)
         Returns:
             torch.Tensor
         """
@@ -159,8 +159,8 @@ def __init__(self, config):
 
     def forward(self, hidden_states):
         hidden_states = hidden_states.reshape(-1, self.hidden_dim)
-        router_logits = F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
-        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
+        router_logits = F.linear(hidden_states, self.weight, self.bias)  # (num_tokens, num_experts)
+        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (num_tokens, top_k)
         router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
         router_scores = router_top_value
         return router_logits, router_scores, router_indices
@@ -434,7 +434,7 @@ class GptOssPreTrainedModel(PreTrainedModel):
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn = True
     _supports_sdpa = False
-    _supports_flex_attn = False
+    _supports_flex_attn = True
 
     _can_compile_fullgraph = True
     _supports_attention_backend = True

src/transformers/models/gpt_oss/modular_gpt_oss.py

@@ -86,8 +86,8 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
 
         Args:
             hidden_states (torch.Tensor): (batch_size, seq_len, hidden_size)
-            selected_experts (torch.Tensor): (batch_size * token_num, top_k)
-            routing_weights (torch.Tensor): (batch_size * token_num, num_experts)
+            selected_experts (torch.Tensor): (batch_size * seq_len, top_k)
+            routing_weights (torch.Tensor): (batch_size * seq_len, top_k)
         Returns:
             torch.Tensor
         """
@@ -157,8 +157,8 @@ def __init__(self, config):
 
     def forward(self, hidden_states):
         hidden_states = hidden_states.reshape(-1, self.hidden_dim)
-        router_logits = F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
-        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
+        router_logits = F.linear(hidden_states, self.weight, self.bias)  # (num_tokens, num_experts)
+        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (num_tokens, top_k)
         router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
         router_scores = router_top_value
         return router_logits, router_scores, router_indices
@@ -354,7 +354,6 @@ def forward(
 class GptOssPreTrainedModel(LlamaPreTrainedModel):
     _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"]
     _supports_sdpa = False
-    _supports_flex_attn = False
     _can_record_outputs = {
         "router_logits": OutputRecorder(GptOssTopKRouter, index=0),
         "hidden_states": GptOssDecoderLayer,

tests/models/ernie4_5_vl_moe/test_modeling_ernie4_5_vl_moe.py

@@ -313,6 +313,7 @@ def load_model(self, dtype, attn_implementation="sdpa"):
             device_map="auto",
             dtype=dtype,
             attn_implementation=attn_implementation,
+            experts_implementation="eager",
             revision="refs/pr/10",
         )
 
@@ -549,6 +550,7 @@ def load_model(self, dtype, attn_implementation="sdpa"):
             device_map="auto",
             dtype=dtype,
             attn_implementation=attn_implementation,
+            experts_implementation="eager",
         )
 
     def test_small_model_integration_test(self):