[BUG] DeepSpeed ZeRO-3 deadlock in engine.step() at step 0 under 2-GPU execution (RTX 3090, torch 2.2.1, DS 0.14.2)

Describe the bug
A minimal DeepSpeed ZeRO stage 3 training script consistently hangs at step 0 inside engine.step() when launched with deepspeed --num_gpus 2. Both ranks reach step_begin, then make no forward progress while the script’s heartbeat confirms they are alive but stuck in the step phase. After 60 seconds, the testcase watchdog exits and prints Test Passed ✅ (this is intentional: the testcase defines the hang as the bug oracle).
This looks like a deadlock/hang during ZeRO-3 optimizer step / communication.

To Reproduce

1. Create/activate an environment with:

- Python 3.10.19
- torch 2.2.1+cu121
- deepspeed 0.14.2

2. Run with 2 GPUs:

# (optional debug; use what you used)
export PYTHONFAULTHANDLER=1
export TORCH_SHOW_CPP_STACKTRACES=1
export TORCH_DISTRIBUTED_DEBUG=DETAIL

# network selection used in reproduction
export NCCL_SOCKET_IFNAME=eno2
export GLOO_SOCKET_IFNAME=eno2
export NCCL_IB_DISABLE=1

# ZeRO stage
export DS_ZERO_STAGE=3

# run
deepspeed --num_gpus 2 testcases/deepspeed_testcase.py

3. Observe logs:
   Both ranks print:

- R0 STEP0 route=0 -> forward
- R1 STEP0 route=0 -> forward
- ... -> backward
- R0 STEP0 route=0 -> step_begin
- R1 STEP0 route=0 -> step_begin
Then repeated:
- HEARTBEAT rank=0 step=0 phase=step_begin age=...
- HEARTBEAT rank=1 step=0 phase=step_begin age=...
After >60s:
- Test Passed ✅
- BUG_SIGNAL: HANG step=0 route=0 timeout>60s last_phase=step_begin

Expected behavior

The script should complete multiple iterations (configured for 6 iters) without hanging—i.e., engine.step() should return, and the run should proceed beyond step 0.

Actual behavior
The run consistently hangs at step 0 inside engine.step(), with both ranks stuck after printing step_begin. No step_end is printed. The heartbeat shows the program is still alive but deadlocked.

ds_report output

DeepSpeed C++/CUDA extension op report
- ninja: OKAY
- async_io: NOT installed / NOT compatible (libaio-dev missing)
- evoformer_attn: NOT installed / NOT compatible (CUTLASS_PATH not set)
- sparse_attn: NOT installed / NOT compatible (torch 2.2 detected; expected <2.0)
Other ops mostly NOT installed but compatible (JIT capable).

DeepSpeed general environment info:
torch install path ............... /home/talha/miniconda3/envs/ds_testcase/lib/python3.10/site-packages/torch
torch version .................... 2.2.1+cu121
deepspeed install path ........... /home/talha/miniconda3/envs/ds_testcase/lib/python3.10/site-packages/deepspeed
deepspeed info ................... 0.14.2, git-hash=unknown, git-branch=unknown
torch cuda version ............... 12.1
nvcc version ..................... 12.0
shared memory (/dev/shm) size .... 62.88 GB

Screenshots
Not applicable (CLI hang with logs attached).

System info (please complete the following information):

- OS: Linux (exact distro not provided; can add if needed)
- Python: 3.10.19
- DeepSpeed: 0.14.2
- PyTorch: 2.2.1+cu121
- torch CUDA: 12.1
- Driver: 550.78
- System CUDA (from nvidia-smi): 12.4
- GPUs: 4× NVIDIA GeForce RTX 3090 (24GB each) available on the node
- Reproduction uses 2 GPUs via deepspeed --num_gpus 2
- Interconnects: single node; using interface eno2; IB disabled (NCCL_IB_DISABLE=1)
- /dev/shm: 62.88 GB

Launcher context
deepspeed --num_gpus 2 testcases/deepspeed_testcase.py

Docker context
Not using Docker (not indicated in environment/logs).

Additional context

• Hang happens immediately at step 0 during engine.step().
• The testcase includes barriers and a heartbeat thread to verify both ranks enter the step together and remain stuck.
• ds_report indicates some ops are not installed and some are not compatible (async_io, evoformer_attn, sparse_attn), but this testcase does not explicitly require those ops.
• Warning present: sparse_attn expects torch < 2.0 but torch 2.2.1 is installed (not sure if relevant to the hang, but included for completeness).

Start with testcases/deepspeed_testcase.py and reproduce the ZeRO-3 run using deepspeed --num_gpus 2, checking the rank logs around engine.step() and step_begin. Trace the relevant ZeRO-3 optimizer and communication path to identify the deadlock. Done means the six-iteration testcase returns from step 0, prints step_end, and completes without the hang signal.