DeepSpeed Stage 3 in lightning leads to Nan and Inf values in the model parameters.

[LittleFlyingSheep]

Bug description

I try to use lightning with DeepSpeed stage 3 to train a model under the precision "16-mixed". However, I find that the model parameters includes Nan and Inf values at the first step. When I change it to DDP, this issue does not exist.

I initialize my trainer as:

trainer = Trainer(
        max_epochs=max_epochs,
        logger=logger,
        callbacks=[checkpoint_callback, lr_monitor],
        sync_batchnorm=sync_batchnorm,
        check_val_every_n_epoch=None,
        val_check_interval=every_n_train_steps * accumulate_grad_batches,
        devices="auto",
        accelerator="gpu",
        precision="16-mixed",
        strategy=deepspeed_stage_3，
        accumulate_grad_batches=accumulate_grad_batches, 
    )

What version are you seeing the problem on?

v2.5

How to reproduce the bug

trainer = Trainer(
        max_epochs=max_epochs,
        logger=logger,
        callbacks=[checkpoint_callback, lr_monitor],
        sync_batchnorm=sync_batchnorm,
        check_val_every_n_epoch=None,
        val_check_interval=every_n_train_steps * accumulate_grad_batches,
        devices="auto",
        accelerator="gpu",
        precision="16-mixed",
        strategy=deepspeed_stage_3，
        accumulate_grad_batches=accumulate_grad_batches, 
    )

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

[lantiga]

hi @LittleFlyingSheep can you provide a minimal reproduction? it will help the investigation

[LittleFlyingSheep]

收到