Lightning-AI / Lightning-AI/pytorch-lightning
Mixed precision autocast is disabled when running FSDP on CPU
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
The `FSDPPrecision` plugin hardcodes the `"cuda"` device type in `forward_context` when using mixed precision (e.g., `bf16-mixed`).
When FSDP is run on CPU with mixed precision, this hardcoded `"cuda"` value causes PyTorch to print a warning and disable autocast. Because autocast is disabled, CPU operations are not cast to `bfloat16` during the forward pass, which can lead to type mismatches (e.g., `Float` inputs meeting `BFloat16` weights) or incorrect benchmarking results.
### What version are you seeing the problem on?
v2.5
### Reproduced in studio
_No response_
### How to reproduce the bug
```python
import torch
from lightning.fabric.plugins.precision.fsdp import FSDPPrecision
# Initialize precision plugin for bf16-mixed
precision = FSDPPrecision(precision="bf16-mixed")
# On a CPU-only environment, enter forward context
with precision.forward_context():
# This will trigger: UserWarning: CUDA is not available... Disabling autocast.
print("Autocast CPU enabled:", torch.is_autocast_enabled("cpu"))
# Expected: True (if we want bf16 mixed on CPU)
# Actual: False
```
### Error messages and logs
```
UserWarning: CUDA is not available or torch_xla is imported. Disabling autocast.
```
This can also lead to:
```
RuntimeError: expected scalar type Float but found BFloat16
```
or similar type mismatch errors when inputs are not auto-casted to match the FSDP sharded parameter type.
### Environment
Current environment
```
- PyTorch Lightning Version: master/latest
- PyTorch Version: 2.0+
- OS: Linux/macOS/Windows
- Accelerator: CPU
```
### More info
`FSDPPrecision` should be updated to accept or dynamically set the target `device_type` (e.g., propagating `self.root_device.type` from the FSDP strategy during setup) and pass it to `torch.autocast`.
cc @ethanwharris
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 in the FSDPPrecision implementation, especially forward_context, and trace how the plugin gets its device from FSDP strategy setup. Reproduce the issue with bf16-mixed on CPU and inspect how torch.autocast receives its device type. Done means CPU autocast is enabled without the CUDA warning while existing CUDA behavior remains correct.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100