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

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