Mixtral's implementation of auxiliary loss seems incorrect
Author: Tobin-rgb
Pull request: no
Labels: name: Text

### System Info

- `transformers` version: 4.41.1
- Platform: Linux-4.18.0-193.28.1.el8_2.x86_64-x86_64-with-glibc2.17
- Python version: 3.10.12
- Huggingface_hub version: 0.23.1
- Safetensors version: 0.4.3
- Accelerate version: 0.30.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)

### 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, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I have recently been trying to train a MOE model, but I found that the following loss function did not balance the experts of each layer adequately.
I read the code and found the operation `torch.cat` of all `layer_gate` confusing me.  It seems to treat experts with the same index in each layer as one expert.

https://github.com/huggingface/transformers/blob/9ba9369a2557e53a01378199a9839ec6e82d8bc7/src/transformers/models/mixtral/modeling_mixtral.py#L78-L151

I noticed that you mentioned the equivalence of the following two formulas here https://github.com/huggingface/transformers/issues/28255 https://github.com/huggingface/transformers/issues/29503

$$ 
\text{total loss} = \sum_{k=1}^{K}\left(\alpha \cdot N \cdot \sum_{i=1}^{N}f_{i,k}\cdot P_{i,k}\right)
$$

$$
\text{total loss} = \alpha \cdot N \cdot \sum_{i=1}^{N}\left( \sum_{k=1}^{K}f_{i,k}\cdot P_{i,k}\right)
$$

However, the following code seems to implement this: 

$$
\text{total loss} = \alpha \cdot N \cdot \sum_{i=1}^{N} \left( \sum_{k=1}^{K}f_{i,k}\cdot  \sum_{k=1}^{K}P_{i,k} \right)
$$

since `sum` and `mean` are just a constant multiple difference for the loss.
```python
# Compute the percentage of tokens routed to each experts
tokens_per_expert = torch.mean(expert_mask.float(), dim=0)

# 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(0)) 
```


### Expected behavior

`load_balancing_loss_func`  needs to be improved.

Comments:
The currently implementation will push a balance to the global tokens routed to each expert in all the layers. However, a correct implementation should push a balance such that the tokens assigned to each expert is balanced inside each layer, instead of globally

If you use the trick from DeepSeekV2. https://huggingface.co/deepseek-ai/deepseek-moe-16b-base/blob/main/modeling_deepseek.py#L341
it turns out that there is no way to align with the current mixtral implementation due to this bug. So, the correct way to implement is not to concat logits from all the layers and then compute aux loss. Instead, we should compute aux loss for each layer and then adds them.

The correct implementation should be something like
```python

def load_balancing_loss_func(
    gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2, attention_mask: Optional[torch.Tensor] = None
) -> 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 model.config.num_hidden_layers tensors of
            shape [batch_size X sequence_length, num_experts].
        attention_mask (`torch.Tensor`, None):
            The attention_mask used in forward function
            shape [batch_size X sequence_length] if not None.
        num_experts (`int`, *optional*):
            Number of experts

    Returns:
        The auxiliary loss.
    """
    if gate_logits is None or not isinstance(gate_logits, tuple):
        return 0

    compute_device = gate_logits[0].device

    layer_aux_loss = []
    for layer_gate in gate_logits:
        routing_weights = torch.nn.functional.softmax(layer_gate, dim=-1)

        _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)

        expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)

        if attention_mask is None:
            # Compute the percentage of tokens routed to each experts
            tokens_per_expert = torch.mean(expert_mask.float(), dim=0)

            # Compute the average probability of routing to these experts
            router_prob_per_expert = torch.mean(routing_weights, dim=0)
        else:
            batch_size, sequence_length = attention_mask.shape
            num_hidden_layers = layer_gate.shape[0] // (batch_size * sequence_length)

            # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
            expert_attention_mask = (
                attention_mask[None, :, :, None, None]
                .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
                .reshape(-1, top_k, num_experts)
                .to(compute_device)
            )

            # Compute the percentage of tokens routed to each experts
            tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
                expert_attention_mask, dim=0
            )

            # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
            router_per_expert_attention_mask = (
                attention_mask[None, :, :, None]
                .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
                .reshape(-1, num_experts)
                .to(compute_device)
            )

            # Compute the average probability of routing to these experts
            router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
                router_per_expert_attention_mask, dim=0
            )

        overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
        layer_aux_loss.append(overall_loss)
    
    overall_loss = sum(layer_aux_loss)
    return overall_loss * num_experts
```

Some operators can be computed in concat form, including `softmax` and `topk` to improve the speed. However, `tokens_per_expert` and `router_prob_per_expert` must be computed layer by layer.

1. Both are correct
2. Depending on the number of experts, one can be favored vs the other
3. A lot of MoE models have been trained and mentioned that they used this exact loss (Ex: Qwen-MoE). 
4. If deepseek v2 uses this, super, we can add with the new model. 

I really don't mind, but there was never once and "official" version of the loss. And only thing was mentionned in the paper + people successfully training the model. 

At this point feel free to open a PR with before vs after losses showing that this is a lot better on the range of MoE we have and it will probably be alright to merge it 🤗 


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.