[BUG] MoE using too much RAM

Describe the bug
I'm using a MoE instead of an MLP in a transformer block, I'm using expert parallelism such that I have one expert per gpu (with 2 gpus) and top-2 routing (so that I use both of them). Theoretically the model should use the exact same amount of ram as before given that the experts are distributed by gpu, but in practice I keep getting CUDA out of memory.

Expected behavior
The model should use exactly the same amount of ram as before instead of doubling.

Launcher context
I'm using the OpenMPI Launcher

Additional context
The model is a transformer with repeated (a lot actually) of blocks, and in each of them I wanted to replace the MLP with a MoE of 2 experts, one per gpu, to increase the model capacity without increasing the memory requirements