Add Support for RMSNorm and optional QKV Biases to Onnx backends - #2427
Open
john-sp wants to merge 4 commits into
Open
Add Support for RMSNorm and optional QKV Biases to Onnx backends#2427john-sp wants to merge 4 commits into
john-sp wants to merge 4 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR extends the ONNX conversion/export path to handle newer attention weight formats by (1) exporting RMSNorm when LayerNorm betas are missing/empty (with an optional discrete fallback for runtimes lacking native RMSNormalization) and (2) supporting attention Q/K/V projections with optional (missing) bias tensors.
Changes:
- Add RMSNorm export support, including a discrete ONNX-node fallback and native RMSNormalization emission for opset 23.
- Detect “LayerNorm with empty betas” and export as RMSNorm instead.
- Support missing Q/K/V projection biases by inferring projection width from
q_wand skipping bias-add nodes when biases are absent.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/neural/onnx/converter.h | Adds alt_rmsnorm option to select discrete RMSNorm fallback behavior. |
| src/neural/onnx/converter.cc | Implements RMSNorm (native + fallback), auto-detects RMSNorm via empty betas, and makes Q/K/V bias adds optional. |
| src/neural/onnx/builder.h | Adds RMSNormalization(...) builder API. |
| src/neural/onnx/builder.cc | Extends supported opset range to 23 and emits native RMSNormalization nodes. |
| src/neural/backends/onnx/network_onnx.cc | Wires opset and alt_rmsnorm into runtime conversion option parsing. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
981
to
983
| converter_options.data_type = | ||
| WeightsToOnnxConverterOptions::StringToDataType(datatype); | ||
| converter_options.opset = opts.GetOrDefault<int>( | ||
| "opset", converter_options.data_type == | ||
| WeightsToOnnxConverterOptions::DataType::kBFloat16 | ||
| ? 22 | ||
| : 17); | ||
|
|
Comment on lines
+48
to
50
| if (ir < 3 || ir > 11) { | ||
| throw Exception("Only ONNX IR between 3 and 10 is supported."); | ||
| } |
Comment on lines
+397
to
+408
| // Only supported since opset 23 | ||
| std::string OnnxBuilder::RMSNormalization(const std::string& name, | ||
| const std::string& input, | ||
| const OnnxConst& scale, int axis, | ||
| float epsilon) { | ||
| auto* node = model_.mutable_graph()->add_node(); | ||
| auto out = PopulateStdNodeFields(node, name, input, "RMSNormalization"); | ||
| node->add_input(AddInitializer(name + "/w/scale", scale)); | ||
| AddIntAttribute(node, "axis", axis); | ||
| AddFloatAttribute(node, "epsilon", epsilon); | ||
| return out; | ||
| } |
Comment on lines
+561
to
569
| // Q/K/V biases are optional in newer weight files. The projection width is | ||
| // normally available from q_b, but must be inferred from q_w when it is | ||
| // omitted. ONNX's MatMul does not need a bias input, so simply skip the | ||
| // corresponding Add in that case. | ||
| const int d_model = layer.mha.q_b.empty() | ||
| ? layer.mha.q_w.size() / embedding_size | ||
| : layer.mha.q_b.size(); | ||
| const int depth = d_model / heads; | ||
|
|
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
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.
Add ONNX conversion support for attention networks that use RMSNorm and omit
Q/K/V projection biases.
By Menkib:
alt_rmsnormfallback that expresses RMSNorm using standard ONNX nodes, keeping it compatible with runtimes that do not support the nativeRMSNormalizationoperator.RMSNormalizationexport support for opset 23.By me:
q_wwhenq_bis absent, and omit Q/K/V bias-add nodes for missing bias tensors.Planned follow up PR: Updating TensorRT version to support native RMSNorm node
This has been tested locally on test networks.