Lightning-AI / Lightning-AI/pytorch-lightning
DeepSpeed Stage 3 in lightning leads to Nan and Inf values in the model parameters.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### 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
```python
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_
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing the reported Trainer configuration with DeepSpeed Stage 3 and 16-mixed precision, then compare its first training step with DDP. The report does not name files or tests and omits environment details and logs; done means identifying why parameters become NaN or Inf under Stage 3 and confirming they remain finite after the fix.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100