Proposal: API to Separate Graph Building from Weight Loading to Reduce Peak Memory Usage

[mtavenrath]

Problem Statement

While analyzing the memory-consumption of the Stable Diffusion demo on a dGPU system, I observed that the CPU memory used during graph building is more than 3x the actual model size. Similar high-memory-usage behavior (requiring multiple times the model size) is also observed on iGPU systems, although the exact multiplier may differ.

While this might be irrelevant for smaller models, it presents a significant overhead for larger ones. For example:

• A 4GB model can consume over 12GB of CPU memory during its load phase on a dGPU.
• This is already critical on systems with 16GB of total memory.

This issue is even more pronounced on UVM (Unified Memory) systems where the CPU and GPU/NPU share memory. This high peak usage during loading unnecessarily consumes memory that could otherwise be allocated for the model weights and intermediate buffers on the device.

---

Proposal

I'd like to propose introducing an API that splits graph creation from weight passing (i.e., loading the constant data).

The primary goal is to enable streaming weights directly into the graph during initialization. This would avoid creating copies of the weights in JavaScript or in the backend's staging memory before the graph is fully built. Ideally, the weight data would only be copied once, from its source buffer directly to its final (e.g., device) location when the data is provided to the graph after it has been built.

This raises a key question for discussion:

• Would all current WebNN backends be able to support a (mostly) "weightless" graph creation, where all tensor shapes and data types are known, but the actual weight data is not provided until a later step?

---

Potential Implementation Idea

One potential solution could look like this:

1. Allow builder.constant() to be called with just an MLOperandDescriptor (defining the shape and data type), without requiring the ArrayBufferView data source.
2. This call would return an MLOperand object that acts as a handle for this "hollow constant" (a tensor without its data).
3. Allow builder.build() to succeed using MLOperands that represent these hollow constants, creating the executable MLGraph.
4. Introduce a new method on the MLGraph interface, such as:
   graph.setConstantData(constantOperand, dataBuffer)
5. After the MLGraph is built (step 3), the user must call this new setConstantData() method for every hollow constant they created. At this stage, the backend can perform any required conversions and transfer the dataBuffer directly to device memory.
6. context.dispatch() would require all constants to be set. Attempting to call dispatch() before all hollow constants have been supplied with data (via setConstantData) would result in an error.

---

Expected Benefits

This change would significantly reduce peak memory pressure during initialization:

