Unnecessary transpose in switch transformer implementation
Author: karan-uppal3
Pull request: no
Labels: name: Core: Modeling

name: bug

### System Info

System Info

- `transformers` version: 4.44.2
- Platform: Linux-6.5.0-35-generic-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.24.6
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config: not found
- PyTorch version (GPU?): 2.3.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: No
- Using GPU in script?: No

### Who can help?

@ArthurZucker @younesbelkada 

There is an additional transpose operation at [link](https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L300) in the Switch Transformer Implementation. The shape of the tensor is `[batch_size, seq_len, num_experts]` and then `.transpose(1,2)` to get the shape `[batch_size, num_experts, seq_len]`, which is later `.reshape(batch_size * seq_len, num_experts)` which is incompatible and gives incorrect result.

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python3
import torch.nn as nn
from transformers import SwitchTransformersConfig

class SwitchTransformersTop1RouterMy(nn.Module):
    
    def __init__(self, config: SwitchTransformersConfig):
        super().__init__()
        self.num_experts = config.num_experts
        self.expert_capacity = config.expert_capacity
        self.classifier = nn.Linear(config.hidden_size, self.num_experts, bias=config.router_bias)
        self.jitter_noise = config.router_jitter_noise
        self.ignore_padding_tokens = config.router_ignore_padding_tokens
        self.dtype = getattr(torch, config.router_dtype)

    def forward(self, router_probs: torch.Tensor):

        expert_index = torch.argmax(router_probs, dim=-1)
        expert_index = torch.nn.functional.one_hot(expert_index, num_classes=self.num_experts)

        token_priority = torch.cumsum(expert_index, dim=-2)
        expert_capacity_mask = token_priority <= self.expert_capacity
        expert_index = expert_index * expert_capacity_mask

        router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
        return expert_index, router_probs

class SwitchTransformersSparseMLPMy(torch.nn.Module):
    
    def __init__(self, config):
        super().__init__()

    def forward(self, router_mask):

        expert_index = torch.argmax(router_mask, dim=-1)

        router_mask = router_mask.bool()
        batch_size, seq_len, num_experts = router_mask.shape
        idx_mask = router_mask.transpose(1, 2).reshape(batch_size * seq_len, num_experts).sum(dim=0)
        idx_mask = torch.nonzero(idx_mask, as_tuple=True)[
            0
        ].tolist()  # length: number of "activated" expert / value: index
        
        return idx_mask

config = SwitchTransformersConfig(
    expert_capacity = 2,
    num_experts = 8
)
router = SwitchTransformersTop1RouterMy(config)
sparse_mlp = SwitchTransformersSparseMLPMy(config)

in_tmp = torch.tensor([[
    [1, 0, 0],
    [0, 0, 1],
    [0, 1, 0]
]])
expert_index, router_probs = router(in_tmp)
out = sparse_mlp(expert_index)

print(out)
```

The output is 
```
[0, 5, 7]
```
which is incorrect for the input given. 

### Expected behavior

Replacing `idx_mask = router_mask.transpose(1, 2).reshape(batch_size * seq_len, num_experts).sum(dim=0)` with `idx_mask = router_mask.reshape(batch_size * seq_len, num_experts).sum(dim=0)` gives the output:

```
[0, 1, 2]
```

Comments:
Hey @karan-uppal3, would you like to open a PR to fix this? 

Sure, can you help me out? This is my first pull request

Absolutely! We wrote a guide to help opening first PRs, you can find it here: https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request