Lightning-AI / Lightning-AI/pytorch-lightning
Device mesh initialization fails 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
When running FSDP on CPU (e.g., for benchmarking or testing), the strategy attempts to initialize the `DeviceMesh` using a hardcoded `"cuda"` device type. This causes a failure on CPU-only environments where CUDA is not available, or when FSDP is intentionally run on CPU.
### 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 import Fabric
from lightning.fabric.strategies import FSDPStrategy
# Configure FSDP with device_mesh on CPU
strategy = FSDPStrategy(parallel_devices=[torch.device("cpu")], device_mesh=(1,))
fabric = Fabric(strategy=strategy, devices=1)
fabric.launch()
```
### Error messages and logs
```
File "src/lightning/fabric/strategies/fsdp.py", line 271, in setup_environment
self._fsdp_kwargs["device_mesh"] = init_device_mesh("cuda", self._fsdp_kwargs["device_mesh"])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
...
RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx
```
### Environment
Current environment
```
- PyTorch Lightning Version: master/latest
- PyTorch Version: 2.2+ (where DeviceMesh is supported in FSDPStrategy)
- OS: Linux/macOS/Windows
- Accelerator: CPU
```
### More info
The hardcoded `"cuda"` string in `setup_environment` should be replaced with `self.root_device.type` to dynamically support the correct device type.
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 src/lightning/fabric/strategies/fsdp.py at setup_environment and reproduce the CPU configuration shown in the issue. Verify that DeviceMesh initialization uses the selected device type and that the CPU run completes without requiring CUDA.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100