Lightning-AI / Lightning-AI/pytorch-lightning

Optimize remote checkpoint loading with parallel multiprocess downloads and zero-copy mmap

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

Nobody has claimed this yet.

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

Description

### Description & Motivation

Currently, loading monolithic, multi-gigabyte checkpoints from remote object stores (such as Google Cloud Storage gs://) via _load() suffers from two significant bottlenecks:

1. Sequential Main-Thread Streaming I/O: Upstream Lightning streams remote checkpoints sequentially over fs.open() on the main execution thread. Network bandwidth is underutilized, leading to high restore latency for large checkpoints.
2. High Peak Heap RAM Utilization & OOM Risks: Streaming directly into torch.load(f, mmap=False) forces Python to allocate heap buffers for incoming bytes while simultaneously constructing unpickled model tensors in memory, creating transient RAM spikes and OOM hazards on standard nodes.

### Pitch

We propose enhancing lightning.fabric.utilities.cloud_io._load() to support high-throughput, multiprocess parallel chunk downloading to node-local cache (/dev/shm or temporary filesystem) followed by zero-copy memory-mapped deserialization (torch.load(..., mmap=True)).

Specifically:
1. Parallel Remote Chunk Fetching: For remote files ≥128 MB, pre-allocate a local cache file and download file chunks concurrently using ProcessPoolExecutor with up to 16 workers.
2. Node-Local Shm Caching & Synchronization: Cache checkpoints in /dev/shm (when sufficient space is available) or /tmp, protected by FileLock so multiple training processes on the same node share a single downloaded artifact safely without redundant downloads.
3. Zero-Copy Memory Mapping: Load the cached checkpoint via torch.load(..., mmap=True), allowing weight tensors to be backed directly by filesystem pages without duplicate heap allocations.
4. Resilient Cleanup: Ensure partial cache files are automatically purged if an exception or download failure occurs.

### 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 reading lightning.fabric.utilities.cloud_io._load() and trace how remote checkpoints are opened and passed to torch.load. Review the requested ProcessPoolExecutor, FileLock, local cache, and mmap behavior, then define tests for shared downloads, failed cleanup, and large remote files. Done means large checkpoints download in parallel, use safe node-local caching, load with mmap, and clean up partial files.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
cloud, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.