Incorrect implementation of auxiliary loss 
Author: liangxuZhang
Pull request: no

### System Info

- `transformers` version: 4.37.0.dev0
- Platform: macOS-13.5-arm64-arm-64bit
- Python version: 3.10.13
- Huggingface_hub version: 0.20.1
- Safetensors version: 0.4.1
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.1.2 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker

### 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

Two issues were found: 
## mixtral's implementation of auxiliary loss is not correct.
I think `load_balancing_loss_func` in `modeling_mixtral` computes auxiliary loss incorrectly
https://github.com/huggingface/transformers/blob/3cefac1d974db5e2825a0cb2b842883a628be7a0/src/transformers/models/mixtral/modeling_mixtral.py#L77-L119
Auxiliary loss is implemented as multiply fraction of tokens dispatched to expert by fraction of the router probability allocated for expert. The fraction of tokens dispatched to expert is calculated as the number of tokens routed to  expert divided by the total number of tokens. The actual implementation is as follows: https://github.com/huggingface/transformers/blob/3cefac1d974db5e2825a0cb2b842883a628be7a0/src/transformers/models/mixtral/modeling_mixtral.py#L109-L113
As we know, the shape of `selected_experts` is `top_k X [batch_size X sequence_length]`，so the shape of `expert_mask` is `[top_k X batch_size X sequence_length, num_experts]`. When we excute `expert_mask = torch.max(expert_mask, dim=-2).values` , the shape of the `expert_mask` becomes `[num_experts]`, which means that whenever a token is routed to an expert, that expert has a value of 1. After the operation of `torch.mean(expert_mask.float(), dim=0)`, `tokens_per_expert` becomes a scaler, which is clearly incorrect, since tokens_per_expert should have a shape of `[num_experts]`. 
**Example**
Example Inputs：
```python
T = 3  # number of tokens [B X S]
num_experts = 8 
top_k = 2 # top_2
gate_logits = torch.randn(T,num_experts)
routing_weights = torch.nn.functional.softmax(gate_logits, dim=-1)
```
Each row of `routing_weights` represents the probability that a token will be routed to an expert
```python
tensor([[0.2551, 0.2519, 0.0357, 0.0830, 0.0897, 0.0981, 0.1565, 0.0299],
        [0.0728, 0.0593, 0.0948, 0.1708, 0.0098, 0.0848, 0.3884, 0.1192],
        [0.0292, 0.0387, 0.0696, 0.1331, 0.6699, 0.0049, 0.0442, 0.0104]])
```
next select experts
```python
 _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)

# treat `top_k` as tokens (shape is `top_k X [batch_size X sequence_length]`)
selected_experts = selected_experts.reshape(-1)

expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
```
we get the following result (shape `[top_k X batch_size X sequence_length, num_experts]`):
```python
tensor([[1, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 1, 0],
        [0, 0, 0, 1, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0, 0]])
```
The final results are as follows
```python
expert_mask = torch.max(expert_mask, dim=-2).values
# tensor([1, 1, 0, 1, 1, 0, 1, 0])
tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
# tensor(0.6250)
router_prob_per_expert = torch.mean(routing_weights, dim=0)
# tensor([0.0746, 0.1031, 0.0448, 0.0804, 0.1823, 0.1216, 0.3112, 0.0820])
overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(-1))
# tensor(0.6250) 
```
Because the sum of the `router_prob_per_expert` is 1,  the final loss value is actually the value of `tokens_per_expert`. As the total number of tokens increases, the value of `tokens_per_expert` will be 1 (each expert has tokens routed to). 

