final revision

Author: cehongwang

.github/workflows/build-test-windows.yml

@@ -135,7 +135,7 @@ jobs:
         pushd .
         cd tests/py/dynamo
         ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_runtime_tests_results.xml runtime/test_000_*
-        ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_partitioning_tests_results.xml partitioning/
+        ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_partitioning_tests_results.xml partitioning/test_000_*
         ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_lowering_tests_results.xml lowering/
         popd
 
@@ -219,6 +219,7 @@ jobs:
         pushd .
         cd tests/py/dynamo
         ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_tests_results.xml runtime/test_001_*
+        ../../../packaging/vc_env_helper.bat python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l1_dynamo_core_partitioning_tests_results.xml partitioning/test_001_*
         popd
 
   L1-dynamo-compile-tests:

py/torch_tensorrt/dynamo/partitioning/_resource_partitioner.py

@@ -59,7 +59,7 @@
 
 logger = logging.getLogger(__name__)
 
-MAX_NUM_OF_ENGINES = 40
+MAX_NUM_OF_ENGINES = 50
 ENGINE_COMPILATION_MEMORY_USAGE_MULTIPLIER = 4
 
 
@@ -127,6 +127,9 @@ def partition_graph(self) -> torch.fx.GraphModule:
             subgraphs, subgraph_size_budget=self.calculate_size_budget()
         )
 
+        if len(subgraphs) == 1:
+            return self.module
+
         # Set the number of TRT engines to be generated
         self.num_trt_accelerated_subgraphs = len([s for s in subgraphs if s.is_acc])
 
@@ -315,7 +318,7 @@ def break_subgraph_by_size(
             )
             new_subgraphs = self.step_and_validate(new_subgraphs, step_size)
             size_0, size_1 = self.size_of_subgraphs(new_subgraphs)
-            if size_0 > size_to_break:
+            if size_0 > size_to_break or size_0 > size_1:
                 break
 
         return new_subgraphs, size_0, size_1

tests/py/dynamo/partitioning/test_001_resource_partitioning.py

@@ -166,7 +166,7 @@ def forward(self, x):
         )
 
         partitioned_module = resource_partition(
-            partitioned_module, cpu_memory_budget=1.7 * 1024 * 1024 * 1024  # 1.7GB,
+            partitioned_module, cpu_memory_budget=1.4 * 1024 * 1024 * 1024  # 1.4GB,
         )
 
         assert (
@@ -177,8 +177,8 @@ def forward(self, x):
                     if "_run_on_acc" in name
                 ]
             )
-            > 3
-        ), "The graph should have more than 3 accelerated subgraphs"
+            == 5
+        ), "The graph should have 5 accelerated subgraphs"
         assert (
             len(
                 [
@@ -298,7 +298,7 @@ def forward(self, x):
         )
 
         partitioned_module = resource_partition(
-            partitioned_module, cpu_memory_budget=1.7 * 1024 * 1024 * 1024  # 1.7GB,
+            partitioned_module, cpu_memory_budget=1.4 * 1024 * 1024 * 1024  # 1.4GB,
         )
 
         assert (
@@ -396,7 +396,7 @@ def forward(self, x):
         )
 
         partitioned_module = resource_partition(
-            partitioned_module, cpu_memory_budget=1.7 * 1024 * 1024 * 1024  # 1.7GB,
+            partitioned_module, cpu_memory_budget=1.4 * 1024 * 1024 * 1024  # 1.4GB,
         )
 
         assert (