Lightning-AI / Lightning-AI/pytorch-lightning
Optimize remote checkpoint loading with parallel multiprocess downloads and zero-copy mmap
Nobody has claimed this yet.
- 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
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 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