• dGPU Systems: In an ideal scenario, this could limit the peak CPU memory overhead to 1x-2x the size of the largest single tensor (for temporary buffering during upload), rather than 3x the entire model.
• iGPU/UVM Systems: The hope is that no temporary CPU-side storage would be needed for the "upload" (as it's shared memory). This would reduce the total peak CPU memory consumption down to roughly Model Size + Max Single Tensor Size.

[mtavenrath]

This issue has been discussed at the TPAC WebNN WG meeting

Key Discussion Points

Current API Capabilities

@reillyeon argued that the current API intends to allow the browser to upload a constant and discard the CPU copy immediately. However, current implementations (both browser-side and JS frameworks) do not handle this efficiently yet.

Backend Constraints

@huningxin pointed out that while DirectML supports "late binding" of weights, ONNX currently requires all weights to be present on the CPU during session creation.

Security vs. Performance

@RafaelCintron noted that browsers use a multiprocess architecture where models run in a different thread for security. While shared memory (zero-copy) is the goal, strict security boundaries make this challenging. Markus agreed that even reducing copies (rather than achieving zero-copy) is a critical first step.

Constant Tensors

There was a debate on whether the existing createConstantTensor API solves this. It was determined that while helpful, it requires backend support to avoid holding data in memory.

Potential Solutions & Next Steps

Spec Changes

@reillyeon suggested adding a streaming constructor for constants and a backpressure mechanism to let developers know when memory is free to load the next weight.

Ecosystem Coordination

The group agreed on the need to identify which operators strictly require weights at build time and to work with backend developers (specifically ONNX) to support weightless graph creation.

Refining Implementation

The consensus is that this is largely an implementation issue across the stack (JS Framework -> WebNN -> Backend) rather than just a spec issue.

[mtavenrath]

Optimization Potentials: Pull-Based Architecture for WebNN Constants

I had implementation discussions with @reillyeon and @huningxin regarding the shift from a Push-based model to provide constants where the App supplies a data buffer to a Pull-based model to provide constants by supplying URL & Offset/Size. The change in the API logic allows for the following optimizations.

1. Latency Hiding via Parallel Compilation

The pull-based API allows the graph topology to be fully defined while at the same time letting the WebNN implementation decide when to load the weights.

• The Optimization: The backend can initiate graph compilation concurrently with the weight download.
• The Benefit: For large models, the compilation time—which can be significant—is effectively "hidden" behind the network latency. Ideally, compilation completes by the time the download finishes.

2. Direct-to-Disk Caching & I/O Alignment

By delegating the download process to the backend, we enable low-level storage optimizations that are inaccessible to JavaScript.

• Sector Alignment: The backend can store downloaded weights with sector page-size alignment.
• Direct I/O (Bypassing FS Cache): This alignment satisfies the strict requirements for Direct I/O. The backend can read weights directly from the disk, bypassing the Operating System’s filesystem cache.
• The Benefit: This prevents large, read-once model files from "thrashing" the OS page cache (preserving system responsiveness) and allows the model to load at raw disk read speeds.

3. Persistent Layout Optimization

Current implementations often require CPU-intensive layout conversions (e.g., NCHW to hardware-specific block formats) every time a model is loaded.

• The Optimization: The backend can perform these layout conversions once—immediately after download and before writing to the persistent disk cache.
• The Benefit: Subsequent loads become strictly I/O bound. The backend reads pre-formatted data directly from the disk into memory, eliminating CPU overhead for layout shuffling during startup.

4. Memory Architecture & UVM Efficiency

A Pull-based model to provide constants minimizes the "High Water Mark" of RAM usage by keeping data flow internal to the native implementation.

• Bypassing CPU Allocations: This architecture avoids the large intermediate allocations typically required in the JavaScript heap and WebNN staging buffers.
• UVM (Unified Virtual Memory): On systems with shared memory architectures (e.g., integrated GPUs, integrated NPUs), the backend can download or map data directly into memory addressable by the accelerator, approaching a true "zero-copy" load path.

5. Dynamic Resource Management

When the backend owns the resource handle (URL/Offset), it can manage memory lifecycles more intelligently than the application can.

• Graceful Degradation: Under memory pressure, the backend can transparently evict weights from GPU memory (re-fetching from disk/network only when needed).
• Model Splitting: For devices with insufficient accelerator RAM to hold the entire model, the backend can orchestrate the swapping of model layers in and out of memory, enabling execution on constrained hardware.

[anssiko]

This topic was discussed at Kobe F2F where the group resolved the following:

RESOLUTION: explore streaming constructor for constants

[mtavenrath]

I had a discussion with @chrisjhebert1973 about the external weights and we were wondering if this feature would also allow remote execution of neural networks, e.g. on a home-server.

The browser has already a protocol to serialize the WebNN graph. Instead of passing the weights a download link is passed which could be used on the remote side. This should result in a quite small representation of the graph for remote compilation and execution.

This leaves only the question how much input/output data there is to make this a viable solution.

[fdwr]

Technically external weights are already achievable via MLTensor and connecting the graph operators to MLGraphBuilder.input() rather than MLGraphBuilder.constant(), where the MLGraph can be built with no weights, and the tensor data written later (but before execution time) via writeTensor. That doesn't mean "problem solved", or that it addresses all your concerns above just that it's possible, and so any considerations of this problem should look at that aspect too.

[reillyeon]

Some of the potential optimizations in #901 (comment) could also apply to WebGPU. I'm curious if @kainino0x or anyone else working on WebGPU has considered a similar ability to directly specify an HTTP request as the source for a resource.

It would be nice for this to cleanly integrate with the Fetch API, something like this:

let constant = builder.constant(new Request("model.bin", {headers: {"Range": "bytes=5435435-5484329"}}));

[kainino0x]

It's been discussed a bit in the context of gpuweb/gpuweb#4185, but we never did very much work to understand how data-streaming APIs on the Web work (existing or proposed). I know we also had some discussion about plain buffer uploads (not just image data) though I don't know if there are notes for that discussion.

Anyway, we are definitely interested in it.