**Solution**
The `tokens_per_expert` calculation should divide the tokens that are routed per expert by the total number of tokens. Specifically, we can sum the columns of `expert_mask` and divide by the total number of tokens.  The following is an implementation
```python
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
expert_mask = expert_mask.reshape(-1, top_k, num_experts)
expert_mask = torch.max(expert_mask, dim=-2).values

# Compute the percentage of tokens routed to each experts
tokens_per_expert = torch.mean(expert_mask.float(), dim=0) / top_k

# Compute the average probability of routing to these experts
router_prob_per_expert = torch.mean(routing_weights, dim=0)

overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert)
return overall_loss * num_experts
```
**Example**
```python
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
expert_mask = expert_mask.reshape(T,top_k,-1)
'''
tensor([[[1, 0, 0, 0, 0, 0, 0, 0],
         [0, 1, 0, 0, 0, 0, 0, 0]],
        [[0, 0, 0, 0, 0, 0, 1, 0],
         [0, 0, 0, 1, 0, 0, 0, 0]],
        [[0, 0, 0, 0, 1, 0, 0, 0],
         [0, 0, 0, 1, 0, 0, 0, 0]]])
'''
expert_mask = torch.max(expert_mask, dim=-2).values
'''
tensor([[1, 1, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 1, 0],
        [0, 0, 0, 1, 1, 0, 0, 0]])
'''
tokens_per_expert = torch.mean(expert_mask.float(), dim=0) / top_k
# tensor([0.1667, 0.1667, 0.0000, 0.3333, 0.1667, 0.0000, 0.1667, 0.0000])
```

