Add option to `torch.load(mmap=True)` to do `MAP_SHARED` rather than `MAP_PRIVATE` #124528

mikaylagawarecki · 2024-04-19T22:13:43Z

🚀 The feature, motivation and pitch

Currently torch.load(mmap=True) calls mmap(2) with MAP_PRIVATE, this ensures that changes to the loaded checkpoint are not written back to the file.

As suggested by @albanD there might be a way to do postprocessing on checkpoints in CPU RAM-constrained environments if we pass MAP_SHARED instead

This would just involve passing shared=True to UntypedStorage.from_file

Alternatives

No response

Additional context

No response

cc @mruberry

The text was updated successfully, but these errors were encountered:

colinacassidy · 2024-04-20T18:57:46Z

Can we claim this task? Looking at getting into contributing to this project and this seems like a good first task. Should have a PR up in a day or so.

mikaylagawarecki · 2024-04-22T14:53:30Z

Hey @colinacassidy, thanks for your interest. Unfortunately I plan to pick this one up :( You could check out the actionable label for a list of issues open to be picked up!

albanD · 2024-04-22T18:02:29Z

I am not a fan of adding an extra argument to load() just for this. So maybe a torch.serialization.default_mmap_visibility would be the best API? Not sure if there are other settings of the mmap call we could expose the same way that might be useful?

Fixes #124528 [ghstack-poisoned]

Fixes #124528 Going over the options for our MapAllocator and what they do, I don't think any other of them need to be piped up to `torch.load` https://github.com/pytorch/pytorch/blob/4f29103749c5011529f1abb10b1508a682588909/aten/src/ATen/MapAllocator.h#L8-L16 However, I wonder if this `MmapVisibility(Enum)` is a good way to represent "or-ing" together of `mmap` flags if we want to extend it in the future. I looked over the flags for [`mmap(2)`](https://man7.org/linux/man-pages/man2/mmap.2.html), and could not immediately see how most of them would be useful for `torch.load` (would maybe `MAP_LOCKED` (like `mlock`) or `MAP_HUGE` ever be worthwhile?) cc albanD [ghstack-poisoned]

Fixes #124528 Going over the options for our MapAllocator and what they do, I don't think any other of them need to be piped up to `torch.load` https://github.com/pytorch/pytorch/blob/4f29103749c5011529f1abb10b1508a682588909/aten/src/ATen/MapAllocator.h#L8-L16 ~However, I wonder if this `MmapVisibility(Enum)` is a good way to represent "or-ing" together of `mmap` flags if we want to extend it in the future. I looked over the flags for [`mmap(2)`](https://man7.org/linux/man-pages/man2/mmap.2.html), and could not immediately see how most of them would be useful for `torch.load` (would maybe `MAP_LOCKED` (like `mlock`) or `MAP_HUGE` ever be worthwhile?)~ Using the flags provided by the python `mmap` library so that we can extend the allowed flags and pipe them down to the cpp `mmap` call if there is a need for other flags in the future cc albanD [ghstack-poisoned]

…True)" Fixes #124528 Going over the options for our MapAllocator and what they do, I don't think any other of them need to be piped up to `torch.load` https://github.com/pytorch/pytorch/blob/4f29103749c5011529f1abb10b1508a682588909/aten/src/ATen/MapAllocator.h#L8-L16 ~However, I wonder if this `MmapVisibility(Enum)` is a good way to represent "or-ing" together of `mmap` flags if we want to extend it in the future. I looked over the flags for [`mmap(2)`](https://man7.org/linux/man-pages/man2/mmap.2.html), and could not immediately see how most of them would be useful for `torch.load` (would maybe `MAP_LOCKED` (like `mlock`) or `MAP_HUGE` ever be worthwhile?)~ Using the flags provided by the python `mmap` library so that we can extend the allowed flags and pipe them down to the cpp `mmap` call if there is a need for other flags in the future cc albanD [ghstack-poisoned]