Add packed 4-bit SafeTensors weight support - #19
Merged
mtavenrath merged 6 commits intoSep 23, 2026
Merged
Conversation
Extend the shared external-weight resolver and writer to support packed logical Int4 and Uint4 tensors in versioned SafeTensors archives. Resolve ordinary and packed constants in one memory-mapped traversal, validate names, metadata, storage types, shapes, and lengths, and write archives atomically without an archive-sized output allocation. Document the .webnn external-weight contract and cover packed, mixed, malformed, sanitized-name, manifest, and raw-weight cases.
## Feature additions - Add a canonical `.webnn` format reference covering graph versions, shapes, initializers, node options, multi-output operations, output bindings, and `GraphJson` round trips. - Add a maintained v2 format example and parser → serializer → parser regression coverage. - Document SafeTensors, packed Int4/Uint4 storage, manifest-backed raw weights, sidecar discovery, mmap ownership, and atomic writer behavior. ## Bugfixes - Make the example build workflow reconstruct untracked tensor inputs from the checked-in weight archive. - Update the example workflow to use the current direct JavaScript emitter command. ## Refactors - Consolidate duplicated format material into canonical documentation pages. - Rewrite the README, dynamic-dimension guide, ONNX-lowering guide, and example documentation against current source behavior. - Remove stale model-specific guidance, unsupported quantitative claims, and obsolete agent-oriented files. ## Behavioral impact and compatibility - No library API, file-format, parser, serializer, or runtime behavior changes are intended. - The example build script now works without locally retained raw tensor files but produces the same documented artifacts. - The removed pages are replaced by canonical format and workflow documentation. ## Validation - `cargo fmt --all -- --check` passed. - `cargo test --all-features` passed, including 195 library tests and the format-reference and ResNet conversion integration tests. - `cargo test --no-default-features` passed, including 72 library tests and the format-reference integration test. - Strict Clippy passed for all targets with all features and without default features. - Parse, serialize, structural validation, manifest validation, pack, unpack, and JavaScript-emitter smoke checks passed. - Relative Markdown links, script syntax, line lengths, stale-language searches, and `git diff --check` passed.
## Feature additions - None. ## Bugfixes - Replace fixed-size `chunks_exact` calls with `as_chunks` across external-weight and ONNX byte decoding, satisfying the new Rust 1.98 Clippy lint enforced by CI. ## Refactors - None. ## Behavioral impact and compatibility - No intended externally observable behavior change; complete chunks are decoded as before and trailing incomplete bytes remain ignored. ## Validation - `RUSTUP_TOOLCHAIN=1.98.1 make fmt-check lint test release` passed with 197 tests. - `RUSTUP_TOOLCHAIN=1.98.1 CARGO_FLAGS=--no-default-features make fmt-check lint test release` passed with 73 tests. - `git diff --cached --check` passed.
mtavenrath
approved these changes
Sep 23, 2026
mtavenrath
pushed a commit
to rustnn/rustnn
that referenced
this pull request
Sep 23, 2026
## Summary RustNN changes to enable full roundtrip numeric validation of onnx2webnn conversion. Tied to [onnx2webnn #5](rustnn/onnx2webnn#5) and dependend on [webnn-graph #19](rustnn/webnn-graph#19). The reference in this PR needs to updated after the webnn-graph was merged. Make serialized WebNN graphs independently reloadable and executable while unifying graph recording and shape inference between `MLGraphBuilder` and the `.webnn` loader. Completed graphs now use an unambiguous shape model: `[]` always means a known rank-zero scalar, bounded dynamic dimensions remain explicit, and unresolved descriptors exist only as temporary internal inference state. ## Feature additions - Add `MLContext::rustnn_build_graph` as a direct compilation entry point for deserialized `GraphInfo`. This avoids constructing a second, unused `GraphRecorder` when compiling a graph reconstructed by the `.webnn` loader. - Add `run_onnx_path_with_inputs` so native ONNX Runtime execution can resolve external-data sidecars relative to the model file. This supports reference validation of filesystem-backed models whose weights are not embedded in the ONNX protobuf. - Serialize complete operation arguments, including reshape targets, Slice parameters, concat axes, permutations, and operand-valued options referenced by stable names. - Reload serialized graphs through the same recording and inference path used by `MLGraphBuilder`. - Add a versioned packed-4-bit Safetensors extension: - Logical `Int4` and `Uint4` dtype and shape remain in `.webnn`. - Packed low-nibble-first bytes are stored as Safetensors `U8`. - Archive metadata identifies the RustNN extension. - Reload validates marker, dtype, shape, byte length, and tensor-name resolution. - Support mixed ordinary and packed-4-bit constants in one Safetensors archive. ## Bugfixes - Infer and record every output of multi-output operations rather than only the first. - Prevent unresolved operands from silently becoming scalar descriptors. - Treat scalar GRU hidden states as rank-zero values and reject them through normal GRU rank validation. - Preserve scalar inputs, constants, intermediates, outputs, quantize/dequantize operands, and converter shape-map entries. - Always serialize required scalar `Reshape` and `Expand` targets as `newShape: []`. - Reject missing or malformed required shape-valued arguments explicitly. - Serialize operand-valued options by stable name instead of unstable numeric operand IDs. - Validate packed 4-bit logical element counts and exact packed storage lengths. - Correct Slice lowering for non-unit strides by deriving backend end indices from `start + extent`; this fixes both the ONNX and LiteRT lowering paths. ## Refactors - Introduce a private shared `GraphRecorder` used by both graph construction and JSON loading. - Centralize descriptor inference, operation insertion, dependency tracking, and graph-output marking. - Make operation insertion atomic: failed inference or validation no longer leaves partially recorded graph state. - Remove the historical empty-vector unknown-shape heuristic and obsolete temporary shape-table workaround. - Keep unresolved shapes represented through `Option` or missing internal entries rather than public `OperandDescriptor` values. - Update setup documentation with concise optional prerequisites for feature-specific backends. ## Behavioral impact and compatibility - `GraphRecorder` remains crate-private and is not a new public API. - `MLContext::rustnn_build_graph` and `run_onnx_path_with_inputs` are additive public entry points. - Every `OperandDescriptor` in completed `GraphInfo` now has a known shape: - `[]` is scalar. - Nonempty shapes may contain bounded `Dimension::Dynamic` entries. - Unknown shapes cannot escape graph construction. - The experimental `.webnn` format has intentional compatibility breaks: - Required shapes may no longer be omitted. - Scalar shape arguments must be serialized explicitly. - Direct input or constant graph outputs are rejected; an explicit operation such as Identity is required. - Ordinary Safetensors remain unchanged. Packed 4-bit tensors require the RustNN metadata marker and are rejected if presented as an unmarked or malformed extension. - Browser WebNN still does not natively execute 4-bit tensors. The ORT backend reconstructs an executable graph using its supported representations. ## Validation - Latest `cargo test --lib`: 379 tests passed. - Default, no-default-feature, `dynamic-inputs`, ONNX Runtime, and available backend-mock checks passed across the refactor. - Focused round trips passed for: - Scalar and bounded-dynamic descriptors. - Reshape, Expand, Slice, Concat, and Gemm bias. - Multi-output inference. - Even- and odd-sized `Int4` and `Uint4` constants. - Mixed native and packed-4-bit Safetensors. - Packed-4-bit graph compilation and ORT execution. - Malformed packed archive rejection. - Downstream onnx2webnn validation passed 743 tests and completed the 52-case skeleton and real-weight sweeps with the expected documented blockers. - Formatting and diff checks passed. - LiteRT regression coverage was added, but the LiteRT feature test could not run locally because `flatc` was unavailable. - Clippy remained blocked by pre-existing `-D warnings` failures outside this change.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR is used and depended on by RustNN #238 and onnx2webnn #5
Summary
Centralize
.webnnexternal-weight reading and writing inwebnn-graph. The resolver now supports versioned packed Int4/Uint4 SafeTensors storage, memory-maps external archives, and shares the same validation rules with a new atomic SafeTensors writer.Feature additions
rustnn.webnn.packed4=1metadata contract.write_external_weights_safetensors, which validates external declarations and writes ordinary and packed tensors through a same-directory temporary file..webnn,@weights, sidecar discovery, dtype mappings, packed-4-bit storage, naming rules, and compatibility guarantees.Bugfixes
Refactors
webnn-graphfor reuse by consumers such as RustNN.Behavioral impact and compatibility
resolve_external_weightsAPI remains compatible.rustnn.webnn.packed4=1and the low-nibble-first layout remain supported.ceil(logical_elements / 2); the.webnndeclaration remains authoritative for their logical dtype and shape.InlineBytes; externally backedGraphJsonstorage is not introduced.Validation
cargo fmt --all -- --check— passed.cargo test external_weights --no-default-features— 13 focused external-weight tests passed.cargo test --all-features— 195 unit tests and the ResNet integration test passed.cargo test --no-default-features— 72 tests passed.cargo clippy --all-targets --no-default-features -- -D warnings— passed.cargo clippy --all-targets --all-features -- -D warnings— reaches a pre-existingclippy::question_markfailure insrc/onnx/shape_inference.rs; no changed external-weight code is implicated.git diff --check— passed.Reviewer focus