Added kwarg to refit

Author: cehongwang

docs/_downloads/7e3a125a2d4ba8274a41b46f5e0723fa/refit_engine_example.py

@@ -78,7 +78,7 @@
 new_trt_gm = refit_module_weights(
     compiled_module=compiled_trt_ep,
     new_weight_module=exp_program2,
-    inputs=inputs,
+    arg_inputs=inputs,
 )
 
 # Check the output

examples/dynamo/refit_engine_example.py

@@ -78,7 +78,7 @@
 new_trt_gm = refit_module_weights(
     compiled_module=compiled_trt_ep,
     new_weight_module=exp_program2,
-    inputs=inputs,
+    arg_inputs=inputs,
 )
 
 # Check the output

py/torch_tensorrt/dynamo/_refit.py

@@ -3,7 +3,7 @@
 import collections.abc
 import copy
 import logging
-from typing import Any, Sequence, Tuple
+from typing import Any, Optional, Sequence, Tuple
 
 import numpy as np
 import tensorrt as trt
@@ -36,7 +36,6 @@
 from torch_tensorrt.dynamo.utils import (
     check_output,
     get_torch_inputs,
-    prepare_inputs,
     set_log_level,
     to_torch_device,
     to_torch_tensorrt_device,
@@ -146,7 +145,8 @@ def _refit_single_trt_engine_with_gm(
 def refit_module_weights(
     compiled_module: torch.fx.GraphModule | ExportedProgram,
     new_weight_module: ExportedProgram,
-    inputs: Tuple[Any, ...],
+    arg_inputs: Optional[Tuple[Any, ...]] = None,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
     verify_output: bool = False,
 ) -> torch.fx.GraphModule:
     """
@@ -208,27 +208,29 @@ def refit_module_weights(
     if settings.debug:
         set_log_level(logger.parent, logging.DEBUG)
 
-    if not isinstance(inputs, collections.abc.Sequence):
-        inputs = [inputs]
-
-    # Prepare torch_trt inputs
-    inputs = prepare_inputs(inputs)
     device = to_torch_tensorrt_device(settings.device)
-    torch_inputs = get_torch_inputs(inputs, device)
+    if arg_inputs:
+        if not isinstance(arg_inputs, collections.abc.Sequence):
+            # Prepare torch_trt inputs
+            arg_inputs = [arg_inputs]
+        torch_inputs = get_torch_inputs(arg_inputs, device)
+
+    if kwarg_inputs:
+        torch_kwarg_inputs = get_torch_inputs(kwarg_inputs, device)
     runtime = trt.Runtime(TRT_LOGGER)
     if not isinstance(new_weight_module, ExportedProgram):
         raise AssertionError(
             f"Input graph should be an ExportedProgram but got type {type(new_weight_module)}"
         )
-    new_weight_module = pre_export_lowering(new_weight_module, torch_inputs)
+    new_weight_module = pre_export_lowering(new_weight_module)
     new_weight_module = new_weight_module.run_decompositions(
         get_decompositions(settings.enable_experimental_decompositions)
     )
     new_gm = new_weight_module.module()
     logger.debug("Input graph: " + str(new_gm.graph))
     # Apply lowering on the graph module
 
-    new_gm = post_lowering(new_gm, torch_inputs)
+    new_gm = post_lowering(new_gm)
 
     logger.info("Compilation Settings: %s\n", settings)
 
@@ -354,11 +356,12 @@ def refit_module_weights(
             refitted_engine = torch.classes.tensorrt.Engine(tuple(new_engine_info))
             setattr(compiled_module, f"{name}_engine", refitted_engine)
 
-    if verify_output:
+    if verify_output and arg_inputs is not None:
         if check_output(
             new_module=new_gm,
             refitted_module=compiled_module,
-            inputs=torch_inputs,
+            arg_inputs=torch_inputs,
+            kwarg_inputs=torch_kwarg_inputs,
         ):
             logger.info("Refitting Succeed!")
         else:

py/torch_tensorrt/dynamo/utils.py

@@ -397,9 +397,12 @@ def function_wrapper(*args: Any, **kwargs: Any) -> Any:
 def check_output(
     new_module: torch.fx.GraphModule,
     refitted_module: torch.fx.GraphModule,
-    inputs: tuple[Any, ...],
+    arg_inputs: Any,
+    kwarg_inputs: Any = None,
 ) -> bool:
-    old_outputs, new_outputs = refitted_module(*inputs), new_module(*inputs)
+    old_outputs, new_outputs = refitted_module(*arg_inputs), new_module(
+        *arg_inputs, **kwarg_inputs
+    )
     for old_output, new_output in zip(old_outputs, new_outputs):
         if isinstance(old_output, torch.Tensor) and isinstance(
             new_outputs, torch.Tensor

tests/py/dynamo/models/test_dtype_support.py

@@ -36,7 +36,7 @@ def forward(self, x):
         exp_mod = torch.export.export(mod, (in_tensor,))
         trt_mod = torch_tensorrt.dynamo.compile(
             exp_mod,
-            arg_inputs=[in_tensor],
+            inputs=[in_tensor],
             pass_through_build_failures=True,
             truncate_double=True,
             min_block_size=1,
@@ -74,7 +74,7 @@ def forward(self, x):
         exp_mod = torch.export.export(mod, (in_tensor,))
         trt_mod = torch_tensorrt.dynamo.compile(
             exp_mod,
-            arg_inputs=[in_tensor],
+            inputs=[in_tensor],
             pass_through_build_failures=True,
             truncate_double=True,
             min_block_size=1,
@@ -118,7 +118,7 @@ def forward(self, x):
         exp_mod = torch.export.export(mod, (in_tensor,))
         trt_mod = torch_tensorrt.dynamo.compile(
             exp_mod,
-            arg_inputs=[in_tensor],
+            inputs=[in_tensor],
             pass_through_build_failures=True,
             truncate_double=False,
             min_block_size=1,
@@ -157,7 +157,7 @@ def forward(self, x):
         exp_mod = torch.export.export(mod, (in_tensor,))
         trt_mod = torch_tensorrt.dynamo.compile(
             exp_mod,
-            arg_inputs=[in_tensor],
+            inputs=[in_tensor],
             pass_through_build_failures=True,
             truncate_double=False,
             min_block_size=1,
@@ -201,7 +201,7 @@ def forward(self, x):
         exp_mod = torch.export.export(mod, (in_tensor,))
         trt_mod = torch_tensorrt.dynamo.compile(
             exp_mod,
-            arg_inputs=[in_tensor],
+            inputs=[in_tensor],
             pass_through_build_failures=True,
             enabled_precisions={torch.float, torch.bfloat16, torch.half},
             min_block_size=1,
@@ -239,7 +239,7 @@ def forward(self, x):
         exp_mod = torch.export.export(mod, (in_tensor,))
         trt_mod = torch_tensorrt.dynamo.compile(
             exp_mod,
-            arg_inputs=[in_tensor],
+            inputs=[in_tensor],
             pass_through_build_failures=True,
             enabled_precisions={torch.float, torch.bfloat16, torch.half},
             min_block_size=1,

tests/py/dynamo/models/test_model_refit.py

@@ -108,7 +108,7 @@ def test_refit_one_engine():
     new_trt_gm = refit_module_weights(
         compiled_module=trt_gm,
         new_weight_module=exp_program2,
-        inputs=inputs,
+        arg_inputs=inputs,
     )
 
     # Check the output
@@ -154,7 +154,7 @@ def test_refit_one_engine_bert():
     new_trt_gm = refit_module_weights(
         compiled_module=trt_gm,
         new_weight_module=exp_program2,
-        inputs=inputs,
+        arg_inputs=inputs,
     )
 
     # Check the output
@@ -203,7 +203,7 @@ def test_refit_one_engine_inline_runtime():
     new_trt_gm = refit_module_weights(
         compiled_module=trt_gm,
         new_weight_module=exp_program2,
-        inputs=inputs,
+        arg_inputs=inputs,
     )
 
     # Check the output
@@ -247,7 +247,7 @@ def test_refit_one_engine_python_runtime():
     new_trt_gm = refit_module_weights(
         compiled_module=trt_gm,
         new_weight_module=exp_program2,
-        inputs=inputs,
+        arg_inputs=inputs,
     )
 
     # Check the output
@@ -313,7 +313,7 @@ def forward(self, x):
     new_trt_gm = refit_module_weights(
         compiled_module=trt_gm,
         new_weight_module=exp_program2,
-        inputs=inputs,
+        arg_inputs=inputs,
     )
 
     # Check the output

tests/py/dynamo/models/test_models_export.py

@@ -222,7 +222,7 @@ def calibrate_loop(model):
             exp_program = torch.export.export(model, (input_tensor,))
             trt_model = torchtrt.dynamo.compile(
                 exp_program,
-                arg_inputs=[input_tensor],
+                inputs=[input_tensor],
                 enabled_precisions={torch.float8_e4m3fn},
                 min_block_size=1,
                 debug=True,