[BUG] Logits are always ZERO(0) at first pass when using ZERO++

Describe the bug
When training, the logits are always zero during the first forward pass of the model. From the subsequent forward passes, the logits are output normally. The DeepSpeed configuration is as follow:

deepspeed:
  train_batch_size: auto
  train_micro_batch_size_per_gpu: auto
  gradient_accumulation_steps: auto
  gradient_clipping: auto
  zero_allow_untested_optimizer: true
  bf16:
    enabled: auto
  zero_optimization:
    stage: 3
    contiguous_gradients: true
    sub_group_size: 1.0e+9
    reduce_bucket_size: auto
    stage3_prefetch_bucket_size: auto
    stage3_param_persistence_threshold: auto
    stage3_max_live_parameters: 1.0e+9
    stage3_max_reuse_distance: 1.0e+9
    stage3_gather_16bit_weights_on_model_save: true
    zero_quantized_weights: false
    zero_hpz_partition_size: 2
    zero_quantized_gradients: false

I tested using huggingface trainer with 2 GPUs. The problem occurs on both of 2 A100s and 2 H100s.

To Reproduce

Full code:

from datasets import load_dataset
from transformers import HfArgumentParser
from trl import SFTConfig, SFTTrainer

class Trainer(SFTTrainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        outputs = super().compute_loss(model, inputs, True)
        # At first step, logits are all zero
        print(outputs[0], outputs[1])
        return outputs if return_outputs else outputs[0]

args = HfArgumentParser(SFTConfig).parse_yaml_file("config.yaml")[0]
train_dataset = load_dataset("imdb", split="train")
eval_dataset = load_dataset("imdb", split="test")

trainer = Trainer(
    "facebook/opt-350m",
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    args=args,
)

trainer.train()

dataset_text_field: text
max_seq_length: 512
output_dir: /tmp

num_train_epochs: 1
per_device_train_batch_size: 1
per_device_eval_batch_size: 1
gradient_accumulation_steps: 1
learning_rate: 2.0e-5
lr_scheduler_type: cosine
optim: adamw_torch_fused

do_train: true
do_eval: true
seed: 42
bf16: true

gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false

deepspeed:
  train_batch_size: auto
  train_micro_batch_size_per_gpu: auto
  gradient_accumulation_steps: auto
  gradient_clipping: auto
  zero_allow_untested_optimizer: true
  fp16:
    enabled: auto
    loss_scale: 0
    loss_scale_window: 1000
    initial_scale_power: 16
    hysteresis: 2
    min_loss_scale: 1
  bf16:
    enabled: auto
  zero_optimization:
    stage: 3
    contiguous_gradients: true
    sub_group_size: 1.0e+9
    reduce_bucket_size: auto
    stage3_prefetch_bucket_size: auto
    stage3_param_persistence_threshold: auto
    stage3_max_live_parameters: 1.0e+9
    stage3_max_reuse_distance: 1.0e+9
    stage3_gather_16bit_weights_on_model_save: true
    zero_quantized_weights: false
    zero_hpz_partition_size: 2
    zero_quantized_gradients: false

logging_strategy: steps
logging_steps: 10
eval_strategy: steps
eval_steps: 100
save_strategy: steps
save_steps: 100

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
--------------------------------------------------
async_io ............... [NO] ....... [OKAY]
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]
 [WARNING]  FP Quantizer is using an untested triton version (3.0.0), only 2.3.0 and 2.3.1 are known to be compatible with these kernels
fp_quantizer ........... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
transformer_inference .. [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.4
 [WARNING]  using untested triton version (3.0.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]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['***']
torch version .................... 2.4.0+cu121
deepspeed install path ........... ['***']
deepspeed info ................... 0.14.5, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.4
deepspeed wheel compiled w. ...... torch 2.4, cuda 12.1
shared memory (/dev/shm) size .... 755.80 GB