Lightning-AI / Lightning-AI/pytorch-lightning

Allow passing custom reader/writer in _distributed_checkpoint_save and _distributed_checkpoint_load.

Open
#20,205 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

### Description & Motivation

[distributed_checkpoint_save](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/strategies/fsdp.py#L861) for sharded dict initializes [torch.distributed.checkpoint.FileSystemWriter](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/strategies/fsdp.py#L876)

Similarly [distributed_checkpoint_load](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/strategies/fsdp.py#L880) and [load_checkpoint](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/strategies/fsdp.py#L613) initializes torch.distributed.checkpoint.FileSystemReader.

If PyTorch-lightining could provide an integration point to plugin custom readers/writers it would be provide user with freedom to read/write from whatever storage they want.

### Pitch

I want to use pytorch lightning to read/write from GCS bucket using FSDP & sharded checkpoint.
This can be done easily in single node checkpointing since it follows checkpoint IO interface and easily allows to read/write checkpoints from anywhere we want.

However this is not the case for sharded checkpoints in multi node scenario. Having the ability to plugin custom read/writer will help us read/write from storage we want.

### Alternatives

_No response_

### Additional context

_No response_

cc @lantiga @borda

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

Read the distributed_checkpoint_save and distributed_checkpoint_load entry points in src/lightning/fabric/strategies/fsdp.py, along with load_checkpoint in src/lightning/pytorch/strategies/fsdp.py, and inspect where FileSystemWriter and FileSystemReader are initialized. Done means users can provide custom reader and writer implementations for sharded, multi-node FSDP checkpoint saving and loading.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.