chore: rebased from opt_out_buffer branch.

Author: keehyuna

core/runtime/TRTEngine.cpp

@@ -310,8 +310,8 @@ bool TRTEngine::set_device_memory_budget(int64_t budget) {
   if (profile_execution) {
     enable_profiling();
   }
-  // Indicates to reevaluate the runtime settings
-  has_context_changed = true;
+
+  runtime_states.set_context_changed();
 
   return result;
 }

core/runtime/TRTEngine.h

@@ -32,33 +32,49 @@ using FlattenedState = std::tuple<
 
 struct RuntimeStates {
   bool need_cudagraphs_record;
+  bool need_cudagraphs_reset;
   bool can_use_pre_allocated_outputs;
 };
 
 struct TorchTRTRuntimeStates {
   // Previous runtime states
-  bool prev_cudagraphs_enabled, prev_pre_allocated_outputs_enabled;
+  bool prev_cudagraphs_enabled = false;
+  bool prev_pre_allocated_outputs_enabled = false;
+  // Indicates to reevaluate the runtime settings as context has changed
+  bool has_context_changed = false;
 
-  // Evaluates whether certain conditions are met to enable CUDA Graph recording or to reuse pre-allocated outputs
+  // Evaluates whether certain conditions are met to enable CUDA Graph recording/reset or to reuse pre-allocated outputs
   // based on the current and previous states, as well as input shape has changed
   RuntimeStates validate_states(bool cudagraphs_enabled, bool pre_allocated_outputs_enabled, bool shape_changed) {
     bool need_cudagraphs_record = false;
     bool can_use_pre_allocated_outputs = false;
+    bool need_cudagraphs_reset = false;
 
     // Cudagraphs record is required if cudagraphs_enabled is switched to True regardless of shape change
-    if (cudagraphs_enabled && (!prev_cudagraphs_enabled || shape_changed)) {
+    // If context is changed by runtime setting like weight streaming, it needs cuda graphs record
+    if (cudagraphs_enabled && (!prev_cudagraphs_enabled || shape_changed || has_context_changed)) {
       need_cudagraphs_record = true;
     }
     // Pre-allocated output can be used when previous and current state are true without shape change
     if (prev_pre_allocated_outputs_enabled && pre_allocated_outputs_enabled && !shape_changed) {
       can_use_pre_allocated_outputs = true;
     }
+
+    if (!cudagraphs_enabled || shape_changed || has_context_changed) {
+      need_cudagraphs_reset = true;
+    }
+
+    // Reset the flag
+    has_context_changed = false;
     prev_cudagraphs_enabled = cudagraphs_enabled;
     prev_pre_allocated_outputs_enabled = pre_allocated_outputs_enabled;
 
-    RuntimeStates values = {need_cudagraphs_record, can_use_pre_allocated_outputs};
+    RuntimeStates values = {need_cudagraphs_record, need_cudagraphs_reset, can_use_pre_allocated_outputs};
     return values;
   }
+  void set_context_changed() {
+    has_context_changed = true;
+  }
 };
 
 struct TRTEngine : torch::CustomClassHolder {

core/runtime/execute_engine.cpp

@@ -205,9 +205,8 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
   // Whether cudagraphs needs to record the graph on this pass
   RuntimeStates states = compiled_engine->runtime_states.validate_states(
       CUDAGRAPHS_MODE, compiled_engine->use_pre_allocated_outputs, shape_changed);
-  bool need_cudagraphs_record = states.need_cudagraphs_record;
 
-  if (!CUDAGRAPHS_MODE || shape_changed) {
+  if (states.need_cudagraphs_reset) {
     compiled_engine->cudagraph.reset();
   }
 
@@ -269,7 +268,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
           std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->input_profile_path);
     }
 
-    setup_input_tensors(inputs, compiled_engine, need_cudagraphs_record);
+    setup_input_tensors(inputs, compiled_engine, states.need_cudagraphs_record);
 
     // Check if input shapes can be inferred.
     int32_t const io_size{compiled_engine->cuda_engine->getNbIOTensors()};
@@ -297,7 +296,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
     for (auto output_indices : compiled_engine->out_binding_map) {
       auto pyt_idx = output_indices.second;
       std::string name = compiled_engine->out_binding_names[pyt_idx];
-      if (need_cudagraphs_record) {
+      if (states.need_cudagraphs_record) {
         // If we are recording the cuda graph then we need to update the persistent output buffer
         compiled_engine->output_buffers[pyt_idx] = std::move(outputs[pyt_idx].clone());
       }
@@ -349,7 +348,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
       // Direct execution uses the caller buffers directly
       compiled_engine->exec_ctx->enqueueV3(compiled_engine->engine_stream);
     } else {
-      if (need_cudagraphs_record) {
+      if (states.need_cudagraphs_record) {
         // If cudagraphs needs to record a graph, capture the enqueueV3 call in a graph
         c10::cuda::CUDAStream recording_stream = compiled_engine->engine_stream;
         compiled_engine->cudagraph.capture_begin();

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -24,23 +24,30 @@
 
 
 class TorchTRTRuntimeStates:
-    def __init__(self, cudagraphs_enabled: bool, pre_allocated_outputs_enabled: bool):
+    def __init__(self, cudagraphs_enabled: bool):
         self.prev_cudagraphs_enabled = cudagraphs_enabled
-        self.prev_pre_allocated_outputs_enabled = pre_allocated_outputs_enabled
+        self.prev_pre_allocated_outputs_enabled = False
+        self.has_context_changed = False
 
     def validate_states(
         self,
         cudagraphs_enabled: bool,
         pre_allocated_outputs_enabled: bool,
         shape_changed: bool,
-    ) -> Tuple[bool, bool]:
-        # Evaluates whether certain conditions are met to enable CUDA Graph recording or to reuse pre-allocated outputs
+    ) -> Tuple[bool, bool, bool]:
+        # Evaluates whether certain conditions are met to enable CUDA Graph recording/reset or to reuse pre-allocated outputs
         # based on the current and previous states, as well as input shape has changed
         need_cudagraphs_record = False
         can_use_pre_allocated_outputs = False
+        need_cudagraphs_reset = False
 
         # Cudagraphs record is required if cudagraphs_enabled is switched to True regardless of shape change
-        if cudagraphs_enabled and (not self.prev_cudagraphs_enabled or shape_changed):
+        # If context is changed by runtime setting like weight streaming, it needs cuda graphs record
+        if cudagraphs_enabled and (
+            not self.prev_cudagraphs_enabled
+            or shape_changed
+            or self.has_context_changed
+        ):
             need_cudagraphs_record = True
 
         # Pre-allocated output can be used when previous and current state are true without shape change
@@ -51,10 +58,19 @@ def validate_states(
         ):
             can_use_pre_allocated_outputs = True
 
+        if not cudagraphs_enabled or shape_changed or self.has_context_changed:
+            need_cudagraphs_reset = True
+
+        # Reset the flag
+        self.has_context_changed = False
         self.prev_cudagraphs_enabled = cudagraphs_enabled
         self.prev_pre_allocated_outputs_enabled = pre_allocated_outputs_enabled
 
-        return need_cudagraphs_record, can_use_pre_allocated_outputs
+        return (
+            need_cudagraphs_record,
+            need_cudagraphs_reset,
+            can_use_pre_allocated_outputs,
+        )
 
 
 class PythonTorchTensorRTModule(Module):  # type: ignore[misc]
@@ -141,15 +157,12 @@ def __init__(
         self.engine = None
         self.weight_name_map = weight_name_map
         self.target_platform = Platform.current_platform()
-<<<<<<< HEAD
+
         self.runtime_states = TorchTRTRuntimeStates(
-            torch_tensorrt.runtime.get_cudagraphs_mode(), False
+            torch_tensorrt.runtime.get_cudagraphs_mode()
         )
         self.pre_allocated_outputs: List[torch.Tensor] = []
         self.use_pre_allocated_outputs = False
-=======
-        self.has_context_changed = False
->>>>>>> 7bb66dac4 (fix: Record cudagraphs when weight streaming budget has changed)
 
         if self.serialized_engine is not None and not self.settings.lazy_engine_init:
             self.setup_engine()
@@ -169,8 +182,7 @@ def set_device_memory_budget(self, budget_bytes: int) -> int:
             del self.context
         budget_bytes = self._set_device_memory_budget(budget_bytes)
         self.context = self.engine.create_execution_context()
-        # Indicates to reevaluate the runtime settings
-        self.has_context_changed = True
+        self.runtime_states.has_context_changed = True
 
         return budget_bytes
 
@@ -360,36 +372,25 @@ def forward(self, *inputs: torch.Tensor) -> torch.Tensor | Tuple[torch.Tensor, .
             self._check_initialized()
 
             cudagraphs_enabled = torch_tensorrt.runtime.get_cudagraphs_mode()
-<<<<<<< HEAD
+
             shape_changed = self.validate_input_shapes(inputs)
-            need_cudagraphs_record, can_use_pre_allocated_outputs = (
-                self.runtime_states.validate_states(
-                    cudagraphs_enabled, self.use_pre_allocated_outputs, shape_changed
-                )
-=======
-            need_cudagraphs_record = cudagraphs_enabled and (
-                not self.cudagraphs_validate_shapes(inputs) or self.has_context_changed
->>>>>>> 7bb66dac4 (fix: Record cudagraphs when weight streaming budget has changed)
+            (
+                need_cudagraphs_record,
+                need_cudagraphs_reset,
+                can_use_pre_allocated_outputs,
+            ) = self.runtime_states.validate_states(
+                cudagraphs_enabled, self.use_pre_allocated_outputs, shape_changed
             )
 
             if need_cudagraphs_record:
-                if self.cudagraph:
-                    self.cudagraph.reset()
                 self._input_buffers = [None] * len(self.input_names)
                 self._output_buffers = [None] * len(self.output_names)
 
-            if self.cudagraph and (not cudagraphs_enabled or self.has_context_changed):
+            if need_cudagraphs_reset and self.cudagraph:
                 self.cudagraph.reset()
                 self.cudagraph = None
 
-<<<<<<< HEAD
             # If in safe mode, check at each iteration for whether a switch is required
-=======
-            # Reset the flag
-            self.has_context_changed = False
-
-            # If in safe mode, check at each iteration for for whether a switch is required
->>>>>>> 7bb66dac4 (fix: Record cudagraphs when weight streaming budget has changed)
             if (
                 torch_tensorrt.runtime._multi_device_safe_mode._PY_RT_MULTI_DEVICE_SAFE_MODE
             ):

tests/py/dynamo/runtime/test_pre_allocated_outputs.py

@@ -1,3 +1,5 @@
+import itertools
+
 import torch
 import torch_tensorrt as torchtrt
 from parameterized import parameterized
@@ -62,67 +64,94 @@ def forward(self, x):
     )
     def test_pre_allocated_outputs_dynamic(self, _, use_python_runtime):
         class SampleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer1 = torch.nn.Linear(100, 128)
+                self.layer2 = torch.nn.Linear(128, 64)
+                self.relu = torch.nn.ReLU()
+
             def forward(self, x):
-                return torch.relu((x + 2) * 0.5)
+                out = self.layer1(x)
+                out = self.relu((out + 2.0) * 0.05)
+                out = self.layer2(out)
+                return out
 
         inputs = torchtrt.Input(
-            min_shape=(1, 3, 128, 224),
-            opt_shape=(8, 3, 192, 224),
-            max_shape=(16, 3, 224, 224),
+            min_shape=(1, 100),
+            opt_shape=(64, 100),
+            max_shape=(128, 100),
             dtype=torch.float,
             name="x",
         )
-        fx_graph = torch.fx.symbolic_trace(SampleModel())
+        model = SampleModel().eval().cuda()
+        fx_graph = torch.fx.symbolic_trace(model)
+
+        input_list = []
+        input_list.append(torch.randn((8, 100)).cuda())
+        input_list.append(torch.randn((12, 100)).cuda())
+        input_list.append(torch.randn((12, 100)).cuda())
+        input_list.append(torch.randn((8, 100)).cuda())
+        input_list.append(torch.randn((8, 100)).cuda())
 
         optimized_model = torchtrt.compile(
             fx_graph,
             "dynamo",
             inputs,
             min_block_size=1,
             pass_through_build_failures=True,
+            use_explicit_typing=True,
+            enable_weight_streaming=True,
             torch_executed_ops={"torch.ops.aten.mul.Tensor"},
             use_python_runtime=use_python_runtime,
         )
 
-        input_list = []
-        ref_out_list = []
-        trt_out_list = []
-        # Alternating cuda_graphs enable and input shapes at every five iterations.
-        for i in [1, 3, 8, 11, 16]:
-            for j in [128, 128, 222, 222, 224]:
-                input_list.append(torch.randn((i, 3, j, 224)).cuda())
+        # List of tuples representing different configurations for three features:
+        # Cuda graphs, pre-allocated output buffer, weight streaming change
+        states = list(itertools.product((True, False), repeat=3))
+        # Create pairs of these configurations, representing an initial state and a changed state
+        states_permutations = itertools.permutations(states, 2)
 
         pre_allocated_output_ctx = torchtrt.runtime.enable_pre_allocated_outputs(
             optimized_model
         )
-        pre_allocated_output = False
-        for enable_cuda_graphs in [False, True]:
-            for i in range(len(input_list)):
-                # Toggles cuda graph at all index in TRIALS
-                if i % TRIALS == i // TRIALS:
-                    cuda_graphs = enable_cuda_graphs
-                else:
-                    cuda_graphs = not enable_cuda_graphs
-                if i % 3 == 0:
-                    pre_allocated_output = not pre_allocated_output
-
+        weight_streaming_ctx = torchtrt.runtime.weight_streaming(optimized_model)
+        streamable_budget = weight_streaming_ctx.total_device_budget
+
+        for init_state, changed_state in states_permutations:
+            for cuda_graphs, pre_allocated_output, weight_streaming in [
+                init_state,
+                changed_state,
+            ]:
                 torchtrt.runtime.set_cudagraphs_mode(cuda_graphs)
                 pre_allocated_output_ctx.set_pre_allocated_output(pre_allocated_output)
 
-                ref_out_list.append(fx_graph(input_list[i]))
-                trt_out_list.append(optimized_model(input_list[i]))
-
-        for torch_model_results, optimized_model_results in zip(
-            ref_out_list, trt_out_list
-        ):
-            torch.testing.assert_close(
-                torch_model_results,
-                optimized_model_results,
-                rtol=5e-03,
-                atol=5e-03,
-                equal_nan=True,
-                check_dtype=True,
-            )
+                if weight_streaming:
+                    weight_streaming_ctx.device_budget = int(streamable_budget * 0.8)
+                else:
+                    weight_streaming_ctx.device_budget = streamable_budget
+
+                ref_out_list = []
+                trt_out_list = []
+                # Input shape changes
+                for i in range(len(input_list)):
+                    if weight_streaming and i == 4:
+                        weight_streaming_ctx.device_budget = int(
+                            streamable_budget * 0.6
+                        )
+                    ref_out_list.append(fx_graph(input_list[i]))
+                    trt_out_list.append(optimized_model(input_list[i]))
+
+                for torch_model_results, optimized_model_results in zip(
+                    ref_out_list, trt_out_list
+                ):
+                    torch.testing.assert_close(
+                        torch_model_results,
+                        optimized_model_results,
+                        rtol=5e-03,
+                        atol=5e-03,
+                        equal_nan=True,
+                        check_dtype=True,
+                    )
         torch._dynamo.reset()