Use torch.compile in heuristic search for max batch size

Author: jkosek

CHANGELOG.md

@@ -21,6 +21,7 @@ limitations under the License.
   enables dynamic shapes user setup for TorchTensorRT compilation
 - new: autocast_dtype added Torch runner configuration to set the dtype for autocast
 - new: New version of Onnx Runtime 1.20 for python version >= 3.10
+- new: Use `torch.compile` path in heuristic search for max batch size
 - change: Removed TensorFlow dependencies for `nav.jax.optimize`
 - change: Removed PyTorch dependencies from `nav.profile`
 - change: Collect all Python packages in status instead of filtered list

model_navigator/pipelines/builders/find_device_max_batch_size.py

@@ -28,7 +28,7 @@
 from model_navigator.pipelines.pipeline import Pipeline
 from model_navigator.runners.onnx import OnnxrtCPURunner, OnnxrtCUDARunner
 from model_navigator.runners.tensorflow import TensorFlowSavedModelCPURunner, TensorFlowSavedModelCUDARunner
-from model_navigator.runners.torch import TorchCPURunner, TorchCUDARunner
+from model_navigator.runners.torch import TorchCompileCPURunner, TorchCompileCUDARunner, TorchCPURunner, TorchCUDARunner
 from model_navigator.utils.config_helpers import do_find_device_max_batch_size
 
 
@@ -77,22 +77,24 @@ def find_device_max_batch_size_builder(
 def _find_max_batch_size_config_for_torch(config: CommonConfig, models_config: Dict[Format, List[ModelConfig]]):
     configurations = []
     for model_cfg in models_config.get(Format.TORCH, []):
-        runner_cls = {
-            DeviceKind.CUDA: TorchCUDARunner,
-            DeviceKind.CPU: TorchCPURunner,
+        runners_cls = {
+            DeviceKind.CUDA: [TorchCUDARunner, TorchCompileCUDARunner],
+            DeviceKind.CPU: [TorchCPURunner, TorchCompileCPURunner],
         }[config.target_device]
 
-        if model_cfg.format != runner_cls.format():
-            raise ModelNavigatorRuntimeError(
-                f"Model config format `{model_cfg.format}` does not match `{runner_cls.format()}`."
+        for runner_cls in runners_cls:
+            if model_cfg.format != runner_cls.format():
+                raise ModelNavigatorRuntimeError(
+                    f"Model config format `{model_cfg.format}` does not match `{runner_cls.format()}`."
+                )
+            mbs_config = FindMaxBatchSizeConfig(
+                format=Format.TORCH,
+                model=config.model,
+                runner_cls=runner_cls,
+                reproduction_scripts_dir=pathlib.Path(model_cfg.key),
             )
-        mbs_config = FindMaxBatchSizeConfig(
-            format=Format.TORCH,
-            model=config.model,
-            runner_cls=runner_cls,
-            reproduction_scripts_dir=pathlib.Path(model_cfg.key),
-        )
-        configurations.append(mbs_config)
+            configurations.append(mbs_config)
+
     for model_cfg in models_config.get(Format.ONNX, []):
         runner_cls = {
             DeviceKind.CUDA: OnnxrtCUDARunner,

tests/unit/base/test_pipelines_builders_find_device_max_batch_size.py

@@ -31,7 +31,7 @@
 from model_navigator.pipelines.builders.find_device_max_batch_size import find_device_max_batch_size_builder
 from model_navigator.runners.onnx import OnnxrtCUDARunner
 from model_navigator.runners.tensorflow import TensorFlowSavedModelCUDARunner
-from model_navigator.runners.torch import TorchCUDARunner
+from model_navigator.runners.torch import TorchCompileCUDARunner, TorchCUDARunner
 
 
 def test_find_device_max_batch_size_builder_return_execution_unit_when_torch_framework_is_used():
@@ -84,12 +84,15 @@ def test_find_device_max_batch_size_builder_return_execution_unit_when_torch_fra
     assert len(pipeline.execution_units) == 1
 
     execution_unit = pipeline.execution_units[0]
-    assert len(execution_unit.kwargs["configurations"]) == 2
+    assert len(execution_unit.kwargs["configurations"]) == 3
 
     configuration = execution_unit.kwargs["configurations"][0]
     assert configuration.runner_cls == TorchCUDARunner
 
     configuration = execution_unit.kwargs["configurations"][1]
+    assert configuration.runner_cls == TorchCompileCUDARunner
+
+    configuration = execution_unit.kwargs["configurations"][2]
     assert configuration.runner_cls == OnnxrtCUDARunner