diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 2b05a96f05..1f7f0305c7 100755 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -367,7 +367,7 @@ def _export_via_legacy( output_names=output_names, dynamic_axes=dynamic_axes, dynamo=False, - opset_version=constants.ONNX_EXPORT_OPSET, + opset_version=constants.ONNX_LEGACY_EXPORT_OPSET, **export_kwargs, ) @@ -416,7 +416,7 @@ def _export_via_dynamo( output_names=output_names, dynamic_axes=None, dynamic_shapes=dynamic_shapes, - opset_version=constants.ONNX_EXPORT_OPSET, + opset_version=constants.ONNX_DYNAMO_EXPORT_OPSET, **export_kwargs, ) if onnx_program is None: @@ -580,6 +580,7 @@ def _resolve_pkv_names(layer_idx, layer_state): "onnx_base_dir": str(export_dir) if needs_external_tensor_data else None, "model_name": self.model_name, "dynamic_axes": None if dynamo else dynamic_axes, # dynamo uses dynamic_shapes, not axes + "onnx_export_opset": constants.get_onnx_export_opset(dynamo), } if onnx_transform_kwargs is not None: transform_kwargs.update(onnx_transform_kwargs) @@ -880,7 +881,7 @@ def _resolve_pkv_names(layer_idx, layer_state): input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, - opset_version=constants.ONNX_EXPORT_OPSET, + opset_version=constants.ONNX_LEGACY_EXPORT_OPSET, dynamo=False, **export_kwargs, ) diff --git a/QEfficient/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py index d49718c33e..eb27051b9d 100644 --- a/QEfficient/base/onnx_transforms.py +++ b/QEfficient/base/onnx_transforms.py @@ -46,8 +46,10 @@ ) # from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func +from QEfficient.customop.onnxscript_utils import get_onnxscript_func from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc -from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, ONNX_EXPORT_OPSET, SIZE_THRESHOLD_DEFAULT +from QEfficient.utils import constants +from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, SIZE_THRESHOLD_DEFAULT logger = logging.getLogger(__name__) @@ -114,13 +116,13 @@ class CustomOpTransform(BaseOnnxTransform): } @classmethod - def apply(cls, model: ModelProto) -> bool: + def apply(cls, model: ModelProto, onnx_export_opset: int = constants.ONNX_LEGACY_EXPORT_OPSET) -> bool: op_applied = False # Register with PyTorch ONNX exporter (for export time) for op_name, (func_class, _) in cls._custom_ops.items(): if hasattr(func_class, "symbolic"): - torch.onnx.register_custom_op_symbolic(f"::{op_name}", func_class.symbolic, ONNX_EXPORT_OPSET) + torch.onnx.register_custom_op_symbolic(f"::{op_name}", func_class.symbolic, onnx_export_opset) used_op_types = {node.op_type for node in model.graph.node} for function_proto in model.functions: @@ -130,7 +132,7 @@ def apply(cls, model: ModelProto) -> bool: existing = {f.name for f in model.functions} for func_name, onnxscript_func in cls._custom_ops.values(): - proto = onnxscript_func.to_function_proto() + proto = get_onnxscript_func(onnxscript_func, onnx_export_opset).to_function_proto() if proto.name not in used_op_types: continue if proto.name not in existing: @@ -643,7 +645,9 @@ def _set_external_data(tensor, file_name): # Non-looping transforms if CustomOpTransform in requested: - applied[CustomOpTransform] = CustomOpTransform.apply(model) + applied[CustomOpTransform] = CustomOpTransform.apply( + model, onnx_export_opset=kwargs.get("onnx_export_opset", constants.ONNX_LEGACY_EXPORT_OPSET) + ) if RenameFunctionOutputsTransform in requested: applied[RenameFunctionOutputsTransform] = RenameFunctionOutputsTransform.apply( diff --git a/QEfficient/customop/ctx_scatter_gather.py b/QEfficient/customop/ctx_scatter_gather.py index aedddb186e..4fceea088f 100644 --- a/QEfficient/customop/ctx_scatter_gather.py +++ b/QEfficient/customop/ctx_scatter_gather.py @@ -8,12 +8,13 @@ import onnxscript import torch +from QEfficient.customop.onnxscript_utils import qeff_custom_op from QEfficient.utils import constants -ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET)) +ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET)) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxScatter(data: onnxscript.FLOAT, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT) -> onnxscript.FLOAT: # Find dims batch_size = ops.Gather(ops.Shape(data), [0]) @@ -56,7 +57,7 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat return g.onnxscript_op(CtxScatter, data, position_ids, updates).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxScatter3D(data: onnxscript.FLOAT, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT) -> onnxscript.FLOAT: # Find dims batch_size = ops.Gather(ops.Shape(data), [0]) @@ -122,7 +123,7 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat return g.onnxscript_op(CtxScatter3D, data, position_ids, updates).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxScatter3DInt( data: onnxscript.INT32, position_ids: onnxscript.INT32, updates: onnxscript.INT32 ) -> onnxscript.INT32: @@ -164,7 +165,7 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat return g.onnxscript_op(CtxScatter3DInt, data, position_ids, updates).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxGather3D(data: onnxscript.FLOAT, ctx_indices: onnxscript.INT32) -> onnxscript.FLOAT: batch_size = ops.Slice(ops.Shape(data), starts=[0], ends=[1], axes=[0]) idx_seq_len = ops.Slice(ops.Shape(ctx_indices), starts=[1], ends=[2], axes=[0]) @@ -215,7 +216,7 @@ def symbolic(g: torch.Graph, data: torch.Value, ctx_indices: torch.Value) -> tor return g.onnxscript_op(CtxGather3D, data, ctx_indices) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxGather( data: onnxscript.FLOAT, ctx_indices: onnxscript.INT32, comp_ctx_len: onnxscript.INT32 ) -> onnxscript.FLOAT: @@ -249,7 +250,7 @@ def symbolic(g: torch.Graph, data: torch.Value, ctx_indices: torch.Value, comp_c return g.onnxscript_op(CtxGather, data, ctx_indices, comp_ctx_len).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxGatherBlockedKV(data: onnxscript.FLOAT, ctx_indices: onnxscript.INT32) -> onnxscript.FLOAT: ctx_indices = ops.Unsqueeze(ctx_indices, [-1]) return ops.GatherND(data, ctx_indices, batch_dims=2) diff --git a/QEfficient/customop/ctx_scatter_gather_cb.py b/QEfficient/customop/ctx_scatter_gather_cb.py index 18e3e3fa2a..f05951175c 100644 --- a/QEfficient/customop/ctx_scatter_gather_cb.py +++ b/QEfficient/customop/ctx_scatter_gather_cb.py @@ -8,12 +8,13 @@ import onnxscript import torch +from QEfficient.customop.onnxscript_utils import qeff_custom_op from QEfficient.utils import constants -ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET)) +ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET)) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxScatterCB( data: onnxscript.FLOAT, batch_index: onnxscript.INT32, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT ) -> onnxscript.FLOAT: @@ -58,7 +59,7 @@ def symbolic( return g.onnxscript_op(CtxScatterCB, data, batch_index, position_ids, updates).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxScatterCB3D( data: onnxscript.FLOAT, batch_index: onnxscript.INT32, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT ) -> onnxscript.FLOAT: @@ -99,7 +100,7 @@ def symbolic( return g.onnxscript_op(CtxScatterCB3D, data, batch_index, position_ids, updates).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxGatherCB( data: onnxscript.FLOAT, batch_index: onnxscript.INT32, ctx_indices: onnxscript.INT32, comp_ctx_len: onnxscript.INT32 ) -> onnxscript.FLOAT: @@ -146,7 +147,7 @@ def symbolic( return g.onnxscript_op(CtxGatherCB, data, batch_index, ctx_indices, comp_ctx_len).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxGatherBlockedKVCB( data: onnxscript.FLOAT, batch_index: onnxscript.INT32, ctx_indices: onnxscript.INT32 ) -> onnxscript.FLOAT: @@ -185,7 +186,7 @@ def symbolic(g: torch.Graph, data: torch.Value, batch_index: torch.Value, ctx_in return g.onnxscript_op(CtxGatherBlockedKVCB, data, batch_index, ctx_indices).setTypeAs(data) -@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1)) +@qeff_custom_op("com.qualcomm.cloud", 1) def CtxGatherCB3D( data: onnxscript.FLOAT, batch_index: onnxscript.INT32, ctx_indices: onnxscript.INT32 ) -> onnxscript.FLOAT: diff --git a/QEfficient/customop/dynamo_ops.py b/QEfficient/customop/dynamo_ops.py index a1f8732e01..e1cbedcbde 100644 --- a/QEfficient/customop/dynamo_ops.py +++ b/QEfficient/customop/dynamo_ops.py @@ -368,21 +368,22 @@ def _(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> CtxScatterCB, CtxScatterCB3D, ) +from QEfficient.customop.onnxscript_utils import get_dynamo_onnxscript_func # noqa: E402 from QEfficient.customop.rms_norm import CustomRMSNorm # noqa: E402 DYNAMO_CUSTOM_OP_TABLE = { - torch.ops.qefficient.rms_norm.default: CustomRMSNorm, - torch.ops.qefficient.ctx_scatter.default: CtxScatter, - torch.ops.qefficient.ctx_scatter_3d.default: CtxScatter3D, - torch.ops.qefficient.ctx_scatter_cb.default: CtxScatterCB, - torch.ops.qefficient.ctx_scatter_cb_3d.default: CtxScatterCB3D, - torch.ops.qefficient.ctx_scatter_3d_int.default: CtxScatter3DInt, - torch.ops.qefficient.ctx_scatter_3d_generalized.default: CtxScatter3D, - torch.ops.qefficient.ctx_gather.default: CtxGather, - torch.ops.qefficient.ctx_gather_3d.default: CtxGather3D, - torch.ops.qefficient.ctx_gather_cb.default: CtxGatherCB, - torch.ops.qefficient.ctx_gather_cb_3d.default: CtxGatherCB3D, - torch.ops.qefficient.ctx_gather_blocked_kv.default: CtxGatherBlockedKV, - torch.ops.qefficient.ctx_gather_blocked_kv_cb.default: CtxGatherBlockedKVCB, - torch.ops.qefficient.ctx_gather_3d_generalized.default: CtxGather3D, + torch.ops.qefficient.rms_norm.default: get_dynamo_onnxscript_func(CustomRMSNorm), + torch.ops.qefficient.ctx_scatter.default: get_dynamo_onnxscript_func(CtxScatter), + torch.ops.qefficient.ctx_scatter_3d.default: get_dynamo_onnxscript_func(CtxScatter3D), + torch.ops.qefficient.ctx_scatter_cb.default: get_dynamo_onnxscript_func(CtxScatterCB), + torch.ops.qefficient.ctx_scatter_cb_3d.default: get_dynamo_onnxscript_func(CtxScatterCB3D), + torch.ops.qefficient.ctx_scatter_3d_int.default: get_dynamo_onnxscript_func(CtxScatter3DInt), + torch.ops.qefficient.ctx_scatter_3d_generalized.default: get_dynamo_onnxscript_func(CtxScatter3D), + torch.ops.qefficient.ctx_gather.default: get_dynamo_onnxscript_func(CtxGather), + torch.ops.qefficient.ctx_gather_3d.default: get_dynamo_onnxscript_func(CtxGather3D), + torch.ops.qefficient.ctx_gather_cb.default: get_dynamo_onnxscript_func(CtxGatherCB), + torch.ops.qefficient.ctx_gather_cb_3d.default: get_dynamo_onnxscript_func(CtxGatherCB3D), + torch.ops.qefficient.ctx_gather_blocked_kv.default: get_dynamo_onnxscript_func(CtxGatherBlockedKV), + torch.ops.qefficient.ctx_gather_blocked_kv_cb.default: get_dynamo_onnxscript_func(CtxGatherBlockedKVCB), + torch.ops.qefficient.ctx_gather_3d_generalized.default: get_dynamo_onnxscript_func(CtxGather3D), } diff --git a/QEfficient/customop/onnxscript_utils.py b/QEfficient/customop/onnxscript_utils.py new file mode 100644 index 0000000000..2103b3af30 --- /dev/null +++ b/QEfficient/customop/onnxscript_utils.py @@ -0,0 +1,72 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import inspect +from collections.abc import Callable + +import onnxscript + +from QEfficient.utils import constants + +_DYNAMO_FUNC_ATTR = "_qeff_dynamo_onnxscript_func" +_MISSING = object() + + +def _onnx_opset(opset_version: int) -> onnxscript.values.Opset: + return getattr(onnxscript, f"opset{opset_version}") + + +def _compile_with_default_opset( + fn: Callable, + custom_opset: onnxscript.values.Opset, + default_opset_version: int, +): + """Compile an ONNXScript function with the requested default-domain opset. + + Existing custom-op definitions refer to a module-level ``ops`` symbol. The + ONNXScript compiler captures module globals while the decorator runs, so we + temporarily point that symbol at the requested default opset and immediately + restore it after compilation. + """ + module = inspect.getmodule(fn) + if module is None: + raise RuntimeError(f"Unable to resolve module for ONNXScript function {fn.__name__}") + + previous_ops = module.__dict__.get("ops", _MISSING) + module.__dict__["ops"] = _onnx_opset(default_opset_version) + try: + return onnxscript.script(custom_opset)(fn) + finally: + if previous_ops is _MISSING: + module.__dict__.pop("ops", None) + else: + module.__dict__["ops"] = previous_ops + + +def qeff_custom_op(domain: str, version: int): + """Compile one custom op body into legacy and dynamo ONNXScript variants.""" + custom_opset = onnxscript.values.Opset(domain, version) + + def decorator(fn: Callable): + legacy_func = _compile_with_default_opset(fn, custom_opset, constants.ONNX_LEGACY_EXPORT_OPSET) + dynamo_func = _compile_with_default_opset(fn, custom_opset, constants.ONNX_DYNAMO_EXPORT_OPSET) + setattr(legacy_func, _DYNAMO_FUNC_ATTR, dynamo_func) + return legacy_func + + return decorator + + +def get_dynamo_onnxscript_func(onnxscript_func): + """Return the dynamo/opset18 variant attached by ``qeff_custom_op``.""" + return getattr(onnxscript_func, _DYNAMO_FUNC_ATTR) + + +def get_onnxscript_func(onnxscript_func, onnx_export_opset: int): + """Return the ONNXScript variant matching the requested export opset.""" + if onnx_export_opset == constants.ONNX_DYNAMO_EXPORT_OPSET: + return get_dynamo_onnxscript_func(onnxscript_func) + return onnxscript_func diff --git a/QEfficient/customop/rms_norm.py b/QEfficient/customop/rms_norm.py index 5d1b5db427..adbb496594 100644 --- a/QEfficient/customop/rms_norm.py +++ b/QEfficient/customop/rms_norm.py @@ -9,13 +9,14 @@ import torch from torch import nn +from QEfficient.customop.onnxscript_utils import qeff_custom_op from QEfficient.customop.utils import select_interface from QEfficient.utils import constants -ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET)) +ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET)) -@onnxscript.script(onnxscript.values.Opset(domain="com.qti.aisw.onnx", version=1)) +@qeff_custom_op("com.qti.aisw.onnx", 1) def CustomRMSNorm(hidden_states: onnxscript.FLOAT, weight: onnxscript.FLOAT, epsilon: float) -> onnxscript.FLOAT: weight = ops.Cast(weight, to=1) variance = ops.ReduceMean(ops.Pow(hidden_states, 2), axes=[-1], keepdims=1) diff --git a/QEfficient/exporter/export_utils.py b/QEfficient/exporter/export_utils.py index 3bc8f1133a..78c4e8c1c6 100644 --- a/QEfficient/exporter/export_utils.py +++ b/QEfficient/exporter/export_utils.py @@ -100,7 +100,7 @@ def export_onnx( input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, - opset_version=constants.ONNX_EXPORT_OPSET, + opset_version=constants.ONNX_LEGACY_EXPORT_OPSET, custom_opsets={"com.qti.aisw.onnx": 1}, ) except Exception as e: diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index e4b42234c6..d931d918a4 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -98,11 +98,17 @@ def get_models_dir(): ONNX_EXPORT_EXAMPLE_MAX_TOP_K_IDS = 512 ONNX_EXPORT_EXAMPLE_TOP_PS = 0.80 ONNX_EXPORT_EXAMPLE_MIN_PS = 0.99 -ONNX_EXPORT_OPSET = 18 +ONNX_LEGACY_EXPORT_OPSET = 17 +ONNX_DYNAMO_EXPORT_OPSET = 18 +ONNX_EXPORT_OPSET = ONNX_LEGACY_EXPORT_OPSET FILE_CHUNK_SIZE_DEFAULT = 10 * 2**30 # 10 GB SIZE_THRESHOLD_DEFAULT = 1024 +def get_onnx_export_opset(dynamo: bool = False) -> int: + return ONNX_DYNAMO_EXPORT_OPSET if dynamo else ONNX_LEGACY_EXPORT_OPSET + + COMPILER = ["/opt/qti-aic/exec/qaic-compile", "-aic-hw"] diff --git a/tests/unit_test/transforms/test_onnx_transforms.py b/tests/unit_test/transforms/test_onnx_transforms.py index 00bccd1dcd..0bb0a9cb3f 100644 --- a/tests/unit_test/transforms/test_onnx_transforms.py +++ b/tests/unit_test/transforms/test_onnx_transforms.py @@ -81,6 +81,24 @@ def test_custom_op_transform_has_apply_method(self): assert hasattr(CustomOpTransform, "apply") assert callable(CustomOpTransform.apply) + def test_custom_op_variants_use_legacy_and_dynamo_opsets(self): + from QEfficient.customop.onnxscript_utils import get_onnxscript_func + from QEfficient.customop.rms_norm import CustomRMSNorm + from QEfficient.utils import constants + + def default_domain_opset(function_proto): + return next(opset.version for opset in function_proto.opset_import if opset.domain == "") + + legacy_proto = get_onnxscript_func( + CustomRMSNorm, constants.get_onnx_export_opset(dynamo=False) + ).to_function_proto() + dynamo_proto = get_onnxscript_func( + CustomRMSNorm, constants.get_onnx_export_opset(dynamo=True) + ).to_function_proto() + + assert default_domain_opset(legacy_proto) == constants.ONNX_LEGACY_EXPORT_OPSET + assert default_domain_opset(dynamo_proto) == constants.ONNX_DYNAMO_EXPORT_OPSET + def test_base_onnx_transform_importable(self): from QEfficient.base.onnx_transforms import BaseOnnxTransform