Replies: 2 comments
-
Note to add to docs that in FX you need classes but in PyTorch you need strings that point to node kinds. |
Beta Was this translation helpful? Give feedback.
0 replies
-
There are two APIs: |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
-
torch_executed_ops
in FXTL;DR
In line with unifying the TorchScript/FX Frontends, the
torch_executed_ops
field from TorchScript should be available for FX use as well. The existing FXleaf_module_list
attribute in the tracer accomplishes a similar result (for modules), however the naming and functionality of the two features should be unified.Goals + Usecases
Specifying that certain operations/modules should be run in Torch, as opposed to the accelerated TensorRT framework, is a common feature among both the FX and TS paths of Torch-TensorRT. However, the method of invoking this feature is different among TS and FX, and could be unified, in line with the ongoing effort to consolidate the frontend interface (RFC #1372, PR #1404). Specifically, while compiling in the TorchScript/FX paths is as easy as toggling
ir="fx"
orir="ts"
intorch_tensorrt.compile(...)
, one cannot do the same fortorch_executed_ops
.Enabling dual TS/FX use of
torch_executed_ops
, alongside other fields currently used exclusively for TorchScript would improve and streamline the existing compilation process.Proposed APIs / UX
Example Workflow
A user would interact with this feature by using the
torch_tensorrt.compile(...)
function with the argumentsir="fx"
and specifying a list of excluded operations, to be executed in Torch (non-accelerated). For instance:Currently, users can exclude modules (like
torch.nn.ReLU
), by setting theleaf_module_list
field of theacc_tracer
, but operations liketorch.add
cannot be excluded this way. The snippet below presents a method to compile a model via the FX path using theacc_tracer
manually.TensorRT/examples/fx/fx2trt_example.py
Lines 23 to 55 in deda87b
Below is a preview of the
leaf_module_list
argument in the tracer.TensorRT/py/torch_tensorrt/fx/tracer/acc_tracer/acc_tracer.py
Lines 274 to 309 in deda87b
Finally, we have the
exclude_support_node_name
argument of theTRTSplitterSetting
:TensorRT/py/torch_tensorrt/fx/tools/trt_splitter.py
Lines 44 to 58 in a343650
Internal Implementation
Design
The design of this feature would begin with TS/FX unification of the
torch_executed_ops
argument. Specifically, this argument should be capable of taking two different types of inputs:ir="ts"
) [Already Supported]torch_executed_ops=["aten::where"]
)ir="fx"
) [To Add]torch_executed_ops=[torch.nn.ReLU, torch.add]
)torch.nn
, such astorch.nn.Softmax
, which has corresponding aten operatoraten::softmax
.Then, for the FX path, the next step would be to add functionality for these operators to be excluded during the tracing/splitting. Specifically, this would include marking certain operations, like
torch.add
to register as unsupported, as though their accelerated counterpart was unimplemented. This would likely involve adding modules toleaf_module_list
and operators toexclude_support_node_name
, and writing functionality to distinguish which modules/operators should go where.Extensions Required to Core API implementations
The existing core Python library would need changes to the
compile
function to support passing arguments fromtorch_executed_ops
to the FX compiler and handling the parsing and proper assignment of those operators to FX.Details specific for TorchScript Support
Implementation exists.
Details specific for FX support
A key challenge to note here is the overload of the terms "module" and "operation". Certain modules, such as
torch.nn.Conv2d
can map to a singleaten::convolution
operator, while the operatortorch.add
maps toaten::add
. In this sense,aten
and the TorchScript path have a clearer notion of a single operation (for example, excluding convolutions and adds), whereas in the FX path, convolution might be considered a "module" while add would be an "operation". Thus, to disabletorch.add
, one would need to employ a different method than addingtorch.add
to theleaf_module_list
, sincetorch.add
is not considered a module.exclude_support_node_name
, as discussed above, is a feasible option for excluding individual operators.Note: Ensure consideration/differentiation of ops in TorchScript (strings like
"aten::add"
), versus Torch objects as needed in FX. Thetorch_executed_ops
field could take a mix of these types.Implementation Phases
Prototype - Extra Small / Small
torch_executed_modules
argument to theleaf_module_list
argument, and allow pass-through of operators between thetorch_tensorrt.compile(ir="fx")
function invocation, and theacc_tracer
invocation.torch_executed_modules
list provided by the user is valid for FX, ifir="fx"
is specified.MVP
1.4.0
- Small / Mediumtorch_executed_ops
for simple operations marked to be executed in Torch, in addition to modulesBeta Was this translation helpful? Give feedback.
All reactions