Add optimization to dynamo-based exporter #1541
Conversation
|
LGTM, cc @jambayk for review |
| @@ -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( | |||
There was a problem hiding this comment.
This optimization includes onnxruntime contrib ops as well or just standard onnx?
There was a problem hiding this comment.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.)
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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.
Describe your changes
Checklist before requesting a review
lintrunner -a(Optional) Issue link