freenode
AI & ML

PyTorch RFC proposes opt-in mmap for distributed checkpoint loads

FileSystemReader could map local .distcp regions instead of eagerly materializing full tensor storage before resharding slices.

PyTorch developers have opened an RFC to add an opt-in memory-mapping path for local distributed checkpoint reads, aimed at cutting the cost of loading large .distcp files when only a slice of each tensor is needed.

Today, FileSystemReader materializes an object's full tensor storage before the planner narrows it to the requested region. That is especially expensive during resharding, when ranks often need only part of a storage. A .distcp file can hold multiple independent torch.save ZIP objects; the runtime already knows each object's path, offset, and length, but torch.load has no clean way to mmap a bounded region of an existing mapping.

The proposal adds a keyword-only mmap=False flag on FileSystemReader. When enabled for eligible local tensor loads, storage becomes a view into a privately mapped checkpoint file rather than a fully copied buffer. The planner, resharding logic, destination copy path, and on-disk format stay the same. Unsupported layouts fall back to the existing loader; corrupt or out-of-bounds regions raise errors.

A prototype implementation is available for review. Feedback is sought on the API shape, bounded-ZIP validation, fallback boundaries, and the trade-off of whole-file mapping.