Mutable module improvement #3394

cehongwang · 2025-02-13T00:22:56Z

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

tests/py/dynamo/runtime/test_mutable_torchtrt_module.py

examples/dynamo/mutable_torchtrt_module_example.py

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

tests/py/dynamo/runtime/test_mutable_torchtrt_module.py

examples/dynamo/mutable_torchtrt_module_example.py

narendasan

Do you have the persistent cache example?

cehongwang · 2025-02-26T01:37:51Z

Do you have the persistent cache example?

I talked to Boris and it seems like save and load is what he is looking for. I am adding engine caching example to MutableTorchTensorRTModule as well.

peri044

Mostly LGTM. Added minor comments

peri044 · 2025-02-27T07:21:02Z

examples/dynamo/mutable_torchtrt_module_example.py

@@ -63,16 +65,14 @@
 # Saving Mutable Torch TensorRT Module
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

-# Currently, saving is only enabled for C++ runtime, not python runtime.
+# Currently, saving is only when "use_python" = False in settings


do you mean use_python_runtime=False ?

peri044 · 2025-02-27T07:22:21Z

examples/dynamo/mutable_torchtrt_module_example.py

    pipe.to(device)

    # The only extra line you need
    pipe.unet = torch_trt.MutableTorchTensorRTModule(pipe.unet, **settings)
-
-    image = pipe(prompt, negative_prompt=negative, num_inference_steps=30).images[0]
+    BATCH = torch.export.Dim("BATCH", min=1 * 2, max=12 * 2)


any reason why it is written as 1*2 and 12*2 ? instead of 2 and 24 ?

peri044 · 2025-02-27T07:28:55Z

examples/dynamo/mutable_torchtrt_module_example.py

+enabled_precisions = {torch.float}
+debug = False
+min_block_size = 1
+use_python_runtime = True


is this necessary ?

peri044 · 2025-02-27T07:35:10Z

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

+        kwargs_dynamic_shape: dict[str, Any],
+    ) -> None:
+        """
+        Set the dynamic shape range. The shape hint should EXACTLY follow arg_inputs and kwarg_inputs passed to the forward function


Can you add this link to reference torch.export's convention : https://pytorch.org/docs/stable/export.html#expressing-dynamism ?

peri044 · 2025-02-27T07:37:11Z

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

+        """
+        assert isinstance(
+            args_dynamic_shape, tuple
+        ), "args dynamic shape has to be a tuple"


Can you add - but the provided type is {type(args_dynamic_shape)} ?

peri044 · 2025-02-27T07:37:21Z

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

+        ), "args dynamic shape has to be a tuple"
+        assert isinstance(
+            kwargs_dynamic_shape, dict
+        ), "args dynamic shape has to be a dictionary"


Can you add - but the provided type is {type(kwargs_dynamic_shape)} ?

peri044 · 2025-02-27T08:00:59Z

tests/py/dynamo/runtime/test_mutable_torchtrt_module.py

+    dynamic_shape = {"a": {1: dim}, "b": [{}, {}], "c": {"a": {}, "b": [{}, {}]}}
+    assertions.assertFalse(
+        torch_trt.MutableTorchTensorRTModule._check_inputs_shape(a, b),
+        msg=f"test_check_output_equal is not correct.",


test_check_input_shape_dynamic

peri044 · 2025-02-27T08:01:04Z

tests/py/dynamo/runtime/test_mutable_torchtrt_module.py

+    )
+    assertions.assertTrue(
+        torch_trt.MutableTorchTensorRTModule._check_inputs_shape(a, b, dynamic_shape),
+        msg=f"test_check_output_equal is not correct.",


test_check_input_shape_dynamic

cehongwang requested review from narendasan and peri044 February 13, 2025 00:22

cehongwang self-assigned this Feb 13, 2025

facebook-github-bot added the cla signed label Feb 13, 2025

github-actions bot requested a review from gs-olive February 13, 2025 00:23

cehongwang force-pushed the mutable_module_improvement branch 2 times, most recently from 210bd3a to 02011eb Compare February 13, 2025 00:30

github-actions bot removed documentation Improvements or additions to documentation component: lowering Issues re: The lowering / preprocessing passes labels Feb 13, 2025

cehongwang force-pushed the mutable_module_improvement branch 3 times, most recently from f183050 to 9580f92 Compare February 13, 2025 21:22

Chengzhe Xu and others added 11 commits February 15, 2025 00:39

init commit for flux torch.compile

248f957

chore: updates

1a81654

chore: updates

cc98a45

chore: updates

46e5170

chore: updates

51ca10f

chore: updates

3551f4b

chore: updates

245c8f2

chore: minor fix

5a2d29b

chore: updates

8435706

Added dynamic shape support to MutableTorchTensorRTModule

6fe31f8

Modified the LoRA example of mutable torch trt module

ec2d674

cehongwang force-pushed the mutable_module_improvement branch from 3db0aec to ec2d674 Compare February 15, 2025 00:39

cehongwang removed the request for review from gs-olive February 18, 2025 18:53

narendasan reviewed Feb 18, 2025

View reviewed changes

cehongwang force-pushed the mutable_module_improvement branch from 0ca42f3 to 0643d96 Compare February 18, 2025 22:40

peri044 reviewed Feb 18, 2025

View reviewed changes

cehongwang force-pushed the mutable_module_improvement branch 6 times, most recently from a5ea97b to a59d92d Compare February 20, 2025 23:04

Fixed the issue in comments

a8e0b48

cehongwang force-pushed the mutable_module_improvement branch from a59d92d to a8e0b48 Compare February 21, 2025 18:49

github-actions bot added the component: conversion Issues re: Conversion stage label Feb 21, 2025

narendasan reviewed Feb 25, 2025

View reviewed changes

examples/dynamo/mutable_torchtrt_module_example.py Show resolved Hide resolved

narendasan reviewed Feb 25, 2025

View reviewed changes

Added dynamic shape support to SDXL example

1a309b8

cehongwang force-pushed the mutable_module_improvement branch from cef1620 to 1a309b8 Compare February 26, 2025 00:22

Fixed the comment and added engine cache example

76437f9

cehongwang force-pushed the mutable_module_improvement branch from 7cf4ad9 to 76437f9 Compare February 26, 2025 02:01

Added support for bool input

3375f49

peri044 reviewed Feb 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Mutable module improvement #3394

Mutable module improvement #3394

cehongwang commented Feb 13, 2025

narendasan left a comment

cehongwang commented Feb 26, 2025

peri044 left a comment

peri044 Feb 27, 2025

peri044 Feb 27, 2025

peri044 Feb 27, 2025

peri044 Feb 27, 2025

peri044 Feb 27, 2025

peri044 Feb 27, 2025

peri044 Feb 27, 2025

peri044 Feb 27, 2025

Mutable module improvement #3394

Are you sure you want to change the base?

Mutable module improvement #3394

Conversation

cehongwang commented Feb 13, 2025

Description

Type of change

Checklist:

narendasan left a comment

Choose a reason for hiding this comment

cehongwang commented Feb 26, 2025

peri044 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment