feat: dynamic shape support for pow/mod/eq operator #2982
Conversation
| @@ -2278,6 +2279,7 @@ def aten_ops_bitwise_not( | |||
|
|
|||
| @dynamo_tensorrt_converter(torch.ops.aten.eq.Tensor) | |||
| @dynamo_tensorrt_converter(torch.ops.aten.eq.Scalar) | |||
| @dynamo_tensorrt_converter(operator.eq, supports_dynamic_shapes=True) | |||
There was a problem hiding this comment.
| @@ -2002,6 +2002,7 @@ def aten_ops_div( | |||
| @dynamo_tensorrt_converter( | |||
| torch.ops.aten.pow.Tensor_Scalar, supports_dynamic_shapes=True | |||
| ) | |||
| @dynamo_tensorrt_converter(operator.pow, supports_dynamic_shapes=True) | |||
There was a problem hiding this comment.
From the PR (#2918), it looks like dynamic testing and registration for operator.mul were added because PyTorch internally uses operator.mul when using torch.nn.Linear layers. This might be why operator.mul was registered with the converter.
If I'm misunderstanding, could you please explain? (cc. @peri044 )
Then, in this PR, why do we need to register operator.pow, operator.eq, and operator.mod with the converter?
There was a problem hiding this comment.
It seems operator.* overrides original Python ops, which results in running any Python op will create a TRT Layer to handle it. However, they are not listed in the schema. I'm also curious if these ops' converters are needed
There was a problem hiding this comment.
Thanks @chohk88 / @zewenli98
If operator.pow is missing and ** is used in module, there is exception from TRTInterpreter.
E torch_tensorrt.dynamo.conversion._TRTInterpreter.UnsupportedOperatorException: Conversion of function _operator.pow not currently supported!
I saw such usage of operator in openchat model. I think we need to register these operators to support such usage.
https://github.com/imoneoi/openchat/blob/master/ochat/models/unpadded_llama.py#L105
There was a problem hiding this comment.
Thanks for the example. If then, it seems to be necessary I think.
cc: @narendasan @dheerajperi
There was a problem hiding this comment.
Yes, openchat model uses operator.pow and there's one other model too I think. We need them to be registered as converters similar to other operator.* variants we have.
There was a problem hiding this comment.
Thanks for the example. LGTM!
| @@ -2002,6 +2002,7 @@ def aten_ops_div( | |||
| @dynamo_tensorrt_converter( | |||
| torch.ops.aten.pow.Tensor_Scalar, supports_dynamic_shapes=True | |||
| ) | |||
| @dynamo_tensorrt_converter(operator.pow, supports_dynamic_shapes=True) | |||
There was a problem hiding this comment.
Yes, openchat model uses operator.pow and there's one other model too I think. We need them to be registered as converters similar to other operator.* variants we have.
5ee1e50 to
bf65839
Compare
This reverts commit 93bb25e.
This reverts commit 93bb25e.
Description
dynamic shape support for pow/mod/eq operator
Fixes # (issue)
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: