[5750013][5591945][5360813]: AutoCast standalone implementation for type inference #719
Conversation
Signed-off-by: Gal Hubara Agam <[email protected]>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #719 +/- ##
==========================================
- Coverage 74.69% 74.62% -0.08%
==========================================
Files 192 192
Lines 18946 19169 +223
==========================================
+ Hits 14152 14305 +153
- Misses 4794 4864 +70 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
created by cursor and committed by mistake Signed-off-by: Gal Hubara Agam <[email protected]>
Signed-off-by: Gal Hubara Agam <[email protected]>
Signed-off-by: Gal Hubara Agam <[email protected]>
Signed-off-by: Gal Hubara Agam <[email protected]>
galagam
changed the title
…add to changelog Signed-off-by: Gal Hubara Agam <[email protected]>
galagam
force-pushed
the
feat/local-type-infer
branch
from
a659cad to
7caedc7
Compare
galagam
changed the title
|
This is not related to this PR, but a shape inference issue I encountered previously. It was happening due to using strict mode in shape inference. Would it be possible to not use the strict mode and use the default mode itself? |
| if use_standalone_type_inference: | ||
| model = onnx_utils.infer_types(model) | ||
| else: | ||
| model = onnx_utils.infer_shapes(model) |
There was a problem hiding this comment.
Could you create a util function for this as it is reused in multiple places?
| if not self.use_standalone_type_inference: | ||
| for idx, d in enumerate(inp.type.tensor_type.shape.dim): | ||
| if d.dim_value: | ||
| inp.type.tensor_type.shape.dim[idx].dim_param = "unk" |
There was a problem hiding this comment.
Similarly for this, can we create a util function?
| return vi.type.tensor_type.elem_type | ||
| return None | ||
|
|
||
| def str_to_tensor_dtype(dtype_str: str) -> onnx.TensorProto.DataType: |
There was a problem hiding this comment.
We have a smaller subset of this mapping here:
Would it be possible to update/reuse this?
| node_name_to_onnx = {node.name: node for node in graph.node} | ||
|
|
||
| # Get nodes to process (from graphsurgeon if available, otherwise from graph directly) | ||
| if gs_graph is not None: |
There was a problem hiding this comment.
Should we convert the graph back to an onnx graph to avoid this logic?
| if not inp_name: | ||
| continue |
There was a problem hiding this comment.
Do we expect any inputs with empty names?
| for attr in node.attribute: | ||
| if attr.name == "value" and attr.type == onnx.AttributeProto.TENSOR: | ||
| if attr.t.HasField("data_type"): | ||
| const_type = attr.t.data_type |
There was a problem hiding this comment.
This pattern is used in multiple places in our codebase, should we create a utility function for it?
| tensor_types[init_name] = init.data_type | ||
|
|
||
| # Helper function to get tensor type | ||
| def get_tensor_type(tensor_name: str) -> int | None: |
There was a problem hiding this comment.
Can we re-use _get_tensor_type() in modelopt.onnx.utils.py or, if not, move this function to modelopt.onnx.utils.py as _get_tensor_type_from_tensor_name or a variation of that?
| return vi.type.tensor_type.elem_type | ||
| return None | ||
|
|
||
| def str_to_tensor_dtype(dtype_str: str) -> onnx.TensorProto.DataType: |
There was a problem hiding this comment.
This can be replaced by onnx_type_str_to_enum() in modelopt.onnx.utils.
| break | ||
| assert const_type is not None | ||
| output_types = [const_type] | ||
| elif node.op_type == "ConstantOfShape": |
There was a problem hiding this comment.
Would this PR fix the issue in bug 5763424?
3 participants
What does this PR do?
Type of change: New feature
Overview:
AutoCast runs full type inference to get the new types after adding casts. ONNX doesn't have a separate function for type inference, and it is done as part of shape inference. Shape inference is a much more complex task than type inference, especially when dynamic shapes are involved. We're seeing some shape inference related bugs in AutoCast. Typically we can WAR, but it's cumbersome. A local implementation might allow users to WAR shape inference related issues. This is opt-in and marked as experimental.
Usage
python -m modelopt.onnx.autocast --onnx_path /path/to/input.onnx [options] --use_standalone_type_inference
Testing
Added use_standalone_type_inference=True to all existing PrecisionConverter tests.
Before your PR is "Ready for review"
Additional Information
A more permanent fix would be to decouple type and shape inference in ONNX, we should invest in that when we have the resources - see onnx/onnx#7100
. This is a quick fix, which is also why it is opt-in and not the default mode.