Added kwarg support for conver_module_to_engine

Author: cehongwang

py/torch_tensorrt/_compile.py

@@ -351,7 +351,7 @@ def convert_method_to_trt_engine(
         torchtrt_inputs = prepare_inputs(inputs)
         exp_program = torch_tensorrt.dynamo.trace(module, torchtrt_inputs, **kwargs)
 
-        return dynamo_convert_module_to_trt_engine(  # type: ignore
+        return dynamo_convert_module_to_trt_engine(
             exp_program,
             inputs=tuple(inputs),
             enabled_precisions=enabled_precisions_set,
@@ -429,9 +429,10 @@ def save(
         raise ValueError(
             "Not all inputs provided are torch.tensors. Please provide torch.tensors as inputs"
         )
-    if kwargs_inputs is not None and not all(
-        value is not None for value in kwargs_inputs.values()
-    ):
+    if kwargs_inputs is None:
+        kwargs_inputs = {}
+
+    if kwargs_inputs and not all(value is not None for value in kwargs_inputs.values()):
         raise ValueError("kwargs should not include None.")
     if output_format not in accepted_formats:
         raise ValueError(

py/torch_tensorrt/dynamo/_compiler.py

@@ -34,6 +34,7 @@
     pre_export_lowering,
 )
 from torch_tensorrt.dynamo.utils import (
+    flatten_dict_value,
     get_torch_inputs,
     parse_complex_tensor_structs,
     prepare_inputs,
@@ -495,6 +496,7 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
 def convert_module_to_trt_engine(
     exported_program: ExportedProgram,
     inputs: Sequence[Any],
+    kwarg_inputs: Optional[dict[str, Any]] = None,
     *,
     enabled_precisions: (
         Set[torch.dtype | dtype] | Tuple[torch.dtype | dtype]
@@ -608,12 +610,15 @@ def convert_module_to_trt_engine(
             stacklevel=2,
         )
 
-    input_list = list(inputs) if inputs is not None else []
+    arg_input_list = list(inputs) if inputs is not None else []
     torch_executed_ops = torch_executed_ops if torch_executed_ops is not None else set()
+    if kwarg_inputs is None:
+        kwarg_inputs = {}
     # Prepare torch_trt inputs
-    input_list = prepare_inputs(input_list)
+    arg_input_list = prepare_inputs(arg_input_list)
+    kwarg_input_list = prepare_inputs(kwarg_inputs)
+    flattened_input_list = arg_input_list + flatten_dict_value(kwarg_input_list)
     device = to_torch_tensorrt_device(device)
-    torch_inputs = get_torch_inputs(input_list, device)
     enabled_precisions = {dtype._from(e) for e in enabled_precisions}
 
     compilation_options = {
@@ -662,8 +667,15 @@ def convert_module_to_trt_engine(
     # Assume converters support dynamic shapes and disable validation
     CONVERTERS.set_dynamic_shape_support(settings.assume_dynamic_shape_support)
 
+    interpreter_result = interpret_module_to_result(
+        gm,
+        inputs=flattened_input_list,
+        arg_inputs=arg_input_list,
+        kwarg_inputs=kwarg_input_list,
+        settings=settings,
+    )
     try:
-        interpreter_result = interpret_module_to_result(gm, input_list, settings)
+        pass
     except UnsupportedOperatorException:
         logger.error(
             f"Conversion of module {gm} not currently fully supported or convertible!",

py/torch_tensorrt/dynamo/_exporter.py

@@ -30,6 +30,8 @@ def export(
         gm (torch.fx.GraphModule): Compiled Torch-TensorRT module, generated by ``torch_tensorrt.dynamo.compile``
         inputs (torch.Tensor): Torch input tensors
     """
+    if kwargs_inputs is None:
+        kwargs_inputs = {}
     patched_module = transform(gm, inputs, kwargs_inputs)
     exp_program = create_trt_exp_program(patched_module)
     return exp_program
@@ -53,8 +55,9 @@ def transform(
     """
     # Make a copy the graph since this function transforms the input graph and changes it's attributes.
     # This transformed graph is meant to be consumed by `create_trt_exp_program`
+    if kwargs_inputs is None:
+        kwargs_inputs = {}
     gm = copy.deepcopy(gm)
-
     # Run shape analysis
     _, outputs_map = partitioning.run_shape_analysis(gm, inputs, kwargs_inputs)
 

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -2,7 +2,7 @@
 
 import io
 import logging
-from typing import List, Sequence
+from typing import Any, List, Optional, Sequence
 
 import torch
 from torch.fx.experimental.proxy_tensor import maybe_disable_fake_tensor_mode
@@ -27,12 +27,16 @@ def infer_module_output_dtypes(
     module: torch.fx.GraphModule,
     inputs: Sequence[Input],
     device: Device,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
     truncate_double: bool = False,
 ) -> List[dtype]:
     with maybe_disable_fake_tensor_mode():
         torch_inputs = get_torch_inputs(inputs, device)
+        if kwarg_inputs is None:
+            kwarg_inputs = {}
+        torch_kwarg_inputs = get_torch_inputs(kwarg_inputs, device)
         module = module.to(device.to(torch.device))
-        module_outputs = module(*torch_inputs)
+        module_outputs = module(*torch_inputs, **torch_kwarg_inputs)
         if not isinstance(module_outputs, (list, tuple)):
             module_outputs = [module_outputs]
 
@@ -62,6 +66,8 @@ def interpret_module_to_result(
     module: torch.fx.GraphModule,
     inputs: Sequence[Input],
     settings: CompilationSettings = CompilationSettings(),
+    arg_inputs: Optional[Sequence[Input]] = None,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
 ) -> TRTInterpreterResult:
     """Interpret an FX module to a TRTInterpreterResult
     Args:
@@ -71,12 +77,22 @@ def interpret_module_to_result(
     Returns:
         TRTInterpreterResult
     """
-    output_dtypes = infer_module_output_dtypes(
-        module,
-        inputs,
-        settings.device,
-        truncate_double=settings.truncate_double,
-    )
+    if arg_inputs is not None:
+        output_dtypes = infer_module_output_dtypes(
+            module,
+            arg_inputs,
+            settings.device,
+            kwarg_inputs=kwarg_inputs,
+            truncate_double=settings.truncate_double,
+        )
+    else:
+        # args and kwargs are combined and flattened to one list
+        output_dtypes = infer_module_output_dtypes(
+            module,
+            inputs,
+            settings.device,
+            truncate_double=settings.truncate_double,
+        )
 
     interpreter = TRTInterpreter(
         module,

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -147,6 +147,8 @@ def get_submodule_io(
         sub_outputs = outputs
         return
 
+    if kwargs_inputs is None:
+        kwargs_inputs = {}
     # Iterate through submodules (both Torch and TRT) and store IO shapes
     for name, _ in parent_module.named_children():
         submodule = getattr(parent_module, name)

tests/py/dynamo/models/test_models_export_kwargs.py

@@ -1,4 +1,6 @@
 # type: ignore
+import os
+import tempfile
 import unittest
 
 import pytest
@@ -62,12 +64,15 @@ def forward(self, x, b=5, c=None, d=None):
     # trt_mod = torchtrt.compile(model, **compile_spec)
 
     exp_program = torch.export.export(model, args=tuple(args), kwargs=kwargs)
-    trt_mod = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    cos_sim = cosine_similarity(model(*args, **kwargs), trt_mod(*args, **kwargs)[0])
+    trt_gm = torchtrt.dynamo.compile(exp_program, **compile_spec)
+    cos_sim = cosine_similarity(model(*args, **kwargs), trt_gm(*args, **kwargs)[0])
     assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
         msg=f"CustomKwargs Module TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
+    # Save the module
+    trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
+    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwargs_inputs=kwargs)
     # Clean up model env
     torch._dynamo.reset()