Is MoE implemented correctly?
MoE FFNs are supposed to be sharded between GPUs. It seems like the current MoE implementation puts all FFNs on each GPU. See e.g. `https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/megatron/mlp.py#L321`. Am I missing something here?