[BUG] convergence issues with `zero_hpz_partition_size`

[stas00]

Describe the bug

I wanted to try ZeRO++ and found that using zero_hpz_partition_size has convergence issues.

The current 0.12.5 version doesn't converge at all

Since then I tried https://github.com/microsoft/DeepSpeed/tree/HeyangQin/mixz_hpz_fix and it's learning, but much slower than w/o it.

Basically, the expectation is that the hybrid and non-hybrid settings should lead to an identical loss curve.

As a base I'm using HF Transformers llama-2-7b on 4 A100 nodes trained from scratch.

with HeyangQin/mixz_hpz_fix it's learning but not well:

2 problems:

1. the loss starts much higher by 1-2 points
2. it's not catching up with the baseline

ds config:

   "zero_optimization": {
       "stage": 3,
       "overlap_comm": true,
       "reduce_scatter": true,
       "reduce_bucket_size": 90000000,
       "contiguous_gradients": true,
       "stage3_gather_16bit_weights_on_model_save": false,
       "stage3_prefetch_bucket_size": 5e7,
       "stage3_param_persistence_threshold": 1e5,
       "stage3_max_live_parameters": 3e9,
       "stage3_max_reuse_distance": 3e9,
       "sub_group_size": 1e9,
       
       "zero_hpz_partition_size": 8,
       
       "offload_optimizer": {
           "device": "none"
       },
       "offload_param": {
           "device": "none"
       }
   },

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 ............... ['/env/lib/conda/tr041-dawn-llama/lib/python3.9/site-packages/torch']
torch version .................... 2.1.1+cu121
deepspeed install path ........... ['/data/env/lib/repos/retro-llama/tr043-dawn-llama-3/DeepSpeed/deepspeed']
deepspeed info ................... 0.12.6+48ddf31d, 48ddf31d, HeyangQin/mixz_hpz_fix
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 .... 669.32 GB

Screenshots If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

• OS: Ubuntu 22.04
• GPU count and types 4 machines with x8 A100s each
• Python version 3.9

@HeyangQin