fix gpt-oss out shape (#40535)

* fix out shape

Signed-off-by: jiqing-feng <[email protected]>

* reset gpt-oss modeling

Signed-off-by: jiqing-feng <[email protected]>

* fix copies

Signed-off-by: jiqing-feng <[email protected]>

* fix tests

Signed-off-by: jiqing-feng <[email protected]>

---------

Signed-off-by: jiqing-feng <[email protected]>
Co-authored-by: Marc Sun <[email protected]>

Author: jiqing-feng

src/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -116,7 +116,7 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
                 glu = gate * torch.sigmoid(gate * self.alpha)
                 gated_output = (up + 1) * glu
                 out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
-                weighted_output = out[0] * routing_weights[token_idx, expert_idx, None]
+                weighted_output = out * routing_weights[token_idx, expert_idx, None]
                 next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
             next_states = next_states.view(batch_size, -1, self.hidden_size)
         else:

src/transformers/models/gpt_oss/modular_gpt_oss.py

@@ -115,7 +115,7 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
                 glu = gate * torch.sigmoid(gate * self.alpha)
                 gated_output = (up + 1) * glu
                 out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
-                weighted_output = out[0] * routing_weights[token_idx, expert_idx, None]
+                weighted_output = out * routing_weights[token_idx, expert_idx, None]
                 next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
             next_states = next_states.view(batch_size, -1, self.hidden_size)
         else:

tests/models/gpt_oss/test_modeling_gpt_oss.py

@@ -128,9 +128,6 @@ def test_flex_attention_with_grads(self):
     def test_generate_compile_model_forward_fullgraph(self):
         return super().test_generate_compile_model_forward_fullgraph()
 
-    def test_batching_equivalence(self, **kwargs):
-        return super().test_batching_equivalence(atol=5e-4, rtol=1e-3)
-
 
 RESULTS_PATH = Path(__file__).parent.parent.parent / "fixtures/gpt_oss/integration_tests.json"