feat-dynamo(0713): Dynamo based ONNX export with subfunction support for Causal-LM models by vbaddi · Pull Request #1178 · quic/efficient-transformers

vbaddi · 2026-07-13T06:01:46Z

Summary

This PR adds dynamo=True export support for CausalLM models, enabling torch.export-based ONNX export with optional repeated-subgraph subfunction extraction (use_onnx_subfunctions=True). It also includes a suite of robustness fixes to the ONNX transform pipeline discovered during integration.

What changed

New infrastructure

QEfficient/utils/custom_op_utils.py: select_interface helper that switches between the torch.autograd.Function path (TorchScript/eager) and torch.ops.qefficient.* path (dynamo) at call time, enabling the same model code to work under both exporters without code duplication.
QEfficient/utils/export_utils.py: export_wrapper refactored into four clearly separated export paths (plain TorchScript, TorchScript+subfunctions, flat dynamo, dynamo+subfunctions). use_dynamo is now an explicit parameter throughout rather than being re-read from kwargs. TorchScript monkey-patches are guarded behind if not use_dynamo: so they are never applied on the dynamo path. Added PyTorch version guard that raises a RuntimeError with install instructions when use_dynamo=True and torch < 2.13.
QEfficient/utils/torch_patches.py: Added temporarily_disable_nested_compile_regions context manager (Path 3) that unwraps static @torch.compiler.nested_compile_region decorators for flat-dynamo export, and temporarily_enable_nested_compile_regions reserved for future use.
QEfficient/base/modeling_qeff.py:
- _prune_unused_fake_initializers — removes FakeTensor-valued initializers that dynamo injects but are not referenced in the graph, preventing downstream ONNX load failures.
- _export extended with dynamic_shapes parameter and full dynamo export path: sets dynamo=True, injects custom_translation_table mapping all torch.ops.qefficient.* ops to their ONNX counterparts, and reorders example_inputs / dynamic_shapes to match the model's forward signature before passing to torch.onnx.export.
QEfficient/base/onnx_transforms.py: Two new transforms for the dynamo+subfunctions path:
- PreserveNestedCacheRetainedStateTransform — exposes CtxScatter outputs from repeated-subgraph function bodies as explicit ONNX outputs and renames KV cache graph inputs from past_key.N_RetainedState → past_key.N so the QAIC compiler can pair inputs and outputs correctly.
- RenameRepeatedSubgraphTransform — renames repeated_subgraphN ONNX Local Functions to the actual decoder layer class names, with multi-pattern matching across known PyTorch naming conventions and logger.warning on mismatch or count mismatch.
QEfficient/transformers/models/modeling_auto.py: convert_dynamic_axes_to_dynamic_shapes converts ONNX dynamic_axes dicts to torch.export Dim-keyed dynamic_shapes dicts, handling nested past_key_values, compressed KV (DeepseekV3), and CB pool dimensions.
QEfficient/utils/constants.py: ONNX export opset bumped from 17 → 18, required by the dynamo ONNX exporter.
pyproject.toml: Updated to torch 2.13.0 stable CPU wheels (minimum required for dynamo export); explicit per-python per-arch wheel URLs for x86_64 and aarch64. onnxscript updated to 0.6.2, compressed-tensors updated to 0.17.0.

CausalLM model wrappers (llama, mistral, qwen2, qwen3, gemma, gemma2, gpt2, gpt_bigcode, gptj, falcon, mpt, phi, phi3, codegen, starcoder2, granite, granitemoe, mixtral_moe, qwen3_moe, olmo2, glm4_moe, grok_1, deepseek_v3, gpt_oss, llama_swiftkv):

QEffDecoderLayer.forward decorated with @torch.compiler.nested_compile_region so dynamo emits repeated-subgraph ONNX Local Functions instead of inlining each layer individually.
torch.tensor(MIN_MASKED_ATTENTION_VALUE) replaced with torch.full_like(attn_weights, ...) throughout — removes a scalar constant that dynamo cannot trace as a dynamic value and that caused shape inference failures during export.
select_interface wired into all custom op call sites (CtxScatter, CtxGather, CustomRMSNorm, and CB/3D variants) so both the TochTrace and dynamo paths dispatch to the correct backend op without branching in model code.

vbaddi · 2026-07-14T08:19:44Z

This review needs to be completed by end of this week to have this in mainline.
@quic-rishinr @ochougul @quic-amitraj can you pls help review this asap and provide feedback. thanks

Re: Backward compatibility, we can revert the torch packages from .toml and have it specific to 2.7.0 and dynamo specific users can enable required packages (torch, onnxscript, compressed-tensors) offline for now. @smedhe

quic-amitraj · 2026-07-14T20:13:58Z

