Load external data without copying it onto the JS heap - #1747
Load external data without copying it onto the JS heap#1747astefanutti wants to merge 9 commits into
Conversation
onnxruntime PR #29477, "[web] Support Blob-backed external data for on-demand loading in JSPI builds", merged on 2026-07-14. The pin here predates it (2026-04-16), so the runtime rejects a Blob handed to it as external data and the change in the next commit has nothing to hand it to. 1.29.0-dev.20260811-e415ef9afd is dated after that merge and is the build the change was developed and measured against.
Loading a model from Cache Storage currently materialises every external-data file into the JS heap before onnxruntime-web ever sees it: `loadResourceFile` reads the `Response` into a `Uint8Array`, and `getModelDataFiles` passes that buffer on. For a multi-gigabyte model that is a full second copy of the weights, live at the same moment the runtime is allocating its own. A Cache Storage `Response.blob()` is a file reference rather than a copy, and onnxruntime-web accepts one for external data as of the runtime bumped in the previous commit. So on a cache hit the bytes can go straight from disk to the runtime and never enter the heap at all. Measured on a 2.887 GB model, same runtime and same session: peak resident goes from 6,056 MB to 4,221 MB. The saving is the size of the weights, so it grows with the model. Two deliberate limits. It only applies on a cache HIT. Reading the stream to report download progress would defeat the purpose — the chunks would be resident twice, once as buffers and once in Blob storage — and while several gigabytes are arriving, progress is worth more than peak memory. Every load after the first takes the new path, which is where the peak actually matters, and is unchanged for Node, which keeps returning a path. And it is opt-in, threaded from the one caller that knows it is looking at external data. `loadResourceFile` is generic — it also serves config.json and tokenizer.json, whose callers parse the result as text — so returning a Blob unconditionally would break them.
…ing it The previous commit made a WARM load cheap: a cached file comes back as a Blob and its bytes never reach the JS heap. A cold one was untouched, and a cold one is the case that fails. `getModelDataFiles` starts every external-data chunk concurrently, and each is read into a `Uint8Array` so download progress can be reported — so a first load peaks at the SUM of the chunks rather than the largest. Measured on a 17 GB model from an empty cache: `Array buffer allocation failed` at 16.2 of 17.0 GB, 13 GB cached, load failed. It succeeds on a later attempt only because each file is cached as it completes, so the retry has less left to buffer. The body can go straight into Cache Storage instead. `cache.put` takes a `Response` and writes it to disk without materialising it, and reading it back with `cache.match` gives the same Blob a warm load gets — so the bytes travel network -> disk -> runtime and never sit on the heap. Progress survives, which is the only reason the buffer existed: a pass-through `TransformStream` counts bytes as they go past. It holds one chunk, not the file, and backpressure keeps it that way. If the cache refuses the write — QuotaExceededError being the expected one — it falls back to the buffered path. That has to re-fetch rather than reuse the response, whose body the failed attempt consumed.
nico-martin
left a comment
nico-martin
left a comment
There was a problem hiding this comment.
Hi @astefanutti, thank you so much for looking into this! I like the streaming direction, but the memory benefit is JSPI-only while the default runtime is still asyncify. The ORT lockfile also needs updating, and as_blob should be part of the in-flight key. Could you address those points and add focused warm/cold cache tests? I'm holding off on merging because the default path does not yet match the memory claim and concurrent callers can receive the wrong type.
| "@huggingface/tokenizers": "^0.1.3", | ||
| "onnxruntime-node": "1.24.3", | ||
| "onnxruntime-web": "1.26.0-dev.20260416-b7804b056c", | ||
| "onnxruntime-web": "1.29.0-dev.20260811-e415ef9afd", |
There was a problem hiding this comment.
Please regenerate pnpm-lock.yaml. The lockfile still resolves the old runtime, so frozen installs fail and validation may not exercise this load-bearing change.
| return_path, | ||
| // In the browser, hand onnxruntime-web a Blob rather than a materialised buffer. | ||
| // Node keeps returning a path, which is cheaper still. | ||
| !return_path, |
There was a problem hiding this comment.
Upstream on-demand Blob loading is JSPI-only, while Transformers.js defaults to asyncify, whose fallback fully materializes the Blob. Please select/gate a supported JSPI path or clearly separate and test the asyncify behavior from the claimed memory result.
| * @param {PretrainedOptions} [options] An object containing optional parameters. | ||
| * @param {boolean} [return_path=false] Whether to return the path of the file instead of the file content. | ||
| * @param {import('./cache.js').CacheInterface | null} [cache] The cache instance to use. | ||
| * @param {boolean} [as_blob=false] Whether to return a `Blob` when the file is served from the cache, |
There was a problem hiding this comment.
This is no longer cache-hit-only: the cold path streams to cache and returns the new Blob in the same call. Please update the parameter and return documentation.
| fatal = true, | ||
| options = {}, | ||
| return_path = false, | ||
| as_blob = false, |
There was a problem hiding this comment.
as_blob changes the promise result type and must be part of this key. Mixed concurrent callers currently share whichever representation was requested first.
`pnpm install --frozen-lockfile` fails on Node 18, 20 and 22 — eight seconds in, before anything is built or tested — because the ORT bump changed `packages/transformers/package.json` and left `pnpm-lock.yaml` pinning 1.26.0-dev.20260416-b7804b056c. Regenerated with `pnpm install --lockfile-only`, so the change is the resolution and nothing else.
`as_blob` changes the RESOLVED TYPE of `getModelFile`, not just how the
bytes are fetched, so two callers asking for the same file with different
values are not asking the same question. Left out of the key, whichever
one arrives first decides for both and the other receives a `Blob` where
it expects a `Uint8Array`, or the reverse — silently, and only under
concurrency. `return_path` is in the key already for exactly this reason.
Adds focused warm/cold tests against an in-memory cache and a stubbed
`env.fetch`, so they assert which branch ran rather than wall-clock
behaviour and touch no network:
· cold — streams into the cache, resolves a Blob, stores exactly once
(not a second time by the block at the end of loadResourceFile), and
still reports progress up to the total, which is the reason the
buffer existed;
· warm — resolves a Blob from the cache without re-fetching, and
reports the single completed event;
· a cache that refuses the write falls back to a buffer and re-fetches,
rather than failing the load;
· `as_blob` unset still resolves a Uint8Array;
· both orderings of a concurrent Blob/bytes pair get the type they
asked for.
The last two fail without the key change and pass with it — checked by
reverting it.
The review is right that the claim and the default path do not match, and
the reason is that this PR has TWO savings which the comments ran
together:
· The DOWNLOAD no longer buffers. Streaming the body into Cache Storage
and reading it back means the file never sits on the JS heap on its
way in. This happens before onnxruntime-web sees anything and holds
whichever build is in use. It is the half that stops a cold load of a
model with several multi-gigabyte chunks dying in getModelDataFiles,
where every chunk starts concurrently and the peak is their sum.
· The SESSION may or may not copy. Handing the runtime a Blob only
avoids materialising it if the runtime can read it that way: the JSPI
build can, the default asyncify build cannot.
So "the bytes never sit on the JS heap at all" was true of the download
and overstated end to end. The comments now say which is which, and say
plainly that on the default build this buys the download and not the
session.
No behaviour change — comments only.
`pnpm build` runs `tsc --build`, and it has been failing on this branch
since the Blob was introduced — five errors, unnoticed because CI never
got past `pnpm install`. Fixing the lockfile in the first commit would
have moved the red from install to build rather than clearing it.
· `buffer` was annotated `Uint8Array` and now legitimately holds a Blob
on both `as_blob` branches. Widened rather than cast away at each
use, so a future branch that forgets which it is fails at build.
· `cache.match` is typed for every backend including the Node file
cache; only a `Response` has `.blob()`, and only a `Response` reaches
that line. Cast instead of the blanket `any` it had.
· `INFLIGHT_LOADS` holds Blob-resolving promises now.
· `getCoreModelFile` never passes `as_blob`, so a Blob cannot come back;
narrowed there rather than widening every caller of a function that
only ever wants bytes or a path.
· The store-at-the-end block cannot see a Blob either — a warm read is a
cache hit, which never sets `toCacheResponse`, and a cold stream
writes the entry itself and then clears the flag. The narrowing says
so, so a third branch that produced a Blob without storing it would
land there loudly.
tsc --build is clean, the bundles build, and tests/utils is 296 passing.
|
Thanks a lot @nico-martin for the review. I've addressed your comments as extra commits rather than a rewrite so the review stays readable: Lockfile (
Warm/cold tests ( The memory claim vs the default runtime (
So "the bytes never sit on the JS heap at all" was true of the download and overstated end to end. The comments now say which is which, and say plainly that on the default build this buys the download and not the session. No behaviour change. If you'd prefer the PR claim only the download half, I'm happy to retitle it accordingly. One more ( |
Two review points that the last round left open. `as_blob`'s documentation still said cache-hit-only on both `loadResourceFile` and `getModelFile`, which stopped being true when the cold path started streaming into the cache and returning the new Blob in the same call. Both now name the two paths it IS honoured on, and say that a Uint8Array comes back otherwise — a cold file with no cache, and the fallback when the cache refuses the write — so callers know they have to handle both. And the asyncify/JSPI split is now tested, not just asserted in a comment. Whether onnxruntime-web copies the Blob at session creation depends on its build and is not observable from this layer; what is observable, and is the saving this file is responsible for, is that the DOWNLOAD never materialises the body. The new case spies on the fetched response and asserts `arrayBuffer()` is never called on the streaming path — paired with the buffered path, where the same spy must SEE the call, so it cannot pass by being wired to the wrong object.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
nico-martin
left a comment
nico-martin
left a comment
There was a problem hiding this comment.
Thanks for the thorough follow-up! LGTM!
3 participants
Description
Loading a model's external data currently routes every byte through the JS heap. On a cache hit,
loadResourceFilereads theResponseinto aUint8ArrayandgetModelDataFilespasses that buffer to onnxruntime-web, which allocates its own — so a multi-gigabyte model is resident twice at the same moment. On a cache miss the same buffer is built in order to report download progress. AndgetModelDataFilesstarts every chunk concurrently, so the peak is the sum of the chunks, not the largest one.For a model with eight ~2 GB chunks that is ~16 GB of
Uint8Arraybefore the runtime has allocated anything, and it fails.A Cache Storage
Response.blob()is a file reference rather than a copy, and onnxruntime-web accepts one for external data as of #29477 (merged 2026-07-14).cache.put()likewise takes aResponseand streams it to disk without materialising it. Between them the bytes can go network → disk → runtime and never touch the heap.Measured
Warm load — same model, same runtime, same session, peak RSS sampled at 200 ms, external data passed as
Uint8ArrayvsBlob, on a 2.887 GB model:The 1,835 MB saved is one full copy of the weights, so it scales with model size. (Caveats: one session rather than the two an app may load concurrently, wasm EP so the GPU's own residency is not counted, and summed RSS double-counts pages shared between processes. The ratio is the finding; the absolutes are indicative.)
Cold load — a 17 GB model (8 external-data files, largest 1.996 GB), empty cache, first attempt:
Array buffer allocation failedBefore this, a first load of a model that size fails and only completes on a later attempt, because each file is cached as it completes so the retry has less left to buffer. Progress reporting still works throughout (1.6 → 4.2 → 6.8 → 9.2 → 11.8 → 14.3 → 16.1 → 17.0 GB).
That model has also run three full conversation suites through the new path — 84 generated answers, no failures.
Design notes
Progress is why the buffer existed, so the cold path keeps it with a pass-through
TransformStreamthat counts bytes as they go past. It holds one chunk, not the file, and backpressure keeps it that way.It is opt-in, threaded from the one caller that knows it is looking at external data.
loadResourceFileis generic — it also servesconfig.jsonandtokenizer.json, whose callers parse the result as text — so returning aBlobunconditionally would break them. Hence theas_blobparameter rather than a change in behaviour. Node is unaffected and keeps returning a path.If the cache refuses the write (
QuotaExceededErrorbeing the expected one) it falls back to the buffered path. That has to re-fetch rather than reuse the response, whose body the failed attempt consumed; the re-fetch goes throughgetFile()so a gated repo keeps itsAuthorizationheader.The onnxruntime-web bump is load-bearing, not incidental. The current pin (2026-04-16) predates #29477, so the runtime materialises any Blob handed to it — which is the peak this removes, and is also why it throws over the 2 GiB
ArrayBufferlimit. The new pin is dated after that merge and is what the numbers above were measured against. The package already pins a dev build, so this is a bump rather than a new kind of dependency.Not covered
getModelDataFilesstill starts every chunk concurrently. That no longer costs heap, but it does mean N simultaneous connections; serialising them is a separate question and not addressed here.