Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add optimization to dynamo-based exporter #1541

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add optimization to dynamo-based exporter #1541

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions olive/passes/onnx/conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,14 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassCon
"dynamic": PassConfigParam(
type_=bool, default_value=True, description=("Whether to export the model with dynamic axes/shapes.")
),
"optimization": PassConfigParam(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This optimization includes onnxruntime contrib ops as well or just standard onnx?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's this one, so no. https://github.com/microsoft/onnxscript/blob/5e7b0e4b626023d9b726b3c8b8336698f1ac4537/onnxscript/optimizer/_optimizer.py#L29

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @gramalingam

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We probably will need to support both? Should we have two flags? Or make this an enumeration instead of bool, so we can decide whether to apply ORT specific optimizations too? Depends on what the intended use of this Olive pass is ... is it intended only for ORT users?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The current design in PyTorch let users to decide whether to optimize or not. This flag simply follows the idea. I personally think adding another optimization flag for ORT is confusing and complicated. If we really want ORT optimization to be an option, I suggest we can change PyTorch exporter to default applying regular optimization, and optimize() method refers to applying ORT specific optimization.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Having optimize to target ort in pytorch is not the most ideal. We could expose a target option in the optimize API. However depending on what Olive wants to provide to its users, I think a similar string option is reasonable. It could default to ORT, similar to what model builder does right now. Although there are pushes for MB to be more generic.

Copy link

@gramalingam gramalingam Jan 8, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that [responding to Titai's message above] is reasonable: specifically, let's forget pytorch here, and focus on olive users. The olive pass can internally always do generic optimization, no need to expose it to user as an option at all. We could expose ort-optimization as an option ... but even that is necessary only if we expect users to use olive for ORT and non-ORT. Unclear if we need that. That's why I am asking about users of this olive pass ... can we assume it is going to be only for ORT users? (We can't make that assumption for pytorch exporter.)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think even for ORT users, contrib operators are not supported for all target use cases. For instance, NPU targets which require QDQ models might not support any contrib operators.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you use as user input to make the decisions? Some combination of (EP, dtype)? As long as we have the necessary information to select the appropriate fusion optimizations, we should be fine

Copy link
Contributor

@jambayk jambayk Jan 8, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In passes like OrtTransformersOptimization, we make some decisions on what optimizations to enable/disable based on the accelerator spec which includes the target device (cpu, gpu, npu) and EP. dtype is already a config option for this conversion pass so that info is also available if needed.

type_=bool,
default_value=True,
description=(
"Whether to optimize the model with constant folding and "
"elimination of redundant operators after conversion."
),
),
}

def _run_for_config(
Expand Down Expand Up @@ -266,6 +274,8 @@ def _export_pytorch_model(
report=logger.isEnabledFor(logging.DEBUG),
)
assert onnx_program is not None
if config["optimization"]:
onnx_program.optimize()
onnx_model = onnx_program.model_proto
else:
# there might be multiple files created during export, so we need to track the dir
Expand Down
Loading