When configured as "zero offload," the computation of backpropagation and the copying from GPU to CPU cannot be parallelized.[BUG]

Describe the bug
Both of my questions occur in the grad_partition function in stage3.py.
Q1: When configured as "zero offload," the computation of backpropagation and the copying of gradients from GPU to CPU cannot be parallelized.
Q2: Why are three cross-side copies needed instead of directly using "grad_partition"?

To Reproduce
out deepspeed config is
json = {
"train_batch_size": 64,
"train_micro_batch_size_per_gpu": 1,
"steps_per_print": 10,
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu"
},
"stage3_param_persistence_threshold": 1.000000e+04,
"stage3_max_live_parameters": 3.000000e+07,
"stage3_prefetch_bucket_size": 3.000000e+07,
"memory_efficient_linear": false
},
"fp16": {
"enabled": true,
"loss_scale_window": 100
},
"gradient_clipping": 1.0,
"prescale_gradients": false,
"wall_clock_breakdown": false,
"activation_checkpointing": {
"partition_activations": false,
"cpu_checkpointing": false,
"contiguous_memory_optimization": false,
"number_checkpoints": null,
"synchronize_checkpoint_boundary": false,
"profile": true
},
"hybrid_engine": {
"enabled": false,
"max_out_tokens": 512,
"inference_tp_size": 1,
"release_inference_cache": false,
"pin_parameters": true,
"tp_gather_partition_size": 8
}
}

Expected behavior
According to the description in "ZeRO-Offload: Democratizing Billion-Scale Model Training," computation and transmission should be parallelizable.

Screenshots

We use Nsight for testing,the result is as follow:

System info (please complete the following information):

• OS: Ubuntu 18.04
• GPU count and type: one machines with x4 V100
• Python version: 3.10.12
• Any other relevant info about your setup
  pytorch==1.13.1
  deepspeed==0.10.1

Launcher context
we use DeepSpeed Examples rep and use cmd as follows:
CUDA_VISIBLE_DEVICES=0,1,2,3 taskset -c 0-13,28-41 python3 train.py --step 1 --deployment-type single_node --actor-model 1.3b

Start in stage3.py at the grad_partition function and reproduce the reported configuration with the provided DeepSpeed Examples command. Use Nsight to inspect whether backpropagation overlaps GPU-to-CPU gradient copies, and trace the three cross-side copies. Done means the parallelization behavior and copy requirements are explained, or a verified fix is accompanied by relevant regression coverage.