hangs during training using deepspeed
Author: mxjmtxrm
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.54.0.dev0
- Platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.33.0
- Safetensors version: 0.5.3
- Accelerate version: 1.8.1
- Accelerate config:    not found
- DeepSpeed version: 0.17.1
- PyTorch version (accelerator?): 2.8.0a0+5228986c39.nv25.05 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@SunMarc @zach-huggingface

### 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

train moe model using run_clm.py with deepspeed zero3.
```
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --config_file "deepspeed.yaml"  run_clm.py \
    --do_train \
    --do_eval \
    --model_name_or_path Qwen3-30B-A3B \
    --trust_remote_code True \
    --torch_dtype auto \
    --train_file train.json \
    --validation_file val.json \
    --block_size 8192 \
    --output_dir ./output \
    --weight_decay 0.01 \
    --learning_rate 5e-6 \
    --lr_scheduler_type "cosine" \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --gradient_accumulation_steps 1 \
    --gradient_checkpointing True \
    --gradient_checkpointing_kwargs '{"use_reentrant": true}' \
    --max_grad_norm 0.5 \
    --warmup_ratio 0.03 \
    --bf16 \
    --num_train_epochs 1 \
    --aux_lambda 0 \
    --eval_strategy "steps" \
    --eval_steps 100 \
    --save_strategy "steps" \
    --save_steps 100 \
    --report_to "tensorboard" \
    --save_total_limit 2 \
    --logging_steps 10 \
    --log_level "info" \
    --include_num_input_tokens_seen True \
    --overwrite_output_dir \
    --save_only_model False 
```

deepspeed.yaml:
```
compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

The training process hangs here:
```
  0%|                                                                                                                            | 1/4966 [00:57<78:36:48, 57.00s/it]
```
And the GPU utilization remains at 100% and does not change.

### Expected behavior

Where is the problem? How to solve it?

Comments:
I encountered the same issue when training with Qwen2-VL
I check the NCCL with `export NCCL_DEBUG=TRACE`.
 I found that during AllReduce, the data size (count) is inconsistent across nodes, which should never happen:
```
hgx13: pro-03:3440984:3440984 [4] NCCL INFO AllReduce: opCount 2dc sendbuff 0x7f6ccf9aca00 recvbuff 0x7f6ccf9aca00 count 487
250944 datatype 9 op 0 root 0 comm 0x5560fa060ea0 [nranks=32] stream 0x556095ee4300                                         
hgx13: pro-03:3440984:3440984 NCCL CALL ncclAllReduce(7f6ccf9aca00,7f6ccf9aca00,487250944,9,0,0,0x5560fa060ea0,0x556095ee43$
0)                                                                                                                          
hgx16: pro-12:1773031:1773031 [5] NCCL INFO AllReduce: opCount 2dc sendbuff 0x7fb4b8000000 recvbuff 0x7fb4b8000000 count 605
64992 datatype 9 op 0 root 0 comm 0x55cbe651d4b0 [nranks=32] stream 0x55cbe7d600c0                                          
hgx16: pro-12:1773031:1773031 NCCL CALL ncclAllReduce(7fb4b8000000,7fb4b8000000,60564992,9,0,0,0x55cbe651d4b0,0x55cbe7d600c0
)       
```
but I have no idea why these things happens, anybody can help? OVO

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.