[BUG] deepspeed zero2 training hangon and timeout after a fixed step

[leeruibin]

Describe the bug
I use deepspeed zero2 to train a transformer-based dit model. However, the script always gets stuck at a fixed step after one hour of training. When I disable deepspeed and use pure pytorch DDP to train the code, the problem disappears. Moreover, even if I change the mixed-precision from fp16 to bf16 or adjust the learning rate, the same problem occurs and the script gets stuck at the same training step.

I have also tried changing the model initialization and resuming the training, but the script still gets stuck after the same training step. For example, if the script gets stuck after training 14K steps, I save the checkpoint at the 10K step and resume training with the 10K-step checkpoint. Then, after training another 14K steps again, the script gets stuck once more.

ds_report output
[rank0]:[E217 07:24:07.938443115 ProcessGroupNCCL.cpp:616] [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=54053, OpType=ALLREDUCE, NumelIn=497464352, NumelOut=497464352, Timeout(ms)=600000) ran for 600087 milliseconds before timing out.
[rank0]:[E217 07:24:07.938551513 ProcessGroupNCCL.cpp:1785] [PG ID 1 PG GUID 1 Rank 0] Exception (either an error or timeout) detected by watchdog at work: 54053, last enqueued NCCL work: 54054, last completed NCCL work: 54052.
[rank0]:[E217 07:24:08.990085728 ProcessGroupNCCL.cpp:1834] [PG ID 1 PG GUID 1 Rank 0] Timeout at NCCL work: 54053, last enqueued NCCL work: 54054, last completed NCCL work: 54052.
[rank0]:[E217 07:24:08.990112551 ProcessGroupNCCL.cpp:630] [Rank 0] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank0]:[E217 07:24:08.990119150 ProcessGroupNCCL.cpp:636] [Rank 0] To avoid data inconsistency, we are taking the entire process down.
[rank0]:[E217 07:24:08.991954167 ProcessGroupNCCL.cpp:1595] [PG ID 1 PG GUID 1 Rank 0] Process group watchdog thread terminated with exception: [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=54053, OpType=ALLREDUCE, NumelIn=497464352, NumelOut=497464352, Timeout(ms)=600000) ran for 600087 milliseconds before timing out.
Exception raised from checkTimeout at /opt/tiger/compile_path/src/code.byted.org/pytorch/pytorch/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:618 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x91 (0x7f6aa3f5fdd1 in /usr/local/lib/python3.11/dist-packages/torch/lib/libc10.so)
frame #1: + 0x1022831 (0x7f6aa4fee831 in /usr/local/lib/python3.11/dist-packages/torch/lib/libtorch_cuda.so)
frame #2: c10d::ProcessGroupNCCL::WorkNCCL::checkTimeout(std::optional<std::chrono::duration<long, std::ratio<1l, 1000l> > >) + 0x22a (0x7f6aa500d27a in /usr/local/lib/python3.11/dist-packages/torch/lib/libtorch_cuda.so)
frame #3: c10d::ProcessGroupNCCL::watchdogHandler() + 0x22e (0x7f6aa500d92e in /usr/local/lib/python3.11/dist-packages/torch/lib/libtorch_cuda.so)
frame #4: c10d::ProcessGroupNCCL::ncclCommWatchdog() + 0x143 (0x7f6aa500f5f3 in /usr/local/lib/python3.11/dist-packages/torch/lib/libtorch_cuda.so)
frame #5: + 0xd44a3 (0x7f6a96cba4a3 in /usr/lib/x86_64-linux-gnu/libstdc++.so.6)
frame #6: + 0x89144 (0x7f6ae9427144 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #7: + 0x1097dc (0x7f6ae94a77dc in /usr/lib/x86_64-linux-gnu/libc.so.6)

terminate called after throwing an instance of 'c10::DistBackendError'

System info (please complete the following information):

• Python version: Python 3.11.2
• Pytorch version: 2.5.1+cu124
• Deepspeed version: 0.16.3
• Any other relevant info about your setup

Launcher context
python -m torch.distributed.launch --nnodes=1 --nproc_per_node=8 --master_port=12345 train_deepspeed.py --xxx

Additional context
Add any other context about the problem here.

[leeruibin]

I inserted some log code into my code to make the script print certain information, aiming to check at which step the code gets stuck. Then, I found that the script hangs during the backward pass on rank 0 after several steps. Here is an example of the log information: the backward pass occurs in step 33417, but not in step 33418.

================================================

