feat: add w_sub ONNX semantic weight + node rename transform - #1230
Draft
vishwasdivakar wants to merge 1 commit into
Draft
feat: add w_sub ONNX semantic weight + node rename transform#1230vishwasdivakar wants to merge 1 commit into
vishwasdivakar wants to merge 1 commit into
Conversation
Adds opt-in semantic renaming for w_sub ONNX exports. When rename_onnx_nodes=True is passed alongside use_onnx_subfunctions=True, opaque ONNX names (e.g., onnx::MatMul_5898) are renamed to their HF equivalents (e.g., model.layers.0.self_attn.q_proj.weight) using 256-byte fingerprint matching against HF safetensors. This makes exported traces self-sufficient for llm_breakdown without a separate rename step. What it does: 1. RenameWsubTransform in onnx_transforms.py (+310 lines) - fingerprint-matches ONNX external weights to HF params, renames weights + anchored MatMul/RMSNorm nodes, stamps metadata. Handles both individual files and consolidated .onnx.data with offsets. Supports BF16/FP16/FP32 auto-detection. 2. Integration in modeling_qeff.py (+74 lines) - adds rename_onnx_nodes param (default False), resolves HF cache path, calls the transform after existing ONNX transforms. Warns if rename_onnx_nodes=True without use_onnx_subfunctions=True. 3. Kwarg threading in modeling_auto.py (+2 lines) - passes rename_onnx_nodes through to the export call. Safety: Default is False - zero impact on existing behavior unless explicitly opted in. Validated Models: LLaMA-3-8B, GPT-OSS 120B and Kimi-K2.5 on AI100
Author
|
@vbaddi, please review! |
ochougul
marked this pull request as draft
August 4, 2026 16:39
Contributor
vbaddi
requested changes
Aug 5, 2026
| "rename_onnx_nodes=True requires use_onnx_subfunctions=True. " | ||
| "Skipping semantic rename." | ||
| ) | ||
| elif rename_onnx_nodes and export_kwargs.get("use_onnx_subfunctions", False): |
Contributor
There was a problem hiding this comment.
export_wrapper removes use_onnx_subfunctions before _export, but here we check export_kwargs["use_onnx_subfunctions"] inside _export. That means rename_onnx_nodes=True would usually skip the actual rename.
| dynamo=dynamo, | ||
| offload_pt_weights=kwargs.get("offload_pt_weights", True), | ||
| prefill_only=prefill_only, | ||
| rename_onnx_nodes=kwargs.get("rename_onnx_nodes", False), |
Contributor
There was a problem hiding this comment.
rename_onnx_nodes was not part of the export hash. A renamed and non-renamed ONNX could collide in the same cache directory.
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| def _resolve_hf_cache_path(pretrained_model_name_or_path: str) -> Optional[str]: |
Contributor
There was a problem hiding this comment.
_resolve_hf_cache_path duplicated Hugging Face Hub cache resolution and could pick the wrong snapshot/revision. Export code should not infer model revision by scanning cache multiple times.
| FINGERPRINT_BYTES = 256 | ||
|
|
||
| # Regex: model.layers.N.<role>.weight -> extract <role> for node rename | ||
| _SEMANTIC_PARAM_RE = re.compile( |
Contributor
There was a problem hiding this comment.
we already have this bunch of data from model. Could you pls re-use.
| return fp_map | ||
|
|
||
| @classmethod | ||
| def _read_hf_fingerprint(cls, shard_path: str, param_name: str, target_dtype) -> Optional[bytes]: |
Contributor
There was a problem hiding this comment.
Generally:
- Pls remove
_resolve_hf_cache_path,rename_onnx_nodes,hf_model_path, andRenameWsubTransformcode which is causing the bloat. - Add a parameter identity renaming to
QEfficient/utils/torch_patches.py. - Add a minimal
RenameWsubNodesTransformfor local-functionMatMul,Gemm, andCustomRMSNormnodes. - We can register the node rename automatically for legacy
use_onnx_subfunctions=True. - Also, Add unit tests.
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.
Adds opt-in semantic renaming for w_sub ONNX exports. When rename_onnx_nodes=True is passed alongside use_onnx_subfunctions=True, opaque ONNX names (e.g., onnx::MatMul_5898) are renamed to their HF equivalents (e.g., model.layers.0.self_attn.q_proj.weight) using 256-byte fingerprint matching against HF safetensors. This makes exported traces self-sufficient for llm_breakdown without a separate rename step.
What it does:
Safety: Default is False - zero impact on existing behavior unless explicitly opted in.
Validated Models: LLaMA-3-8B, GPT-OSS 120B and Kimi-K2.5 on AI100