Training hangs at the first gradient syncing of an MoE model while using deepspeed

[negar-foroutan]

System Info

• transformers version: 4.39.3
• Platform: Linux-5.15.0-94-generic-x86_64-with-glibc2.35
• Python version: 3.12.2
• Huggingface_hub version: 0.22.2
• Safetensors version: 0.4.2
• Accelerate version: 0.28.0
• PyTorch version (GPU?): 2.2.2+cu121 (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?:
• Using distributed or parallel set-up in script?:

Who can help?

@pacman100

Information

• The official example scripts
• My own modified scripts

Tasks

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

Reproduction

I'm training a customized MoE language model using 8 GPUs on one node, and it works fine using Accelerate (without DeepSpeed). However, when I enable DeepSpeed (both ZerO 1 and 2), the training hangs at the first gradient syncing and will end up with an NCCL timeout after a few minutes.
The MoE is built based on the XGLM architecture and the task as language modeling.
DeepSpeed works fine when I train a dense model (not MoE).
I also set deepspeed_moe_layer_cls_names to my MoE block, but it doesn't seem to work.
I use the accelerate launch command to run my experiments.
I guess the problem with MoE is that not all GPUs use the same parameters (experts) in one forward pass, and that's why GPUs are waiting for each other to receive all gradients.

Expected behavior

The gradient syncing and training continue properly, and we are able to train an MoE (sparse network) using DeepSpeed.