Incorrect router probability calculation
Author: lhallee
Pull request: no

### System Info

transformers version 4.36.0

### Who can help?

 @ArthurZucker and @younesbelkada

### Information

- [X] 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

I think load_balancing_loss_func in modeling_mixtral creates router_prob_per_group_and_expert incorrectly
https://github.com/huggingface/transformers/blob/v4.36.0/src/transformers/models/mixtral/modeling_mixtral.py#L120

Trying to multiply something batch_size * num_hidden_layers, num_experts by batch_size * num_hidden_layers, topk, 1

`torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert.unsqueeze(-1)) * (num_experts**2)`

Correct creation of routing_weights should likely be from gate_logits, which ensures it is the correct size

`routing_weights = gate_logits.softamx(dim=-1)`

The unsqueeze(-1) is necessary with this. Also router_prob_per_group_and_expert should average over axis=-2

`router_prob_per_group_and_expert = torch.mean(routing_weights, axis=-2)`

This follows the previous implementation in modeling_switch_transformers
https://github.com/huggingface/transformers/blob/v4.36.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L91

### Expected behavior

Something like this would fix it

```
def router_loss_func_test(gate_logits: torch.Tensor, top_k=2) -> float:
    if gate_logits is None:
        return 0

    if isinstance(gate_logits, tuple):
        # cat along the layers?
        gate_logits = torch.cat(gate_logits, dim=0) # batch_size * num_hidden_layers, sequence_length, num_experts

    num_experts = gate_logits.shape[-1]

    _, expert_indicies = torch.topk(gate_logits, top_k, dim=-1)  # this is done so you don't need to pass expert_indicies
    routing_probs = gate_logits.softmax(dim=-1) # routing probs

    if expert_indicies.dtype != torch.int64: # cast the expert indices to int64, otherwise one-hot encoding will fail
        expert_indicies = expert_indicies.to(torch.int64)

    if len(expert_indicies.shape) == 2:
        expert_indicies = expert_indicies.unsqueeze(2)

    expert_mask = torch.nn.functional.one_hot(expert_indicies, num_experts)

    # For a given token, determine if it was routed to a given expert.
    expert_mask = torch.max(expert_mask, axis=-2).values

    expert_mask = expert_mask.to(torch.float32) # cast to float32 otherwise mean will fail
    tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)

    router_prob_per_group_and_expert = torch.mean(routing_probs, axis=-2)

    loss = torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)
    return loss
```

Comments:
Sorry could you either show the issue or detail where you had a problem? The computation is different because the output shape are also different, the routing mecanism is also different. 🤗  

Sure! @ArthurZucker

Here's the current loss function for convenience
```
def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2) -> float:
    r"""
    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.

    See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
    experts is too unbalanced.

    Args:
        gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
            Logits from the `gate`, should be a tuple of tensors. Shape: [batch_size, seqeunce_length, num_experts].
        num_experts (`int`, *optional*):
            Number of experts

    Returns:
        The auxiliary loss.
    """
    if gate_logits is None:
        return 0

    if isinstance(gate_logits, tuple):
        # cat along the layers?
        gate_logits = torch.cat(gate_logits, dim=0)

    routing_weights, selected_experts = torch.topk(gate_logits, top_k, dim=-1)
    routing_weights = routing_weights.softmax(dim=-1)

    # cast the expert indices to int64, otherwise one-hot encoding will fail
    if selected_experts.dtype != torch.int64:
        selected_experts = selected_experts.to(torch.int64)

    if len(selected_experts.shape) == 2:
        selected_experts = selected_experts.unsqueeze(2)

    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)

    # For a given token, determine if it was routed to a given expert.
    expert_mask = torch.max(expert_mask, axis=-2).values

    # cast to float32 otherwise mean will fail
    expert_mask = expert_mask.to(torch.float32)
    tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)

    router_prob_per_group_and_expert = torch.mean(routing_weights, axis=-1)
    return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert.unsqueeze(-1)) * (num_experts**2)
```

An example

```
num_hidden_layers=30
batch_size = 16
seq_len = 32
num_experts = 8
gate_logits = tuple(torch.randn(batch_size, seq_len, num_experts) for _ in range(num_hidden_layers))
load_balancing_loss_func(gate_logits=gate_logits, num_experts=num_experts)
```
Shape error

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[c:\Users\Logan](file:///C:/Users/Logan) Hallee\Desktop\MOE-PLM\moesm_testing.ipynb Cell 13 line 6
      [3](vscode-notebook-cell:/c%3A/Users/Logan%20Hallee/Desktop/MOE-PLM/moesm_testing.ipynb#X15sZmlsZQ%3D%3D?line=2) num_experts = 8
      [5](vscode-notebook-cell:/c%3A/Users/Logan%20Hallee/Desktop/MOE-PLM/moesm_testing.ipynb#X15sZmlsZQ%3D%3D?line=4) gate_logits = tuple(torch.randn(batch_size, seq_len, num_experts) for _ in range(30))
----> [6](vscode-notebook-cell:/c%3A/Users/Logan%20Hallee/Desktop/MOE-PLM/moesm_testing.ipynb#X15sZmlsZQ%3D%3D?line=5) load_balancing_loss_func(gate_logits=gate_logits, num_experts=8)

[c:\Users\Logan](file:///C:/Users/Logan) Hallee\Desktop\MOE-PLM\moesm_testing.ipynb Cell 13 line 4
     [42](vscode-notebook-cell:/c%3A/Users/Logan%20Hallee/Desktop/MOE-PLM/moesm_testing.ipynb#X15sZmlsZQ%3D%3D?line=41) tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
     [44](vscode-notebook-cell:/c%3A/Users/Logan%20Hallee/Desktop/MOE-PLM/moesm_testing.ipynb#X15sZmlsZQ%3D%3D?line=43) router_prob_per_group_and_expert = torch.mean(routing_weights, axis=-1)
---> [45](vscode-notebook-cell:/c%3A/Users/Logan%20Hallee/Desktop/MOE-PLM/moesm_testing.ipynb#X15sZmlsZQ%3D%3D?line=44) return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert.unsqueeze(-1)) * (num_experts**2)

RuntimeError: The size of tensor a (480) must match the size of tensor b (32) at non-singleton dimension 1
```
 

The loss is made to be used with the outputs of the model, which merge batch and sequence length 😉

It looks like the documentation is wrong then. Could you clarify where the merge happens and the correct shape of the input?

Hello~ does this function "load_balancing_loss_func" really work？ It  always output a constant for me.

> Hello~ does this function "load_balancing_loss_func" really work？ It always output a constant for me.

Same to me, and the grad norm is 0. @ArthurZucker 

Thanks all for the feedback I'll check it and update the doc with an example! 
The merge happens in the forward of the `MixtralSparseMoeBlock` here: https://github.com/huggingface/transformers/blob/cfd3e8d1e05e11b12bf50efb90691a4ad1f68926/src/transformers/models/mixtral/modeling_mixtral.py#L706

> Thanks all for the feedback I'll check it and update the doc with an example! The merge happens in the forward of the `MixtralSparseMoeBlock` here: https://github.com/huggingface/transformers/blob/cfd3e8d1e05e11b12bf50efb90691a4ad1f68926/src/transformers/models/mixtral/modeling_mixtral.py#L706

Hi, have you fix the constant loss problem ?

Yes, #28115 fixes this sorry everyone ! 🤗