Revert back to support lazy init while reducing the memory consumption

Author: cehongwang

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -734,11 +734,8 @@ def run(
                     if interpreter_result is not None:  # hit the cache
                         return interpreter_result  # type: ignore[no-any-return]
 
-        _LOGGER.debug(
-            f"CPU memory usage before network construction: {get_cpu_memory_usage()} MB"
-        )
         self._construct_trt_network_def()
-        _LOGGER.debug(
+        _LOGGER.info(
             f"CPU memory usage after network construction: {get_cpu_memory_usage()} MB"
         )
 
@@ -758,16 +755,16 @@ def run(
         self._create_timing_cache(
             builder_config, self.compilation_settings.timing_cache_path
         )
-        _LOGGER.debug(
-            f"CPU memory usage before engine building: {get_cpu_memory_usage()} MB"
-        )
+
         cuda_engine = self.builder.build_engine_with_config(
             self.ctx.net, builder_config
         )
         assert cuda_engine
+
         _LOGGER.debug(
             f"CPU memory usage after engine building: {get_cpu_memory_usage()} MB"
         )
+
         _LOGGER.info(
             f"Build TRT engine elapsed time: {datetime.now() - build_engine_start_time}"
         )

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import io
 import logging
 from typing import Any, List, Optional, Sequence
 
@@ -33,7 +34,7 @@ def infer_module_output_dtypes(
     """
     outputs = [node for node in module.graph.nodes if node.op == "output"]
     outputs = outputs[0].args
-    return get_output_dtypes(outputs, truncate_double)
+    return get_output_dtypes(outputs, truncate_double)  # type: ignore[no-any-return]
 
 
 def interpret_module_to_result(
@@ -113,11 +114,16 @@ def convert_module(
             delattr(module, attr)
     release_memory()
     logger.debug(
-        f"CPU memory usage after clearing frozen parameters and building memory: {get_cpu_memory_usage()} MB"
+        f"CPU memory usage after clearing frozen parameters and building memory in conversion: {get_cpu_memory_usage()} MB"
     )
 
+    serialized_engine = interpreter_result.engine.serialize()
+    with io.BytesIO() as engine_bytes:
+        engine_bytes.write(serialized_engine)
+        serialized_engine = engine_bytes.getvalue()
+    breakpoint()
     return rt_cls(
-        cuda_engine=interpreter_result.engine,
+        serialized_engine=serialized_engine,
         input_binding_names=list(interpreter_result.input_names),
         output_binding_names=list(interpreter_result.output_names),
         name=name,

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -123,7 +123,6 @@ class PythonTorchTensorRTModule(Module):  # type: ignore[misc]
 
     def __init__(
         self,
-        cuda_engine: trt.ICudaEngine = None,
         serialized_engine: Optional[bytes] = None,
         input_binding_names: Optional[List[str]] = None,
         output_binding_names: Optional[List[str]] = None,
@@ -183,19 +182,7 @@ def __init__(
         # Unused currently - to be used by Dynamic Shape support implementation
         self.memory_pool = None
 
-        if cuda_engine:
-            assert isinstance(
-                cuda_engine, trt.ICudaEngine
-            ), "Cuda engine must be a trt.ICudaEngine object"
-            self.engine = cuda_engine
-        elif serialized_engine:
-            assert isinstance(
-                serialized_engine, bytes
-            ), "Serialized engine must be a bytes object"
-            self.engine = serialized_engine
-        else:
-            raise ValueError("Serialized engine or cuda engine must be provided")
-
+        self.serialized_engine = serialized_engine
         self.input_names = (
             input_binding_names if input_binding_names is not None else []
         )
@@ -217,6 +204,7 @@ def __init__(
             else False
         )
         self.settings = settings
+        self.engine = None
         self.weight_name_map = weight_name_map
         self.target_platform = Platform.current_platform()
         self.runtime_states = TorchTRTRuntimeStates(
@@ -231,7 +219,7 @@ def __init__(
         self.output_allocator: Optional[DynamicOutputAllocator] = None
         self.use_output_allocator_outputs = False
 
-        if self.engine and not self.settings.lazy_engine_init:
+        if self.serialized_engine is not None and not self.settings.lazy_engine_init:
             self.setup_engine()
 
     def get_streamable_device_memory_budget(self) -> Any:
@@ -272,22 +260,13 @@ def set_default_device_memory_budget(self) -> int:
         return self._set_device_memory_budget(budget_bytes)
 
     def setup_engine(self) -> None:
-
-        if isinstance(self.engine, trt.ICudaEngine):
-            pass
-        elif isinstance(self.engine, bytes):
-            runtime = trt.Runtime(TRT_LOGGER)
-            self.engine = runtime.deserialize_cuda_engine(self.engine)
-        else:
-            raise ValueError(
-                "Expected engine as trt.ICudaEngine or serialized engine as bytes"
-            )
-
         assert (
             self.target_platform == Platform.current_platform()
         ), f"TensorRT engine was not built to target current platform (target: {self.target_platform}, current: {Platform.current_platform()})"
 
         self.initialized = True
+        runtime = trt.Runtime(TRT_LOGGER)
+        self.engine = runtime.deserialize_cuda_engine(self.serialized_engine)
         if self.settings.enable_weight_streaming:
             self.set_default_device_memory_budget()
         self.context = self.engine.create_execution_context()
@@ -323,7 +302,7 @@ def _check_initialized(self) -> None:
             raise RuntimeError("PythonTorchTensorRTModule is not initialized.")
 
     def _on_state_dict(self, state_dict: Dict[str, Any], prefix: str, _: Any) -> None:
-        state_dict[prefix + "engine"] = self.engine
+        state_dict[prefix + "engine"] = self.serialized_engine
         state_dict[prefix + "input_names"] = self.input_names
         state_dict[prefix + "output_names"] = self.output_names
         state_dict[prefix + "platform"] = self.target_platform
@@ -338,7 +317,7 @@ def _load_from_state_dict(
         unexpected_keys: Any,
         error_msgs: Any,
     ) -> None:
-        self.engine = state_dict[prefix + "engine"]
+        self.serialized_engine = state_dict[prefix + "engine"]
         self.input_names = state_dict[prefix + "input_names"]
         self.output_names = state_dict[prefix + "output_names"]
         self.target_platform = state_dict[prefix + "platform"]

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -2,12 +2,10 @@
 
 import base64
 import copy
-import io
 import logging
 import pickle
 from typing import Any, List, Optional, Tuple, Union
 
-import tensorrt as trt
 import torch
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import Platform
@@ -78,7 +76,6 @@ class TorchTensorRTModule(torch.nn.Module):  # type: ignore[misc]
 
     def __init__(
         self,
-        cuda_engine: Optional[trt.ICudaEngine | bytes] = None,
         serialized_engine: Optional[bytes] = None,
         input_binding_names: Optional[List[str]] = None,
         output_binding_names: Optional[List[str]] = None,
@@ -126,22 +123,8 @@ def __init__(
         """
         super(TorchTensorRTModule, self).__init__()
 
-        if serialized_engine:
-            assert isinstance(
-                serialized_engine, bytes
-            ), "Serialized engine must be a bytes object"
-            self.serialized_engine = serialized_engine
-
-        elif cuda_engine:
-            assert isinstance(
-                cuda_engine, trt.ICudaEngine
-            ), "Cuda engine must be a trt.ICudaEngine object"
-            serialized_engine = cuda_engine.serialize()
-            with io.BytesIO() as engine_bytes:
-                engine_bytes.write(serialized_engine)  # type: ignore
-                self.serialized_engine = engine_bytes.getvalue()
-        else:
-            raise ValueError("Serialized engine or cuda engine must be provided")
+        if not isinstance(serialized_engine, bytearray):
+            ValueError("Expected serialized engine as bytearray")
 
         self.input_binding_names = (
             input_binding_names if input_binding_names is not None else []
@@ -153,11 +136,12 @@ def __init__(
         self.hardware_compatible = settings.hardware_compatible
         self.settings = copy.deepcopy(settings)
         self.weight_name_map = weight_name_map
+        self.serialized_engine = serialized_engine
         self.engine = None
         self.requires_output_allocator = requires_output_allocator
 
         if (
-            self.serialized_engine
+            serialized_engine
             and not self.settings.lazy_engine_init
             and not self.settings.enable_cross_compile_for_windows
         ):