## Note
On the other hand, in switch transformer ((https://arxiv.org/abs/2101.03961), auxiliary loss  should converge to 1 when the load is balanced. However, the top 1 strategy is used in the paper, so the maximum value is taken when calculating tokens_per_expert. In the top_k strategy, this corresponds to `top_k*T` tokens being routed to the experts, so tokens_per_expert should be divided by `top_k`. Otherwise the final converged value should be `top_k`.

By the way, the unit test should determine if the loss is close to 1 instead of 8.
https://github.com/huggingface/transformers/blob/3cefac1d974db5e2825a0cb2b842883a628be7a0/tests/models/mixtral/test_modeling_mixtral.py#L477

## Should the output of each layer of the gated network be concatenated into a tensor?

https://github.com/huggingface/transformers/blob/3cefac1d974db5e2825a0cb2b842883a628be7a0/src/transformers/models/mixtral/modeling_mixtral.py#L98C5-L101C1 

Before calculating the auxiliary loss, the routing outputs of the different transformer layers of the expert layer are concatenated into a tensor. This implies that the routing outputs of different layers are mapped to the same expert, and in fact the routing outputs of each layer should be mapped to its own layer of experts. So should the auxiliary loss be calculated for each layer independently, rather than concatenated into a tensor？

### Expected behavior

I expect to examine the problem and review my solution for the first issue and have a discussion about the second issue, as I'm not sure if it makes more sense to calculate loss separately.

Comments:
Hi, liangxu! I also noticed this difference between the implementation by mixtral and the equations 4-5 in (https://arxiv.org/pdf/2101.03961.pdf). I am wondering if this has something to do with "The objective can also be differentiated as the P-vector is differentiable, but the f-vector is not." (line1, page8)?

Hey! Thanks you for the deep-dive. 
TLDR: there is still a small bug in our implementation, I'll review the PR, and yes given that the random model is init with uniform distrib the test values should be closer to unifrom distrib. 

However this is wrong : 
>  in fact the routing outputs of each layer should be mapped to its own layer of experts 

The paper states the following: (page 7)
> For each Switch layer, this auxiliary loss is added to the total model loss during training. 

This means that as long as we sum the auxiliary losses of each layers, we should be good to go.

The total auxiliary loss across all layers $\(K\)$ is given by:

$$
\text{total loss} =  \sum_{k=1}^{K} \left( \alpha \cdot N \cdot \sum_{i=1}^{N} f_{i,k} \cdot P_{i,k} \right)
$$

Now, let's factorize the expression and use the associativity of the addition:

$$
\text{total loss} =  \alpha \cdot N \cdot \sum_{i=1}^{N}   \left(   \sum_{k=1}^{K}  f_{i,k} \cdot P_{i,k} \right)
$$

## The actual question
Now your question can basically be summed up as: **how do we deal with top2 vs top1** ?

- **top1 | top2**   either we balance the distribution of top1 + top2 routing, what we are doing. (meaning that top1 and top2 should overall give a uniform repartition, on average, the fraction of tokens dispatched to each expert across both top-1 and top-2 will be balanced)
- **top1 & top2** or we separately balance the distribution of top1 and the distribution of top2 (meaning that top1 should uniformly route to the 8 experts, and top2 should also uniformly route to 8 experts:  the fraction of tokens dispatched to each expert is balanced independently for top-1 and top-2).

Mistral did not really share what they did so I am down to support both (meaning let the user decide what he want to optimize!). 

At the end of the day what is important is that we have:
```python
tokens_per_expert * router_prob_per_expert.unsqueeze(0)
tensor([[0.0180, 0.0179, 0.0157, 0.0152, 0.0117, 0.0244, 0.0053, 0.0172],   # top1 
        [0.0256, 0.0172, 0.0130, 0.0186, 0.0151, 0.0098, 0.0094, 0.0165]],  # top2
       grad_fn=<MulBackward0>)
```
Let's write explicitly the sums,either we want: (**top1 | top2** ) 
```python
torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0), dim=0).sum() * num_experts
```
close to 1. Or we want: (**top1 & top2** ) 
```python
(torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0), dim=0)/top_k).sum() * num_experts
```
to be close to 1. 

I used this implementation:
```python
def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2) -> float:
    if isinstance(gate_logits, tuple):
        compute_device = gate_logits[0].device
        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)

    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) # [batch_size X sequence_length, top_k]
    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) # [batch_size X sequence_length, top_k, num_experts]
    tokens_per_expert = torch.mean(expert_mask.float(), dim=0) # [top_k, num_experts]
    # Compute the average probability of routing to these experts
    router_prob_per_expert = torch.mean(routing_weights, dim=0) # [num_experts]
    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) # / top_k
    return overall_loss * num_experts
```

Mark

I have been following this discussion about the load balancing loss loosely, since I have my own custom training script where I'm trying to add proper support for Mixtral. Isn't the load balancing loss still wrong, since it should be computed per layer and then averaged at the end?

The [paper](https://arxiv.org/abs/2101.03961) states "For each Switch layer, this auxiliary loss is added to the total model loss during training". To me, this seems to imply computing the loss separately per layer, then summing or averaging at the end.

I found [this implementation](https://github.com/google/flaxformer/blob/main/flaxformer/architectures/moe/routing.py) from Google that supports this. It looks like the auxiliary load balancing loss is computed per-layer as a scalar, which must then be combined at some point (I haven't read the code that closely yet).

Intuitively, it also makes sense: you would not want unbalanced expert assignments from one layer to "cancel out" those from another layer, which can happen (and seems likely to happen) if you just concatenate everything across all layers at the beginning, which is what is done now.

@ArthurZucker I know that in those equations you can compute either the sum over experts or the sum over layers first. But, you still need the $f_{i,k}$ per layer. If you concatenate all the gate logits at the beginning, then do a one_hot somewhere and average to get expert assignment fractions, you are losing the per-layer breakdown. So I think we need to explicitly keep everything separated by layer until the very end.

For my custom training script, I implemented the load balancing loss like this, it seems to work well:
```python
def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2) -> float:
    if isinstance(gate_logits, tuple):
        compute_device = gate_logits[0].device
        stacked_gate_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)

    routing_weights = torch.nn.functional.softmax(stacked_gate_logits, dim=-1) # [num_layers, num_tokens, num_experts]
    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) # [num_layers, num_tokens, top_k]
    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) # [num_layers, num_tokens, top_k, num_experts]
    # For a given token, determine if it was routed to a given expert. Think of this as a collection of top_k-hot vectors.
    expert_mask = torch.max(expert_mask, dim=-2).values.float() # [num_layers, num_tokens, num_experts]
    tokens_per_layer_and_expert = torch.mean(expert_mask, dim=-2) # [num_layers, num_experts]
    router_prob_per_layer_and_expert = torch.mean(routing_weights, dim=-2) # [num_layers, num_experts]
    return torch.mean(tokens_per_layer_and_expert * router_prob_per_layer_and_expert) * num_experts**2
```

I'm not 100% sure this is correct, so please check closely. Feel free to use this or reference it if it looks right. Note that the minimum loss with the above implementation is top_k, not sure if that is desired or if there should be a divide by top_k somewhere in there so the minimum loss is always 1.

@ArthurZucker  Sorry to reply so late. The new solution looks good to me. About how do we deal with top2 vs top1, I investigated some other implementations. In Google's implementation, the final loss is not divided by $top_k$ https://github.com/google/flaxformer/blob/main/flaxformer/architectures/moe/routing.py#L744.

```python
expert_mask = jax.nn.one_hot(expert_indices, num_experts, dtype=jnp.int32)
# For a given token, determine if it was routed to a given expert.
# Shape: [num_groups, tokens_per_group, num_experts]
expert_mask = jnp.max(expert_mask, axis=-2)
''' 
top1 and top2
 tensor([[0, 0, 0, 0, 0, 1, 0, 1],       
        [1, 0, 1, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 1]])
'''
tokens_per_group_and_expert = jnp.mean(
    expert_mask, dtype=jnp.float32, axis=-2)
# tensor([0.6667, 0.0000, 0.3333, 0.0000, 0.0000, 0.3333, 0.0000, 0.6667]) 
router_prob_per_group_and_expert = jnp.mean(
    router_probs, dtype=jnp.float32, axis=-2)
# tensor([0.1267, 0.1277, 0.1109, 0.1204, 0.1274, 0.1236, 0.1479, 0.1153])
return (
    jnp.mean(  # pytype: disable=bad-return-type  # jnp-type
        tokens_per_group_and_expert * router_prob_per_group_and_expert,
        dtype=jnp.float32,
    )
    * num_experts**2
  )
# tensor(2.0203) close to top2
```

In the implementation of deepspeed, only the top1 tokens participate in calculating the auxiliary loss. 
https://github.com/microsoft/DeepSpeed/blob/1787673edc7e45cd79fe10b95f92a02d3eb91505/deepspeed/moe/sharded_moe.py#L282C1-L314C60

```python
gates = F.softmax(logits, dim=1)
indices1_s = torch.argmax(gates, dim=1)
num_experts = int(gates.shape[1])
mask1 = F.one_hot(indices1_s, num_classes=num_experts)
...
# Compute l_aux
me = torch.mean(gates, dim=0)
ce = torch.mean(mask1.float(), dim=0)
l_aux = torch.mean(me * ce) * num_experts * num_experts

```

> I have been following this discussion about the load balancing loss loosely, since I have my own custom training script where I'm trying to add proper support for Mixtral. Isn't the load balancing loss still wrong, since it should be computed per layer and then averaged at the end?
> 
> The [paper](https://arxiv.org/abs/2101.03961) states "For each Switch layer, this auxiliary loss is added to the total model loss during training". To me, this seems to imply computing the loss separately per layer, then summing or averaging at the end.
> 
> I found [this implementation](https://github.com/google/flaxformer/blob/main/flaxformer/architectures/moe/routing.py) from Google that supports this. It looks like the auxiliary load balancing loss is computed per-layer as a scalar, which must then be combined at some point (I haven't read the code that closely yet).
> 
> Intuitively, it also makes sense: you would not want unbalanced expert assignments from one layer to "cancel out" those from another layer, which can happen (and seems likely to happen) if you just concatenate everything across all layers at the beginning, which is what is done now.
> 
> @ArthurZucker I know that in those equations you can compute either the sum over experts or the sum over layers first. But, you still need the fi,k per layer. If you concatenate all the gate logits at the beginning, then do a one_hot somewhere and average to get expert assignment fractions, you are losing the per-layer breakdown. So I think we need to explicitly keep everything separated by layer until the very end.
> 
> For my custom training script, I implemented the load balancing loss like this, it seems to work well:
> 
> ```python
> def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2) -> float:
>     if isinstance(gate_logits, tuple):
>         compute_device = gate_logits[0].device
>         stacked_gate_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
> 
>     routing_weights = torch.nn.functional.softmax(stacked_gate_logits, dim=-1) # [num_layers, num_tokens, num_experts]
>     _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) # [num_layers, num_tokens, top_k]
>     expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) # [num_layers, num_tokens, top_k, num_experts]
>     # For a given token, determine if it was routed to a given expert. Think of this as a collection of top_k-hot vectors.
>     expert_mask = torch.max(expert_mask, dim=-2).values.float() # [num_layers, num_tokens, num_experts]
>     tokens_per_layer_and_expert = torch.mean(expert_mask, dim=-2) # [num_layers, num_experts]
>     router_prob_per_layer_and_expert = torch.mean(routing_weights, dim=-2) # [num_layers, num_experts]
>     return torch.mean(tokens_per_layer_and_expert * router_prob_per_layer_and_expert) * num_experts**2
> ```
> 
> I'm not 100% sure this is correct, so please check closely. Feel free to use this or reference it if it looks right. Note that the minimum loss with the above implementation is top_k, not sure if that is desired or if there should be a divide by top_k somewhere in there so the minimum loss is always 1.

@tdrussell  @ArthurZucker  I also have doubts about whether to concatenate the gate logits of all layers to one tensor. 
Consider that there are only two load-unbalanced moe layers, the toekn in the first layer is routed to the first two experts, and the token in the second layer is routed to the last two experts. I did a test. If calculated individually, each layer has a loss of 2.6796(average or sum), and if made a concatenation, the final loss is 2.1001. Concatenating the outputs of each layer into a tensor would make the unbalanced load more balanced, which seems unreasonable.

However, the above test is based on the premise that the experts between layers are independent of each other. In MOE, I don't know if a token is routed to the same expert in different layers. If it is routed to the same expert, then concatenating the routing outputs of each layer into the same tensor is actually the same as computing them separately. However if not the final loss will be inaccurate, although the difference is small when the hidden state dimension is large.

1. @tdrussell, pretty sure my previous answer covers the interpretation of `For each Switch layer, this auxiliary loss is added to the total model loss during training` which means that you sum the loss computed for each layer, and I showed that this can be computed in a vectorized manner. 
2. @liangxuZhang I don't understand the concept of first and second layer of unbalanced MoE expert. But in the Mistral implementation, the routing is "independant" (compared to switch transformers where you have expert capacity for example). So mentioning my previous comment, concatenating is IMO the way to go 😉 


Maybe I am missing something. There are two statistics needed for the loss: expert assignment fractions, and average router probability per expert. If you concatenate all layers at the beginning (it is a tuple of layer logits as input to the function, right?), then these two statistics are computed without regard to which layer the logits came from. This basically says, we don't care if individual layers have unbalanced expert assignments, as long as across the whole model, the assignments are balanced on average.

It might be a matter of interpretation as to whether the loss should be that, but other implementations look like they balance the expert assignments per layer. [Here](https://gist.github.com/tdrussell/0529afd8d280fbe2c1c582d8f865e909) is a script comparing the two versions of the loss, showing the difference. I made some contrived logits that are unbalanced for each layer, but cancel out if you combine the logits from all layers.

I don't mind updating it! Will also prevent having to gather the logits on a single device, compute on the layer device then cpu transfer. It's really up to the community / the results! As I am not sure the paper / mistral shared their implementation! 🤗 

Computing this way is simpler, feel free to open a PR if you want for the per-layer computation! WOuld be nice to have results!