[BUG] Failed to checkpoint with deepspeed 0.12.4

Describe the bug
I encountered an issue when using DeepSpeed 0.12.4 with the OpenChat trainer, where checkpointing failed and raised an NCCL error. However, the checkpoints work fine when using DeepSpeed 0.12.2-0.12.3. The trainer utilizes the save_pretrained from the Transformers library and deepspeed.checkpoint.utils.clone_tensors_for_torch_save, as seen here.

Interestingly, the checkpoint only fails if the training is conducted with more than a certain number of steps (likely >100 steps) using ZeRO. Training with fewer steps does not result in a failure.

Here is the error message:

[10:38:00<41:53:38, 10.39s/it][E ProcessGroupNCCL.cpp:474] [Rank 7] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1, OpType=
REDUCE, NumelIn=2, NumelOut=2, Timeout(ms)=1800000) ran for 1800475 milliseconds before timing out.                                                                                                                                                                  
[E ProcessGroupNCCL.cpp:488] 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.                                                                        
[E ProcessGroupNCCL.cpp:494] To avoid data inconsistency, we are taking the entire process down.                                                                                                                                                                     
[E ProcessGroupNCCL.cpp:915] [Rank 7] NCCL watchdog thread terminated with exception: [Rank 7] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1, OpType=REDUCE, NumelIn=2, NumelOut=2, Timeout(ms)=1800000) ran for 1800475 milliseconds before timing
 out.  

Expected behavior
Checkpoint should work correctly with deepspeed 0.12.4

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/share/tianshikeji/syh18_tsinghua/miniconda3/envs/openchat_deepspeed0124/lib/python3.11/site-packages/torch']
torch version .................... 2.1.1+cu121
deepspeed install path ........... ['/home/share/tianshikeji/syh18_tsinghua/miniconda3/envs/openchat_deepspeed0124/lib/python3.11/site-packages/deepspeed']
deepspeed info ................... 0.12.4, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 11.8
deepspeed wheel compiled w. ...... torch 2.1, cuda 12.1
shared memory (/dev/shm) size .... 1007.76 GB

System info (please complete the following information):

• OS: Ubuntu 20.04.6 LTS
• H100x8 NVLink (single machine)
• Also observed this issue on A40x8

Launcher context
deepspeed launcher on H100x8

deepspeed --module ochat.training_deepspeed.train \
	--model_type openchat_v3.2_mistral \
	--model_path imone/Mistral_7B_with_EOT_token \
	--data_prefix DATA_TOKENIZED \
	--save_path SAVE_PATH \
	--save_every 1 \
	--epochs 5 \
	--batch_max_len 77824 \
	--deepspeed \
	--deepspeed_config ochat/training_deepspeed/deepspeed_config.json