Please remove all the commented code from all the files.

ochougul · 2026-07-16T08:06:32Z

+            if getattr(self.model.config, "model_type", None) == "gpt_oss" and not self.continuous_batching:
+                bs = 1


why is this?

Signed-off-by: Sharvari Medhe <smedhe@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

Solved review comments for #1178 --------- Signed-off-by: Sharvari Medhe <smedhe@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

Review comment addressed for #1178 --------- Signed-off-by: smedhe <smedhe@qti.qualcomm.com> Signed-off-by: Sharvari Medhe <smedhe@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

## Stacked On This PR is stacked on #1194 by targeting `quic:quic-enable-dynamo-for-causallm`. ## Problem Legacy TorchScript export and dynamo export need different ONNX opsets. The previous code used a single `ONNX_EXPORT_OPSET` value everywhere, so both export paths, custom-op symbolic registration, ONNXScript function protos, and the dynamo custom translation table all effectively shared the same opset-backed custom op definitions. For this branch we need: - legacy export (`dynamo=False`) to stay on ONNX opset 17 - dynamo export (`dynamo=True`) to use ONNX opset 18 - QEff custom ops (`CtxScatter`, `CtxGather`, `CustomRMSNorm`, etc.) to emit function bodies built against the matching default-domain opset ## Approach Instead of duplicating `ctx_scatter_gather.py`, `ctx_scatter_gather_cb.py`, and `rms_norm.py` into separate opset-17/opset-18 files, this PR keeps one source definition for each custom op and compiles two ONNXScript variants from it. The helper in `QEfficient/customop/onnxscript_utils.py` temporarily binds the module-level `ops` symbol to the requested default ONNX opset while the ONNXScript decorator compiles the function. The default returned function is the legacy/opset-17 variant, and the dynamo/opset-18 variant is attached to it for explicit lookup. This avoids code duplication while ensuring both export modes get different ONNXScript function protos. ## Key Changes - Adds explicit constants: - `ONNX_LEGACY_EXPORT_OPSET = 17` - `ONNX_DYNAMO_EXPORT_OPSET = 18` - `get_onnx_export_opset(dynamo)` - Updates legacy export paths to pass `opset_version=17`. - Updates dynamo export to pass `opset_version=18`. - Makes `CustomOpTransform.apply(..., onnx_export_opset=...)` mode-aware so it registers custom symbolics and appends ONNXScript function protos for the selected opset. - Passes the selected opset from export into the ONNX transform pipeline. - Updates `DYNAMO_CUSTOM_OP_TABLE` to point at the opset-18 ONNXScript variants, because dynamo uses that table during export before post-export transforms run. - Adds a regression test proving the same `CustomRMSNorm` source can produce opset-17 and opset-18 function protos. ## Why This Design Duplicating the custom-op files would work, but it creates long-term maintenance risk: every custom-op behavior fix would need to be made in two places. This PR centralizes only the opset selection while keeping the custom-op implementation itself single-source. The mode boundary is explicit at export time and transform time. --------- Signed-off-by: Sharvari Medhe <smedhe@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com> Co-authored-by: Sharvari Medhe <smedhe@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

Signed-off-by: Sharvari Medhe <smedhe@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

The PR removed _layerwise_safe_export_passes_enabled() guard from layerwise_safe_onnx_export_patches, causing peephole ONNX passes to be disabled for all TorchScript exports. Phi3's SplitToSequence_182 was no longer folded to a constant Split, making the QAIC compiler reject it with 'Non-constant split tensor not supported' when -sub-functions was active. Gate the patches in _export_via_legacy by _layerwise_active, restoring the original behavior: passes are disabled only during layerwise export. Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

…ments (#1211) fix#1178 :solving multi device compile issue and addressed review comments --------- Signed-off-by: Sharvari Medhe <smedhe@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

…ged two stage of dynamo testing into one Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>

quic-amitraj · 2026-07-23T20:16:59Z

`gpt_oss` disabled in `test_dynamo_cb_generate`

Temporarily skipped gpt_oss in tests/dynamo/test_on_qaic.py::test_dynamo_cb_generate to unblock CI.

Failure: The QAIC compiler rejects GptOssForCausalLM.onnx when -sub-functions is active:

vbaddi assigned vbaddi and smedhe Jul 13, 2026

vbaddi added enhancement New feature or request onnx.dynamo 1.23 Release 1.23 Features labels Jul 13, 2026

amarquic mentioned this pull request Jul 13, 2026

feat(weight-free): weight-free ONNX export for causal LMs (dynamo path) #1181

Closed

3 tasks

vbaddi requested review from anujgupt-github, ochougul, quic-amitraj, quic-hemagnih and quic-rishinr July 14, 2026 08:04