Repository: huggingface/transformers (Python) — Bug Fix

Issue:
 Incorrect implementation of auxiliary loss 
### 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.


Discussion:

Resolving patch:
diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py
--- a/src/transformers/models/mixtral/modeling_mixtral.py
+++ b/src/transformers/models/mixtral/modeling_mixtral.py
@@ -103,11 +103,7 @@ def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tenso
 
     _, 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)
-    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)
@@ -115,7 +111,7 @@ def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tenso
     # 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.unsqueeze(-1))
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
     return overall_loss * num_experts
 
 


Test patch:
diff --git a/tests/models/mixtral/test_modeling_mixtral.py b/tests/models/mixtral/test_modeling_mixtral.py
--- a/tests/models/mixtral/test_modeling_mixtral.py
+++ b/tests/models/mixtral/test_modeling_mixtral.py
@@ -474,7 +474,7 @@ def test_load_balancing_loss(self):
         model.eval()
         result = model(input_ids, attention_mask=attention_mask)
         self.assertEqual(result.router_logits[0].shape, (91, config.num_local_experts))
-        torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(8, dtype=torch.float32))
+        torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2)
 
 
 @require_torch