Skip to content

feat: optimize ONNX peak RSS with mmap'd external-data weights - #138

Merged
ivan-digital merged 2 commits into
soniqo:mainfrom
Avadhut2:feat/onnx-external-data
Aug 21, 2026
Merged

feat: optimize ONNX peak RSS with mmap'd external-data weights#138
ivan-digital merged 2 commits into
soniqo:mainfrom
Avadhut2:feat/onnx-external-data

Conversation

@Avadhut2

Copy link
Copy Markdown
Contributor

Summary

Optimizes ONNX peak RSS on memory-constrained devices by memory-mapping external data.

What Changed

  • Added scripts/export_external_data.py to extract and 4096-byte page-align heavy ONNX tensors.
  • Added documentation in models.md explaining how to bundle memory-mapped heavy models.
  • Updated AGENTS.md with guidelines making external data mandatory for models >= 100MB.

Test Plan

  • Validated python export script locally on dummy model with threshold=0.
  • Verified .onnx.data payload offset alignments to 4096-byte boundaries.
  • Confirmed session.use_device_allocator_for_initializers properly loads aligned format into mapped memory natively.

@ivan-digital

ivan-digital commented Aug 21, 2026

Copy link
Copy Markdown
Member

Thanks @Avadhut2 for the PR and for identifying the ONNX external-data approach. It produced a meaningful peak-memory improvement, and the follow-up needed for merge is now complete.

What changed:

  • Corrected the ONNX repacker: 64 KiB alignment, empty-tensor handling, bounded streaming for existing sidecars, atomic no-clobber publication, and public ONNX/protobuf APIs only.
  • Re-exported the variants with measurable benefit as additive external-v2/ graph/sidecar pairs while preserving the existing inline artifacts.
  • Updated download and inference paths to use immutable revisions, fetch sidecars before graph installation, and retain the filenames expected by the runtimes.
  • Added aligned bundles for Parakeet INT8, Whisper Small INT8, Canary INT8/FP32, Sortformer, and ReDimNet.
  • Fixed Sortformer symbolic prediction-axis handling and added concrete runtime output-shape validation.

Compared with the existing inline ONNX files, load-peak RSS was reduced by moving non-empty tensor weights into 64 KiB-aligned external-data sidecars. This lets ONNX Runtime map or read the weights separately from the graph protobuf, avoiding an additional large inline tensor payload during session initialization. Model architecture and precision were unchanged.

Measured load-peak RSS reductions:

  • Parakeet INT8: 701 MiB lower
  • Whisper Small INT8: 511 MiB lower
  • Canary INT8/FP32: 281/259 MiB lower
  • Sortformer: 208 MiB lower
  • ReDimNet: 79 MiB lower

WER parity on the same 50-utterance FLEURS English slice:

  • Parakeet: 5.92% → 5.92%
  • Whisper Small: 7.14% → 7.14%
  • Canary: 10.42% → 10.42%
  • All 150 paired hypotheses were identical.

Validation:

  • All 11 GitHub Actions jobs passed across Linux, macOS, Windows, sanitizers, ONNX packaging, LiteRT, examples, and aarch64
  • 13/13 repacker tests passed locally
  • 32/32 default tests passed locally, including pipeline e2e and stress
  • 41/41 ONNX tests passed locally with the real external bundles
  • Sortformer inline and external inference both passed
  • Shell syntax and ShellCheck passed

These results support keeping Parakeet as the general ASR default. Canary remains opt-in for its smaller bundle and multilingual transcription/translation support.

Thanks again for the contribution. I am merging with a normal merge commit so the original author commit remains in the project history.

@ivan-digital
ivan-digital merged commit ff1c3e2 into soniqo:main Aug 21, 2026
11 checks passed
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.

2 participants