pass index selection on device 0 step 33417
pass index selection on device 2 step 33417
pass index selection on device 5 step 33417
pass index selection on device 6 step 33417
pass index selection on device 3 step 33417
pass index selection on device 1 step 33417
pass index selection on device 7 step 33417
pass index selection on device 4 step 33417

pass forward on device 2 step 33417
pass forward on device 0 step 33417
pass forward on device 6 step 33417
pass forward on device 1 step 33417
pass forward on device 3 step 33417
pass forward on device 5 step 33417
pass forward on device 7 step 33417
pass forward on device 4 step 33417

pass backward on device 7 step 33417
pass backward on device 1 step 33417
pass backward on device 6 step 33417
pass backward on device 0 step 33417
pass backward on device 3 step 33417
pass backward on device 2 step 33417
pass backward on device 4 step 33417
pass backward on device 5 step 33417

pass optimization step on device 2 step 33417
pass optimization step on device 7 step 33417
pass optimization step on device 6 step 33417
pass optimization step on device 5 step 33417
pass optimization step on device 4 step 33417
pass optimization step on device 3 step 33417
pass optimization step on device 0 step 33417
pass optimization step on device 1 step 33417
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:7', dtype=torch.int32), running_loss: 5.311360597610474, device: 7
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:0'), running_loss: 5.180153250694275, device: 0
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:1', dtype=torch.int32), running_loss: 5.240574240684509, device: 1
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:2', dtype=torch.int32), running_loss: 5.239009439945221, device: 2
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:5', dtype=torch.int32), running_loss: 5.259356379508972, device: 5
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:3', dtype=torch.int32), running_loss: 5.346103191375732, device: 3
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:6', dtype=torch.int32), running_loss: 5.2149001359939575, device: 6
train_steps: 33417, shared_interval_idx: tensor([1], device='cuda:4', dtype=torch.int32), running_loss: 5.125726819038391, device: 4

pass index selection on device 4 step 33418
pass index selection on device 2 step 33418
pass index selection on device 3 step 33418
pass index selection on device 5 step 33418
pass index selection on device 1 step 33418
pass index selection on device 7 step 33418
pass index selection on device 6 step 33418

pass forward on device 2 step 33418
pass forward on device 7 step 33418
pass forward on device 4 step 33418
pass forward on device 1 step 33418
pass forward on device 5 step 33418
pass forward on device 3 step 33418
pass forward on device 6 step 33418
pass forward on device 0 step 33418

pass backward on device 2 step 33418
pass backward on device 6 step 33418
pass backward on device 5 step 33418
pass backward on device 7 step 33418
pass backward on device 4 step 33418
pass backward on device 1 step 33418
pass backward on device 3 step 33418

pass optimization step on device 5 step 33418
pass optimization step on device 4 step 33418
pass optimization step on device 2 step 33418
pass optimization step on device 6 step 33418
pass optimization step on device 3 step 33418
pass optimization step on device 7 step 33418
pass optimization step on device 1 step 33418

train_steps: 33418, shared_interval_idx: tensor([2], device='cuda:5', dtype=torch.int32), running_loss: 6.078900039196014, device: 5
train_steps: 33418, shared_interval_idx: tensor([2], device='cuda:4', dtype=torch.int32), running_loss: 5.950639605522156, device: 4
train_steps: 33418, shared_interval_idx: tensor([2], device='cuda:2', dtype=torch.int32), running_loss: 6.0558441281318665, device: 2
train_steps: 33418, shared_interval_idx: tensor([2], device='cuda:3', dtype=torch.int32), running_loss: 6.203619062900543, device: 3
train_steps: 33418, shared_interval_idx: tensor([2], device='cuda:6', dtype=torch.int32), running_loss: 6.059363543987274, device: 6
train_steps: 33418, shared_interval_idx: tensor([2], device='cuda:7', dtype=torch.int32), running_loss: 6.090326964855194, device: 7
train_steps: 33418, shared_interval_idx: tensor([2], device='cuda:1', dtype=torch.int32), running_loss: 6.035567224025726, device: 1

