[BUG] Deepspeed zero++ hpz hangs forever

Describe the bug

After https://github.com/microsoft/DeepSpeed/pull/4906, zero++ hpz training hangs forever (with 2 nodes with 8 A100s each). I reverted the commit and the training did not hang anymore, so it is very likely that the commit caused the issue.

To Reproduce

Hardware Settings:

A100 nodes * 2, each with 8 A100 GPUs, requested 240 CPUs and 1920GB RAM each.

Python Dependencies:

1. transformers==4.35.2
2. torch==2.0.0+cu118
3. flash-attn==2.3.6
4. local built deepspeed on master

Code

1. Copy from Gist
2. Replace the following offline path in run.sh with the online one

      --model_path /shared/public/models/models-do-not-use-for-production/llama-2-7b-hf \
      --data_path /shared/public/data/mmlu \

Error

Dump from one of the processes in the master node:

$ py-spy dump  --pid  7114
Process 7114: /export/apps/python/3.10/bin/python3 -u src/train.py --model_path /shared/public/models/models-do-not-use-for-production/llama-2-7b-hf --data_path /shared/public/data/mmlu --logging_dir /shared/user/mmlu-output/f6eb8c3d53ff44973a4f-2024-01-29T21:54:38+0000-log --bf16 True --max_steps 50 --per_device_train_batch_size 16 --per_device_eval_batch_size 16 --gradient_checkpointing --max_seq_length 2048 --save_strategy no --learning_rate 2e-5 --logging_steps 1 --lora False --report_to none --output_dir /shared/user/mmlu-output/f6eb8c3d53ff44973a4f-2024-01-29T21:54:38+0000 --deepspeed src/baseline_hpz.json
Python v3.10.2 (/export/apps/python/3.10.2/bin/python3.10)

Thread 7114 (active): "MainThread"
    forward (deepspeed/runtime/zero/linear.py:57)
    decorate_fwd (torch/cuda/amp/autocast_mode.py:98)
    apply (torch/autograd/function.py:506)
    zero3_linear_wrap (deepspeed/runtime/zero/linear.py:109)
    forward (torch/nn/modules/linear.py:114)
    _call_impl (torch/nn/modules/module.py:1538)
    forward (transformers/models/llama/modeling_llama.py:366)
    _call_impl (torch/nn/modules/module.py:1538)
    forward (transformers/models/llama/modeling_llama.py:672)
    _call_impl (torch/nn/modules/module.py:1538)
    forward (torch/utils/checkpoint.py:107)
    apply (torch/autograd/function.py:506)
    checkpoint (torch/utils/checkpoint.py:249)
    forward (transformers/models/llama/modeling_llama.py:912)
    _call_impl (torch/nn/modules/module.py:1538)
    forward (transformers/models/llama/modeling_llama.py:1034)
    _call_impl (torch/nn/modules/module.py:1538)
    forward (deepspeed/runtime/engine.py:1842)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
    _call_impl (torch/nn/modules/module.py:1501)
    compute_loss (transformers/trainer.py:2748)
    training_step (transformers/trainer.py:2725)
    _inner_training_loop (transformers/trainer.py:1859)
    train (transformers/trainer.py:1555)
    train (trl/trainer/sft_trainer.py:323)
    main (train.py:102)
    <module> (train.py:108)
Thread 8433 (idle): "Thread-1"
    wait (threading.py:324)
    wait (threading.py:600)
    run (tqdm/_monitor.py:60)
    _bootstrap_inner (threading.py:1009)
    _bootstrap (threading.py:966)

Dump from one of the processes in the worker node:

$ py-spy dump --pid 5952
Process 5952: /export/apps/python/3.10/bin/python3 -u src/train.py --model_path /shared/public/models/models-do-not-use-for-production/llama-2-7b-hf --data_path /shared/public/data/mmlu --logging_dir /shared/user/mmlu-output/f6eb8c3d53ff44973a4f-2024-01-29T21:54:39+0000-log --bf16 True --max_steps 50 --per_device_train_batch_size 16 --per_device_eval_batch_size 16 --gradient_checkpointing --max_seq_length 2048 --save_strategy no --learning_rate 2e-5 --logging_steps 1 --lora False --report_to none --output_dir /shared/user/mmlu-output/f6eb8c3d53ff44973a4f-2024-01-29T21:54:39+0000 --deepspeed src/baseline_hpz.json
Python v3.10.2 (/export/apps/python/3.10.2/bin/python3.10)

Thread 5952 (idle): "MainThread"
    backward (torch/autograd/__init__.py:200)
    backward (torch/_tensor.py:487)
    backward (deepspeed/runtime/fp16/loss_scaler.py:63)
    backward (deepspeed/runtime/zero/stage3.py:2152)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
    backward (deepspeed/runtime/engine.py:1964)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
    backward (accelerate/utils/deepspeed.py:167)
    backward (accelerate/accelerator.py:1958)
    training_step (transformers/trainer.py:2734)
    _inner_training_loop (transformers/trainer.py:1859)
    train (transformers/trainer.py:1555)
    train (trl/trainer/sft_trainer.py:323)
    main (train.py:102)
    <module> (train.py:108)
Thread 7269 (idle): "Thread-1"
    wait (threading.py:324)
    wait (threading.py:600)
    run (tqdm/_monitor.py:60)
    _bootstrap_inner (threading.py:1009)
    _bootstrap (threading.py:966)
Thread 7716 (idle)
Thread 7718 (idle)
Thread 7717 (idle)
Thread 7719 (idle)
Thread 7720 (idle)
Thread 7721 (idle)
    reduce_scatter_tensor (torch/distributed/distributed_c10d.py:2987)
    wrapper (torch/distributed/distributed_c10d.py:1436)
    reduce_scatter_tensor (deepspeed/comm/torch.py:255)
    reduce_scatter_tensor (deepspeed/comm/comm.py:289)
    log_wrapper (deepspeed/comm/comm.py:117)
    reduce_scatter_fn (deepspeed/comm/comm.py:257)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
    _torch_reduce_scatter_fn (deepspeed/runtime/comm/coalesced_collectives.py:23)
    reduce_scatter_coalesced (deepspeed/runtime/comm/coalesced_collectives.py:120)
    decorate_context (torch/utils/_contextlib.py:115)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
    __avg_scatter_grads (deepspeed/runtime/zero/stage3.py:1282)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
    __reduce_and_partition_ipg_grads (deepspeed/runtime/zero/stage3.py:1213)
    decorate_context (torch/utils/_contextlib.py:115)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
    reduce_independent_p_g_buckets_and_remove_grads (deepspeed/runtime/zero/stage3.py:1163)
    reduce_ready_partitions_and_remove_grads (deepspeed/runtime/zero/stage3.py:1422)
    reduce_partition_and_remove_grads (deepspeed/runtime/zero/stage3.py:1129)
    wrapped_fn (deepspeed/utils/nvtx.py:15)
Thread 7722 (idle)
Thread 7723 (idle)