load_balancing_loss in mixtral model
Author: 1773226512
Pull request: no

### System Info

torch '1.13.0+cu117'

### Who can help?

@ArthurZucker and @younesbelkada

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

The balancing loss function always return a constant.
Here is the official code:
```
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)
```
Here is my code:
```
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)
```
It always return 4.



### Expected behavior

please anwser this question

Comments:
same question.

Hi everyone, https://github.com/huggingface/transformers/pull/28115 should fix this issue

The aux loss is constant too(now is 8 but not 4, and no grad when BP) after using PR #28115 
@younesbelkada @ArthurZucker 

I also meet this. And there is no `grad_fn` in the loss tensor.

> The aux loss is constant too(now is 8 but not 4, and no grad when BP) after using PR #28115 @younesbelkada @ArthurZucker



Not sure I understand. Here is a small repro:
```python
from transformers import MixtralForCausalLM, MixtralConfig
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
config = MixtralConfig(num_hidden_layers=2, hidden_size=1024,  )
model = MixtralForCausalLM(config)

inputs = tokenizer("This sentence should be properly", return_tensors="pt")
inputs["labels"] = inputs["input_ids"]
outs = model(**inputs)
print(outs.loss)
outs.loss.backward()

inputs = tokenizer("A prioris cela devrait changer", return_tensors="pt")
inputs["labels"] = inputs["input_ids"]
outs = model(**inputs)
print(outs.loss)
```
```
tensor(10.5236, grad_fn=<NllLossBackward0>)
tensor(10.8024, grad_fn=<NllLossBackward0>)
```
could both of you open a new issue and elaborate? 

If I understand correctly, the loss consists of two parts, the autoregressive categorical loss on the one hand, and balancing the loss of each expert on the other.
I print each of the above two losses before adding up the final one, however, only the first one has a `grad_fn` of `NllLossBackward0`, the second one is just a tensor without `grad_fn` . That's why the `grad_fn` of the final loss is `NllLossBackward0` instead of `AddBackward`.
And the balancing expert's loss doesn't change during the training process. Maybe the changes you see are just changes due to the autoregressive categorical loss.

I open a new issue https://github.com/huggingface/transformers/issues/28205.