[rank0]:[E218 08:45:14.863205760 ProcessGroupNCCL.cpp:616] [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=54043, OpType=ALLREDUCE, NumelIn=497464352, NumelOut=497464352, Timeout(ms)=600000) ran for 600008 milliseconds before timing out.
[rank0]:[E218 08:45:14.865658012 ProcessGroupNCCL.cpp:1785] [PG ID 1 PG GUID 1 Rank 0] Exception (either an error or timeout) detected by watchdog at work: 54043, last enqueued NCCL work: 54044, last completed NCCL work: 54042.
[rank0]:[E218 08:45:14.865679631 ProcessGroupNCCL.cpp:1834] [PG ID 1 PG GUID 1 Rank 0] Timeout at NCCL work: 54043, last enqueued NCCL work: 54044, last completed NCCL work: 54042.
[rank0]:[E218 08:45:14.865685345 ProcessGroupNCCL.cpp:630] [Rank 0] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank0]:[E218 08:45:14.865690479 ProcessGroupNCCL.cpp:636] [Rank 0] To avoid data inconsistency, we are taking the entire process down.
[rank0]:[E218 08:45:14.867338996 ProcessGroupNCCL.cpp:1595] [PG ID 1 PG GUID 1 Rank 0] Process group watchdog thread terminated with exception: [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=54043, OpType=ALLREDUCE, NumelIn=497464352, NumelOut=497464352, Timeout(ms)=600000) ran for 600008 milliseconds before timing out.

[leeruibin]

I output more logs and found that the error was caused by the mismatch in different [components?]. However, I have no idea why this happens. I am training a moe-model with a fixed routing. For example, I route the computation to the same expert by first setting an index to 0 (which means routing to the first expert). Then, before the forward pass, I broadcast the index to all processes to ensure that all processes activate the same expert. Here is the example code:

if dist.get_rank() == 0:
    probs = torch.tensor([0.2395, 0.521, 0.2395])
    shared_interval_idx = torch.multinomial(probs, 1).to("cuda")
else:
    shared_interval_idx = torch.zeros(1, dtype=torch.int32, device="cuda")

dist.broadcast(shared_interval_idx, src=0)
model.forward(x,y,shared_interval_idx)

When I output the shared_interval_idx, the number is the same for all the processes. However, the script crashes after several steps, accompanied by the following log. Does anyone have any advice?

Rank 0: Before broadcast
Rank 2: Before broadcast
Rank 4: Before broadcast
Rank 6: Before broadcast
Rank 3: Before broadcast
Rank 5: Before broadcast
Rank 1: Before broadcast
Rank 7: Before broadcast

Rank 0: Tensor shape torch.Size([1]), dtype torch.int32
Rank 2: Tensor shape torch.Size([1]), dtype torch.int32

Rank 4: Tensor shape torch.Size([1]), dtype torch.int32
Rank 6: Tensor shape torch.Size([1]), dtype torch.int32

Rank 3: Tensor shape torch.Size([1]), dtype torch.int32
Rank 5: Tensor shape torch.Size([1]), dtype torch.int32
Rank 7: Tensor shape torch.Size([1]), dtype torch.int32
Rank 1: Tensor shape torch.Size([1]), dtype torch.int32

Rank 1: After broadcast
Rank 4: After broadcast
Rank 0: After broadcast
Rank 3: After broadcast
Rank 6: After broadcast
Rank 5: After broadcast
Rank 2: After broadcast
Rank 7: After broadcast
pass index selection on device 1 step 3942 with shared_interval_idx as tensor([1], device='cuda:1', dtype=torch.int32)
pass index selection on device 2 step 3942 with shared_interval_idx as tensor([1], device='cuda:2', dtype=torch.int32)
pass index selection on device 4 step 3942 with shared_interval_idx as tensor([1], device='cuda:4', dtype=torch.int32)
pass index selection on device 0 step 3942 with shared_interval_idx as tensor([1], device='cuda:0', dtype=torch.int32)
pass index selection on device 3 step 3942 with shared_interval_idx as tensor([1], device='cuda:3', dtype=torch.int32)
pass index selection on device 6 step 3942 with shared_interval_idx as tensor([1], device='cuda:6', dtype=torch.int32)
pass index selection on device 5 step 3942 with shared_interval_idx as tensor([1], device='cuda:5', dtype=torch.int32)
pass index selection on device 7 step 3942 with shared_interval_idx as tensor([1], device='cuda:7', dtype=torch.int32)
pass forward on device 3 step 3942
pass forward on device 2 step 3942
pass forward on device 4 step 3942
pass forward on device 5 step 3942
pass forward on device 1 step 3942
pass forward on device 6 step 3942
pass forward on device 7 step 3942
pass forward on device 0 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 2 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 344502304vs 497464352
pass backward on device 2 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 6 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 344502304vs 497464352
pass backward on device 6 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 3 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 344502304vs 497464352
pass backward on device 3 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 4 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 344502304vs 497464352
pass backward on device 4 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 1 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 344502304vs 497464352
pass backward on device 1 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 1 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 497464352vs 344502304
pass backward on device 0 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 5 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 344502304vs 497464352
pass backward on device 5 step 3942
Error during backward on rank0: Detected mismatch between collectives on ranks. Rank 7 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[344502304], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=16150, OpType=ALLREDUCE, TensorShape=[497464352], TensorDtypes=Half, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: Tensor Tensor shapes: 344502304vs 497464352
pass backward on device 7 step 3942
[rank1]: Traceback (most recent call last):
[rank1]: File "/mnt/bn/genai-nebula/liruibin/codes/DiT/train_deepspeed_debug.py", line 488, in
[rank1]: main(args)
[rank1]: ^^^^^^^^^^
[rank1]: File "/mnt/bn/genai-nebula/liruibin/codes/DiT/train_deepspeed_debug.py", line 320, in main
[rank1]: model_engine.step()
[rank1]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/engine.py", line 2249, in step
[rank1]: self._take_model_step(lr_kwargs)
[rank1]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/engine.py", line 2152, in _take_model_step
[rank1]: self.optimizer.step()
[rank1]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1831, in step
[rank1]: self.check_overflow()
[rank1]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2069, in check_overflow
[rank1]: self._check_overflow(partition_gradients)
[rank1]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1983, in _check_overflow
[rank1]: self.overflow = self.has_overflow(partition_gradients)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2003, in has_overflow
[rank1]: overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial()
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1996, in has_overflow_partitioned_grads_serial
[rank1]: for j, grad in enumerate(self.averaged_gradients[i]):
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: TypeError: 'NoneType' object is not iterable
[rank5]: Traceback (most recent call last):
[rank5]: File "/mnt/bn/genai-nebula/liruibin/codes/DiT/train_deepspeed_debug.py", line 488, in
[rank5]: main(args)
[rank5]: ^^^^^^^^^^
[rank5]: File "/mnt/bn/genai-nebula/liruibin/codes/DiT/train_deepspeed_debug.py", line 320, in main
[rank5]: model_engine.step()
[rank5]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/engine.py", line 2249, in step
[rank5]: self._take_model_step(lr_kwargs)
[rank5]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/engine.py", line 2152, in _take_model_step
[rank5]: self.optimizer.step()
[rank5]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1831, in step
[rank5]: self.check_overflow()
[rank5]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2069, in check_overflow
[rank5]: self._check_overflow(partition_gradients)
[rank5]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1983, in _check_overflow
[rank5]: self.overflow = self.has_overflow(partition_gradients)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2003, in has_overflow
[rank5]: overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial()
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/usr/local/lib/python3.11/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1996, in has_overflow_partitioned_grads_serial
[rank5]: for j, grad in enumerate(self.averaged_gradients[i]):
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: TypeError: 'NoneType' object is not iterable

[cedricgsh]

@leeruibin I met the same problem, have you figure it out ?

[cedricgsh]

I increase reduce bucket size. Model can train with longer time but still meet these question. I use H20 GPU

[leeruibin]

@cedricgsh Not yet. I have experimented with various parameters and batch sizes, but the issue persists. I am actively seeking a solution. For now, I am using PyTorch DDP (Distributed Data Parallel) to train the model instead of integrating DeepSpeed, and the training process is functioning normally.

[tjruwase]

python -m torch.distributed.launch --nnodes=1 --nproc_per_node=8 --master_port=12345 train_deepspeed.py --xxx

@leeruibin, @cedricgsh are you able to share full repro steps?

[cedricgsh]

@leeruibin But in this way, you cannot speed up your training process

[leeruibin]

@tjruwase Here is a mini demo of my training code. I use deepspeed zero 2 to train the model.

import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
import torch.distributed as dist
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import argparse
import deepspeed
import torch.nn as nn
import torch.nn.functional as F
import torch.profiler
from torch.utils.data import Dataset

#################################################################################
#                             Training Helper Functions                         #
#################################################################################


def requires_grad(model, flag=True):
    """
    Set requires_grad flag for all parameters in a model.
    """
    for p in model.parameters():
        p.requires_grad = flag


class DemoDataset(Dataset):
    def __init__(self, ):
        self.samples = [torch.randn([3,224,224]) for i in range(1000)]

    def __getitem__(self, index):
        img = self.samples[index]
        target = 100
        return img, target
    def __len__(self):
        return len(self.samples)

from timm.layers.patch_embed import PatchEmbed
from timm.models.vision_transformer import PatchEmbed, Attention, Mlp

