Lightning-AI / Lightning-AI/pytorch-lightning

Device mesh initialization fails when running FSDP on CPU

Open Beginner friendly
#21,814 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug ver: 2.5.x ver: 2.6.x
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.