Lightning-AI / Lightning-AI/pytorch-lightning

Enable FSDP strategy on CPU (accelerator="cpu")

Open
#21,812 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature strategy: fsdp
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Motivation

Currently, attempting to use FSDPStrategy with CPU in PyTorch Lightning raises an artificial ValueError during trainer initialization:

ValueError: The strategy `fsdp` requires a GPU accelerator, but received `accelerator='cpu'`.

Allowing FSDP to run on CPU is highly valuable for several key workflows:

1. Storage I/O Macrobenchmarking: Enabling CPU-based FSDP allows us to run macrobenchmarks on checkpoint save/restore storage I/O without consuming expensive GPU hours. Isolating storage network and disk throughput from GPU compute provides clear signals for future performance improvements in distributed checkpointing.
2. Local Development & Debugging: Developers can build, debug, and verify distributed FSDP code locally on laptops without requiring GPU access.
3. CI & Unit Testing: Enables fast, inexpensive multi-process FSDP end-to-end testing in CPU-only CI environments (e.g., GitHub Actions).

### Related Issues

- This feature proposal is blocked by / related to #21804, which tracks the crash/segmentation fault when saving full state dict checkpoints under CPU FSDP. Once it is resolved, CPU support can be fully enabled and tested for both full and sharded checkpoint formats.
- Precursor work in #21774 resolved PyTorch >= 2.5 contract issues by passing explicit CPU device_id s to FSDP, unblocking CPU compatibility at the PyTorch wrapper level.

### Pitch

I propose officially enabling accelerator="cpu" support for FSDPStrategy :

1. Remove Accelerator Restrictions: Update _AcceleratorConnector to accept "cpu" and CPUAccelerator instances when using FSDPStrategy .
2. Update Error Messaging: Revise validation error messages to state that FSDP accepts CUDA, GPU, or CPU accelerators.
3. Add Comprehensive E2E Testing: Introduce multi-rank CPU training tests (e.g., test_fsdp_cpu_trainable ) that verify:
- Genuine FSDP layer wrapping ( FullyShardedDataParallel ) on CPU across multiple ranks.
- Training loss convergence and optimizer stepping under the Gloo distributed backend.
- Checkpoint saving and round-trip parameter verification for both sharded and full ( state_dict_type ) checkpoint formats once #21804 is resolved.

### Alternatives

_No response_

### Additional context

_No response_

cc @lantiga

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 by inspecting _AcceleratorConnector and FSDPStrategy validation to trace the CPU accelerator restriction and current error messaging. Add the proposed test_fsdp_cpu_trainable coverage for multi-rank CPU training with the Gloo backend, including genuine FSDP wrapping and optimizer progress. Checkpoint round-trip coverage for full state dicts depends on the resolution of issue #21804.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.