[CHM-209] Add opt-in memory-mapped model loading - #41
Draft
enra64 wants to merge 1 commit into
Draft
Conversation
- Exposes MobileModuleLoadOptions::USE_MMAP from the patched LibTorch-Lite as `useMmap`, so weight storages alias a read-only mapping of the .ptl instead of being copied onto the heap. Mapped pages stay clean and file-backed, so they do not count against the iOS per-process memory footprint. - Defaults to off, via TorchModuleDefaults so a host app can switch it from one place instead of threading it through every model that loads a TorchModule. - The option is detected at compile time rather than named directly, so the wrapper still builds against an unpatched LibTorch-Lite (including the published pod); there it resolves to 0 and logs that mmap is unavailable instead of silently loading the copying way. - Honoured on iOS only; Android accepts and ignores it, as LiteModuleLoader exposes no load options. - Adds PLMMemoryInfo so callers can measure phys_footprint and verify the effect rather than assume it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
enra64
force-pushed
the
arneherdick/chm-209-mmap-model-loading-gh
branch
from
August 12, 2026 09:57
38d21a3 to
799d04c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Exposes memory-mapped model loading through the wrapper, so weight storages can alias a read-only mapping of the
.ptlinstead of being copied onto the heap. Mapped pages stay clean and file-backed, so on iOS they do not count towardphys_footprint— the figure the OS compares against the per-process memory limit before terminating a process.Defaults to off. The process-wide default exists so a host app can drive this from one place — e.g. a remote feature flag — instead of threading the flag through every class that loads a
TorchModule. It is read when aTorchModuleis constructed, so changing it only affects modules loaded afterwards, which suits a rollout switch.Platform support
MobileModuleLoadOptions::USE_MMAPLiteModuleLoaderexposes no load optionsAndroid would additionally need the patch ported to its own libtorch build, which is a different version from the iOS one.
Requires a patched LibTorch-Lite
USE_MMAPdoes not exist in upstream PyTorch. It comes from a patch inpytorch-lite-builds, so this needs aVoizeLibTorch-Litebuilt with it. Until that is published,useMmap = truehas no effect — and the ObjC change will not compile against an unpatched LibTorch-Lite, since it references the option.Safety
Enabling it cannot break loading. Only records serialized uncompressed and aligned can be mapped; anything else is copied as usual, decided per record. Archives written by PyTorch qualify — its writer stores records uncompressed and 64-byte aligned precisely so they can be mapped — while one repacked by an ordinary zip tool may not.
The mapping is
PROT_READ, so an in-place write to a weight faults rather than silently dirtying the page and giving the memory saving back.Weights that a backend repacks at load time — byte-quantized embedding tables, XNNPACK-prepacked layers — are materialized on the heap regardless, so models dominated by those save proportionally less.
PLMMemoryInfoAdds an accessor for the process's
phys_footprintand resident size, so callers and tests can assert on the effect rather than assume it. Worth having because the intuitive metric is misleading here: mapped weights stay resident, often at 100%, while contributing nothing to footprint. Reasoning from resident size suggests the optimization did nothing.Verification
Against 15 real TorchScript archives, on a simulator: every mappable record byte-compared against the copying reader — 715 MB of record bytes, zero mismatches, bit-identical tensor hashes in both modes, no load failures. A forward pass produces identical output either way, and the footprint advantage survives inference. With all 15 models loaded in one process: 845 MB footprint growth copying vs 330 MB mapped.
No model re-export is required: every tensor record in all 15 was already stored-uncompressed and 64-byte aligned, with zero unaligned records.
Status: draft — one result is unexplained
A physical-device measurement does not yet agree with the simulator: in a whole-app trace, total dirty memory fell by only ~51 MiB where ~500 MB was expected, while ~500 MiB of clean mapped file appeared. Either the aliasing is not taking effect in that build, or the comparison run had fewer models resident. Under investigation — this should not merge, and the flag should not be enabled anywhere, until that is understood.