Skip to content

[CHM-209] Add opt-in memory-mapped model loading - #41

Draft
enra64 wants to merge 1 commit into
mainfrom
arneherdick/chm-209-mmap-model-loading-gh
Draft

[CHM-209] Add opt-in memory-mapped model loading#41
enra64 wants to merge 1 commit into
mainfrom
arneherdick/chm-209-mmap-model-loading-gh

Conversation

@enra64

@enra64 enra64 commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

What

Exposes memory-mapped model loading through the wrapper, so weight storages can alias a read-only mapping of the .ptl instead of being copied onto the heap. Mapped pages stay clean and file-backed, so on iOS they do not count toward phys_footprint — the figure the OS compares against the per-process memory limit before terminating a process.

// per module
val module = TorchModule(path, useMmap = true)

// or set a process-wide default once at startup, so existing call sites are unchanged
TorchModuleDefaults.useMmap = true

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 a TorchModule is constructed, so changing it only affects modules loaded afterwards, which suits a rollout switch.

Platform support

Target Behaviour
iOS Honoured — maps MobileModuleLoadOptions::USE_MMAP
Android Accepted and ignored: LiteModuleLoader exposes no load options

Android 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_MMAP does not exist in upstream PyTorch. It comes from a patch in pytorch-lite-builds, so this needs a VoizeLibTorch-Lite built with it. Until that is published, useMmap = true has 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.

PLMMemoryInfo

Adds an accessor for the process's phys_footprint and 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.

- 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
enra64 force-pushed the arneherdick/chm-209-mmap-model-loading-gh branch from 38d21a3 to 799d04c Compare August 12, 2026 09:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant