[BUG] The forward hook function of Zero stege3 might misjudge the current step of the training process whether it is in the forward computation or in the backward recomputation

Describe the bug
pre_sub_module_forward_function https://github.com/microsoft/DeepSpeed/blob/9dfb06de36bb29293b1e94dc1e48d6f2adf54d2c/deepspeed/runtime/zero/parameter_offload.py#L493 The function utilizes prev_grad_state to store the result of torch.is_grad_enabled, and this prev_grad_state is then passed as an argument to fetch_sub_module to determine whether it is executed in the forward pass. In other words, the execution of fetch_sub_module in forward mode depends on the outcome of torch.is_grad_enabled.

However, when activation-checkpoint is enabled, torch.is_grad_enabled returns False, even though the computation is taking place during the forward pass. This leads to fetch_sub_module mistakenly assuming that the current program is not in the process of executing forward computation. More critically, it triggers the HPZ optimization of ZeRO++ during the forward pass, whereas this behavior is intended to occur during the backward pass.

To Reproduce
Steps to reproduce the behavior:

1. Go to deepspeed-megatron set --zero-stage 3 --deepspeed-activation-checkpointing and ds_config.json set zero_hpz_partition_size 2
2. Go to [deepspeed/deepspeed/runtime/zero/parameter_offload.py in line 505. ]https://github.com/microsoft/DeepSpeed/blob/9dfb06de36bb29293b1e94dc1e48d6f2adf54d2c/deepspeed/runtime/zero/parameter_offload.py#L493
3. log the prev_grad_state
4. kick off training
5. See error: prev_grad_state will be set as false in forward step which leads all_gather_coalesced(https://github.com/microsoft/DeepSpeed/blob/9dfb06de36bb29293b1e94dc1e48d6f2adf54d2c/deepspeed/runtime/zero/partition_parameters.py#L1123C25-L1123C25) use zero_param_process_group. likewise, prev_grad_state will be set as true in backward step which leads all_gather_coalesced use ds_process_group.

Expected behavior
all_gather_coalesced use ds_process_group in forward step and zero_param_process_group for backward.

ds_report output
{
"train_batch_size": 16,
"gradient_accumulation_steps": 1,
"steps_per_print": 1,
"zero_optimization": {
"stage": 3,
"allgather_partitions": true,
"allgather_bucket_size": 5.000000e+08,
"overlap_comm": true,
"reduce_scatter": true,
"stage3_max_live_parameters": 1.000000e+06,
"stage3_max_reuse_distance": 1.000000e+06,
"stage3_prefetch_bucket_size": 1.000000e+05,
"reduce_bucket_size": 5.000000e+08,
"contiguous_gradients": true,
"zero_hpz_partition_size": 2,
"zero_quantized_weights": false,
"zero_quantized_gradients": false
},
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.00015,
"max_grad_norm": 1.0,
"betas": [0.9, 0.95]
}
},
"gradient_clipping": 1.0,
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"wall_clock_breakdown": true
}