MLX: bound the lazy graph so a method's peak memory is not its whole graph - #22932
msluszniak wants to merge 5 commits into
Conversation
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
…graph Backport of pytorch/executorch#22932. MLX is lazy: dispatch() only builds graph nodes and nothing is materialized until MLXBackend::execute calls async_eval on the outputs, so every intermediate in a method is live at once. Whisper-small encode (495 instructions) peaks at 1105.6 MB against 94.8 MB steady active; on iOS that peak lands in the app's footprint. Evaluating every N instructions bounds the pending graph: tiny 550 -> 166 MB, base 1071 -> 230 MB, small 1106 -> 443 MB, at 1.01-1.08x the speed and with bit-identical output. ET_MLX_EVAL_EVERY overrides; 0 restores the old path.
…graph
MLX is lazy: Interpreter::dispatch only builds graph nodes, and nothing is
materialized until MLXBackend::execute calls async_eval on the method outputs.
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 with mlx::core::get_{active,peak,cache}_memory on macOS:
peak 1105.6 MB against 94.8 MB of steady-state active memory. On iOS that peak
lands in the app's footprint and is what makes the model unusable there
(pytorch#22513).
Evaluate the live per-execution tensors once the intermediates produced since
the last barrier exceed a byte budget. Each barrier costs a GPU sync, so the
cost tracks the NUMBER of barriers; budgeting bytes rather than counting
instructions puts them only in the methods that actually allocate. Whisper-small
at 512 MB takes 12 barriers in encode and 0 in decode.
iPhone 16, whisper-small int8 through the full pipeline, medians of interleaved
rounds:
peak MB peak-loaded pipeline ms encode decode x10
no barrier 1194.4 763.7 885.2 463.2 472.6
byte budget 512MB 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:
peak MB (off -> 512) speed vs off
small encode 1105.6 -> 350.1 1.05x
small decode 258.0 -> 258.0 0.97x
tiny encode 550.4 -> 343.2 1.07x
tiny decode 61.8 -> 61.8 0.93x
SmolLM2 fwd 1196.5 -> 496.8 1.00x
Outputs bit-identical to the unbounded path on whisper tiny/base/small x
{encode, decode} and on two MLX LLMs. ET_MLX_EVAL_BUDGET_MB overrides the
budget; 0 restores the previous behaviour.
debdde0 to
aa47f06
Compare
nil-is-all
added
the
module: mlx
|
Thanks for the PR @msluszniak! Could we expose this as a per-handle runtime backend option named eval_threshold_bytes, rather than an env-only setting? MLX already has the LoadBackendOptionsMap / get_runtime_spec pattern for clear_cache_interval. I’d follow that pattern with kEvalThresholdBytesKey, validate the value at load time, and configure the interpreter before the init chain runs. Also document this new option in the our docs. The default should be 0 (disabled), preserving existing behavior. Disabled should mean no meaningful overhead from this mechanism: no tensor traversal, nbytes() queries, byte accumulation, or evaluation-root collection. Please guard the accounting itself, not just the final evaluation. Existing op-internal eval calls should remain unchanged. For nested chains, we could share size_t& pending_bytes through run_chain, exec_if, and exec_scan. Ordinary dispatch handlers wouldn’t need to change; the surrounding loop would accumulate their estimates only when the option is enabled. Add an accumulate_only parameter to run_chain, with nested calls passing true: After the existing SCAN/IF/dispatch block, the outer loop would look conceptually like this; helper names are illustrative: This avoids double-counting IF/SCAN’s child instructions at the parent level. SCAN’s final stack operation should be accounted for separately, also guarded by the option being enabled. Accumulation should be overflow-safe. accumulate_only suppresses only threshold-triggered evaluation, not op-internal evaluations. Deferring the threshold check until exec_scan returns means its collected outputs have been stacked into state and are reachable from the evaluation roots. Evaluating inside the body could otherwise reset the shared counter while leaving earlier outputs retained only in collected unevaluated. “Threshold” is intentional: this is best-effort evaluation scheduling, not a hard memory limit. A long SCAN or IF branch can exceed the threshold before returning, and internal evaluations can make the estimate overcount pending work. Please document those limitations. I’d add focused tests for nested accumulation and threshold triggering, output equivalence, independent per-handle settings, invalid values, and the disabled path, specifically verifying that disabled execution never invokes the accounting helper. Representative memory and latency measurements should validate the enabled behavior. |
Replaces the ET_MLX_EVAL_BUDGET_MB env var with a per-handle eval_threshold_bytes runtime spec, following the clear_cache_interval pattern. It is read in init() before the init chain runs, validated there, and defaults to 0 (disabled). Disabled now means no work at all: the traversal, the nbytes() queries, the accumulation and the evaluation-root collection all sit behind the threshold check, not just the evaluation itself. Nested chains share the caller's counter through run_chain's new pending_bytes/accumulate_only parameters, so IF branches and SCAN bodies accumulate but do not trigger an evaluation of their own; the enclosing chain checks once control returns to it, by which point a SCAN's collected outputs are stacked into state and reachable from the roots. SCAN and IF are not charged at the parent level, since their children already were. SCAN's final stack is charged separately. Accumulation saturates so it cannot wrap back under the threshold. Documents that this is a threshold and not a hard limit, and the three ways peak can exceed it.
|
@metascroy could you look if now the implementation follows the established flow? |
|
Eh, apparently I cannot fix the linter 😅, my bad, will fix this in sec. |
The new test block was inserted between that comment and its et_cxx_test, which also left cmake-format wanting to reflow the two comments into one.
|
Please update the PR description to match the implementation: eval_threshold_bytes is now a runtime option, disabled by default. The description still refers to ET_MLX_EVAL_BUDGET_MB. |
|
I'll let @metascroy approve, but looks good to me |
BackendOptions comes from runtime/backend/options.h; backend_options.h only carries the key.
4 participants
Fixes the root cause of #22513.
The problem
MLX is lazy.
Interpreter::dispatchonly builds graph nodes; nothing is materialized untilMLXBackend::executecallsasync_evalon 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
encodeis 495 instructions built before a single byte is evaluated. Measured on macOS withmlx::core::get_{active,peak,cache}_memory():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_bytesruntime 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
encodeand 0 indecode. 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:
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:
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_bytesruntime option, read per delegate handle alongsideclear_cache_intervaland validated at load. It defaults to 0, which disables the mechanism entirely: no traversal, nonbytes()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:
set_cache_limit(256MB)at init. ForcingET_MLX_CACHE_LIMIT_MB=16moves cache 262 -> 103 MB and leaves peak bit-identical at 1105.6 MB.ExecutionState::tensorsholds only 7 slots for those 495 instructions; the AOT side already plans and reuses them. A last-use release built onfor_each_tidfreed 5 slots and moved peak by 0 MB. The arrays are pinned by the pending graph, not the slot table.cc @metascroy