Stall during code execution and continuous increase in CPU usage

Describe the bug
I want to apply zero-3 and offload the relevant parameters to NVME. I found during debugging in PyCharm that after running the “assert (swap_handle.async_pwrite(buffer, path) == 0)” code, the CPU memory began to increase at a rate of approximately 0.1-0.2GiB/s. And when it reaches “assert self.pending_writes == self.aio_write_handle.wait()”, the code comes to a standstill and CPU usage continues to grow

To Reproduce
Steps to reproduce the behavior:

1. Install related dependencies (including libaio) and set CFLAGS and LDFLAGS when installing deepspeed
2. git clone https://github.com/microsoft/Megatron-DeepSpeed.git
3. cd /home/Megatron-DeepSpeed/examples_deepspeed/offload_pp
4. RUN /home/yaxuanw/anaconda3/envs/deepspeed/bin/python3 /home/yaxuanw/microsoftds/Megatron-DeepSpeed/pretrain_gpt.py --override-opt_param-scheduler --adam-beta1 0.9 --adam-beta2 0.95 --tensor-model-parallel-size 1 --moe-expert-parallel-size 1 --num-experts 1 --moe-loss-coeff 0.01 --moe-train-capacity-factor 1.0 --moe-eval-capacity-factor 1.0 --moe-min-capacity 4 --init-method-std 0.014 --lr-decay-tokens 260000000000 --lr-warmup-tokens 375000000 --micro-batch-size 2 --exit-duration-in-mins 30000000 --rampup-batch-size 32 32 1953125 --global-batch-size 256 --num-layers 24 --hidden-size 1024 --num-attention-heads 16 --seq-length 2048 --max-position-embeddings 2048 --train-tokens 300000000000 --train-samples 439453125 --lr 3.0e-4 --min-lr 3.0e-5 --lr-decay-style cosine --split 98,2,0 --log-interval 1 --eval-interval 100 --eval-iters 10 --save-interval 1000 --weight-decay 0.1 --clip-grad 1.0 --hysteresis 2 --num-workers 0 --fp16 --load /home/yaxuanw/microsoftds/Megatron-DeepSpeed/examples_deepspeed/offload_pp/output/checkpoint/gpt-0.35B-lr-3.0e-4-minlr-3.0e-5-bs-256-gpus-16-mp-1-pp-1 --save /home/yaxuanw/microsoftds/Megatron-DeepSpeed/examples_deepspeed/offload_pp/output/checkpoint/gpt-0.35B-lr-3.0e-4-minlr-3.0e-5-bs-256-gpus-16-mp-1-pp-1 --tensorboard-queue-size 1 --log-timers-to-tensorboard --timing-log-level 1 --no-pipeline-parallel --cpu-optimizer --log-batch-size-to-tensorboard --log-validation-ppl-to-tensorboard --tensorboard-dir /home/yaxuanw/microsoftds/Megatron-DeepSpeed/examples_deepspeed/offload_pp/output/tensorboard/gpt-0.35B-lr-3.0e-4-minlr-3.0e-5-bs-256-gpus-16-mp-1-pp-1_SINIAN-01_2023.12.19-23.34.15 --checkpoint-activations --vocab-file /home/yaxuanw/microsoftds/Megatron-DeepSpeed/dataset/gpt2-vocab.json --merge-file /home/yaxuanw/microsoftds/Megatron-DeepSpeed/dataset/gpt2-merges.txt --data-path /home/yaxuanw/microsoftds/Megatron-DeepSpeed/dataset/meg-gpt2-oscar-en-10k_text_document --data-impl mmap --deepspeed --deepspeed_config /home/yaxuanw/microsoftds/Megatron-DeepSpeed/examples_deepspeed/offload_pp/ds_config_gpt_gpt-0.35B-lr-3.0e-4-minlr-3.0e-5-bs-256-gpus-16-mp-1-pp-1.json --pipeline-model-parallel-size 1 --deepspeed-activation-checkpointing
5. My configuration file is as follows：
   {
   "train_batch_size" : 256,
   "train_micro_batch_size_per_gpu": 1,
   "fp16": {
   "enabled": true
   },
   "optimizer": {
   "type": "Adam"
   },
   "zero_optimization": {
   "stage": 3,
   "offload_optimizer": {
   "device": "nvme",
   "nvme_path": "/home/yaxuanw/microsoftds/Megatron-DeepSpeed/nvme",
   "pin_memory": true,
   "buffer_count": 4,
   "fast_init": false,
   "ratio": 0.1
   },
   "offload_param": {
   "device": "nvme",
   "nvme_path": "/home/yaxuanw/microsoftds/Megatron-DeepSpeed/nvme",
   "pin_memory": true,
   "buffer_count": 10,
   "buffer_size": 1e8,
   "max_in_cpu": 1e8
   },
   "overlap_comm": true,
   "contiguous_gradients": true,
   "stage3_max_live_parameters": 1e9,
   "stage3_max_reuse_distance": 1e9,
   "reduce_bucket_size": "auto",
   "stage3_prefetch_bucket_size": "auto",
   "stage3_param_persistence_threshold": "auto"
   },
   "aio": {
   "block_size": 1048576,
   "queue_depth": 32,
   "thread_count": 8,
   "single_submit": true,
   "overlap_events": true
   },
   "activation_checkpointing": {},
   "flops_profiler": {
   "enabled": true,
   "profile_step": 1,
   "module_depth": -1,
   "top_modules": 1,
   "detailed": true,
   "output_file": null
   }
   }

Expected behavior
The training process runs normally and the relevant parameters are offloaded to NVME according to the configuration

ds_report output
[2023-12-26 00:52:10,139] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)

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

async_io ............... [YES] ...... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [YES] ...... [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/yaxuanw/anaconda3/envs/deepspeed/lib/python3.10/site-packages/torch']
torch version .................... 2.1.2+cu121
deepspeed install path ........... ['/home/yaxuanw/anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.12.6, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.2
deepspeed wheel compiled w. ...... torch 2.1, cuda 12.1
shared memory (/dev/shm) size .... 251.78 GB

System info (please complete the following information):

• OS: Ubuntu 22.04.3
• one machine with x4 V100s
• Python version: 3.10.0

Docker context
no

Start with pretrain_gpt.py in the examples_deepspeed/offload_pp reproduction and the ds_config_gpt_gpt-0.35B-lr-3.0e-4-minlr-3.0e-5-bs-256-gpus-16-mp-1-pp-1.json configuration. Reproduce the stall around swap_handle.async_pwrite and aio_write_handle.wait using the supplied command, then trace the async I/O entry point. Done means training completes normally without unbounded CPU-memory growth while parameters are offloaded to NVME.