You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit 067632d
Browse filesBrowse the repository at this point in the historyBrowse files
MLX: bound the lazy graph so a method's peak memory is not its whole graph (pytorch#22932)
Fixes the root cause of pytorch#22513.
## The problem
MLX is lazy. `Interpreter::dispatch` only builds graph nodes; nothing is
materialized until `MLXBackend::execute` calls `async_eval` on the
method outputs (`MLXBackend.cpp:557`). For a long instruction chain that
means **every intermediate in the method is live at the same instant**.
Whisper-small `encode` is 495 instructions built before a single byte is
evaluated. Measured on macOS with
`mlx::core::get_{active,peak,cache}_memory()`:
| | whisper-small encode |
| --- | --- |
| MLX peak | **1105.6 MB** |
| MLX active, steady | 94.8 MB |
| MLX cache | 262.3 MB (at the existing 256 MB cap) |
On iOS that peak lands in the app's footprint, which is what makes the
model unusable there.
## The fix
Evaluate the live per-execution tensors once the intermediates produced
since the last barrier exceed a byte threshold. Off by default; opt in
per model with the `eval_threshold_bytes` runtime option.
Each barrier costs a GPU sync, so the cost tracks the **number** of
barriers, not the interval. Budgeting bytes rather than counting
instructions puts barriers only in the methods that actually allocate:
whisper-small at 512 MB takes **12** barriers in `encode` and **0** in
`decode`. An every-32-instruction rule took 15 and 22, and those 22
bought 50 MB on a method that peaks at 258 MB while costing 21% on
device.
**iPhone 16, whisper-small int8, full pipeline, medians of interleaved
rounds:**
| arm | peak MB | peak-loaded | pipeline ms | encode | decode x10 |
| --- | --- | --- | --- | --- | --- |
| no barrier | 1194.4 | 763.7 | 885.2 | 463.2 | 472.6 |
| **eval_threshold_bytes 512 MB** | **692.8** | **261.0** | **831.3** |
**420.5** | **410.4** |
1.72x lower peak, 2.9x lower execute-phase footprint, and 1.06x
**faster** than the unbounded path.
**macOS, medians of 4 interleaved rounds x 10 executions, round 1
discarded:**
| model | method | peak off -> 512 MB | speed vs off |
| --- | --- | --- | --- |
| whisper-small | encode | 1105.6 -> 350.1 MB | 1.05x |
| whisper-small | decode | 258.0 -> 258.0 MB | 0.97x |
| whisper-tiny | encode | 550.4 -> 343.2 MB | 1.07x |
| whisper-tiny | decode | 61.8 -> 61.8 MB | 0.93x |
| SmolLM2-135M | forward | 1196.5 -> 496.8 MB | 1.00x |
Outputs are **bit-identical** to the unbounded path on whisper
tiny/base/small x {`encode`, `decode`} and on two MLX LLMs.
The threshold is the `eval_threshold_bytes` runtime option, read per
delegate handle alongside `clear_cache_interval` and validated at load.
It defaults to **0, which disables the mechanism entirely**: no
traversal, no `nbytes()` queries, no accumulation, no evaluation-root
collection. The measurements above are at 512 MB.
## What this is not
Two other candidates were measured and ruled out:
- **Buffer cache.** Already capped by `set_cache_limit(256MB)` at init.
Forcing `ET_MLX_CACHE_LIMIT_MB=16` moves cache 262 -> 103 MB and leaves
peak bit-identical at 1105.6 MB.
- **Slot retention.** `ExecutionState::tensors` holds only **7** slots
for those 495 instructions; the AOT side already plans and reuses them.
A last-use release built on `for_each_tid` freed 5 slots and moved peak
by 0 MB. The arrays are pinned by the pending graph, not the slot table.
cc @metascroy
---------
Co-authored-by: Scott Roy <scroy@meta.com>
0 commit comments