def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class Block(nn.Module):
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, use_flash_attn=False, **block_kwargs):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)

        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)


    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

class FinalLayer(nn.Module):
    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)

    def forward(self, x):
        x = self.norm_final(x)
        x = self.linear(x)
        return x

class DemoModel(nn.Module):
    def __init__(self,):
        super().__init__()
        self.hidden_size = 512
        self.first_layer = PatchEmbed(img_size=224,patch_size=2,in_chans=3,embed_dim=self.hidden_size,bias=True)
        self.blocks_1 = nn.ModuleList([
            Block(self.hidden_size, 4, mlp_ratio=4.0) for _ in range(6)
        ])
        self.blocks_2 = nn.ModuleList([
            Block(self.hidden_size, 4, mlp_ratio=4.0) for _ in range(6)
        ])
        self.blocks_3 = nn.ModuleList([
            Block(self.hidden_size, 4, mlp_ratio=4.0) for _ in range(6)
        ])
        self.final_layer = FinalLayer(self.hidden_size,patch_size=2,out_channels=3)
    
    def unpatchify(self, x):
        """
        x: (N, T, patch_size**2 * C)
        imgs: (N, H, W, C)
        """
        c = 3
        p = 2
        h = w = int(x.shape[1] ** 0.5)
        assert h * w == x.shape[1]

        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
        return imgs
    
    def forward(self, x, index):
        x = self.first_layer(x)
        if index == 0:
            for block in self.blocks_1:
                x = block(x)
        elif index == 1:
            for block in self.blocks_2:
                x = block(x)
        else:
            for block in self.blocks_3:
                x = block(x)
        
        x = self.final_layer(x)
        x = self.unpatchify(x)   
        return x


#################################################################################
#                                  Training Loop                                #
#################################################################################

def main(args):
    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
    
    deepspeed.init_distributed()

    rank = args.local_rank
    device = rank % torch.cuda.device_count()
    seed = args.global_seed * dist.get_world_size() + rank
    torch.manual_seed(seed)
    torch.cuda.set_device(device)
    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
    

    model = DemoModel()
    dataset = DemoDataset()
    sampler = DistributedSampler(
        dataset,
        num_replicas=dist.get_world_size(),
        rank=rank,
        shuffle=True,
        seed=args.global_seed
    )
    loader = DataLoader(
        dataset,
        batch_size=args.train_batch_size,
        shuffle=False,
        sampler=sampler,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=True
    )

    model_engine, opt, _, __ = deepspeed.initialize(
        args=args, model=model, model_parameters=model.parameters())

    
    train_steps = 0

    for epoch in range(1000):
        sampler.set_epoch(epoch)
        data_iter_step = 0
        for x, y in loader: 
            model_engine.train() 
            x = x.to(device,dtype=torch.float16) 
            y = y.to(device) 

            if dist.get_rank() == 0:
                probs = torch.tensor([0.2395, 0.521, 0.2395])
                shared_interval_idx = torch.multinomial(probs, 1).to("cuda",dtype=torch.int32)
            else:
                shared_interval_idx = torch.zeros(1, dtype=torch.int32, device="cuda")

            dist.broadcast(shared_interval_idx, src=0)
            output = model(x,shared_interval_idx)
            loss = F.mse_loss(output,x)

            model_engine.backward(loss) 
            model_engine.step()
            if rank==0:
                out_loss = loss.item()
                print(f"training loss: {out_loss}")
            train_steps += 1
            data_iter_step += 1



if __name__ == "__main__": 
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_batch_size", type=int, default=32)
    parser.add_argument("--global-seed", type=int, default=2023) 
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument('--accum_iter', default=1, type=int,)
    parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher') 
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args() 

    print(args)
    main(args) 

And here is the deepspeed config

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1,
        "verbose": true
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 0.0001,
            "betas": [
                0.9,
                0.999
            ],
            "eps": 1e-08,
            "weight_decay": 0
        }
    },
    "bf16": {
        "enabled": false
    },
    "train_micro_batch_size_per_gpu": 32,
    "gradient_accumulation_steps": 1,
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1000000000.0
    },
    "logging": {
        "level": "DEBUG",
        "print_rank": "all"
    },
    "steps_per_print": 10000
}

[JiaruiFeng]

I am facing the same issue, is there any update?

[goyp695]

I am facing the same issue, is there any update?

[sfc-gh-truwase]

@googe859 I am unable to repro on master branch with above example. Can you please try with master branch or share a different repro if possible?

[tjruwase]

Closing due to lack of response. Please re-open if needed.