Load balancing loss in Switch Transformer
Author: ToheartZhang
Pull request: no

### System Info

- `transformers` version: 4.36.2
- Platform: Linux-5.15.0-92-generic-x86_64-with-glibc2.35
- Python version: 3.10.11
- Huggingface_hub version: 0.20.2
- Safetensors version: 0.3.1
- Accelerate version: 0.22.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- 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?: True
- Using distributed or parallel set-up in script?: True

### Who can help?

@ArthurZucker @younesbelkada

### 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 think there is a bug in load_balancing loss. Here `router_probs` is [batch_size, seqeunce_length, num_experts], which does not consider different layers. However, experts are different across layers.
https://github.com/huggingface/transformers/blob/092f1fdaa4224fdd88c616dc9678e6fcb37bfffd/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L91-L107

I think the error source is from the unpacking.
https://github.com/huggingface/transformers/blob/092f1fdaa4224fdd88c616dc9678e6fcb37bfffd/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L1699-L1707
The layer dimension is eliminated by `torch.cat`.

### Expected behavior

I think `_unpack_router_logits` should use `stack` instead, and `load_balancing_loss_func` should consider layer dimension since experts from different layers are different. Otherwise, the loss may be meaningless.

Comments:
There was a good debate here: https://github.com/huggingface/transformers/issues/28255#issuecomment-1874241942 
Might shed some light as to why we don't consider layer-wise 

Thank you for your explanation! However, I still have some doubts. Suppose there are 4 layers and 4 experts in each layer, if the 1st layer all routes to expert 1, the 2nd layer all routes to expert2, etc. It seems quite balanced for the concatenating solution, but not for the per-layer solution.

For the formula
$$\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 concatenating solution is actually
$$\text{total loss} =  \alpha \cdot N \cdot \sum_{i=1}^{N}   \left(   \sum_{k=1}^{K}  f_{k} \cdot P_{i,k} \right). $$
Therefore, it is not equivalent to the original loss:
$$\text{total loss} =  \sum_{k=1}^{K} \left( \alpha \cdot N \cdot \sum_{i=1}^{N} f_{i,k} \cdot P_{i,k} \right)$$

> the concatenating solution is actually 

That does not look right to me, the way we compute it takes into account $$ f_{i,k} $$.

Now again as I stated before in my comment, if we go for solution 1 vs solution 2 I don't mind as long as we have proof / certitude that this brings better results. The paper did not share a code and the interpretation we have right now is mathematically correct, the other interpretation is also correct, but the first 1 (that we are currently doing) is closer to what is mentioned in the paper. 

TLDR; if it **is** better, let's add a section in the README or let's advertise the other loss. We usually leave it to the user to define the loss function and this one can be overwritten 😉 

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.