Repository: huggingface/transformers
Pull request: Fix shapes in modular_gpt_oss.py

# What does this PR do?

This PR fixes a comment in `modular_gpt_oss.py` which has the incorrect shape written in the description for GptOssExperts.  I fix the annotated shape for routing_experts from `(batch_size * token_num, num_experts)` to `(batch_size * token_num, top_k)`. Looking at the git blame, I think the original implementation of GPT OSS had the original shape which was later modified but the comment was not edited. I also made a few other edits to the shape annotations for MoE in this file for consistency (assuming `num_tokens = batch_size * seq_len`). Please let me know if the changes make sense!

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@ArthurZucker 

Changed files:
filename: src/transformers/models/gpt_oss/modeling_gpt_oss.py
patch: @@ -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

filename: src/transformers/models/gpt_oss/modular_gpt_oss.py
patch: @@ -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