Initial attempt to delete one copy of the model

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -702,7 +702,8 @@ def compile(
 
     # Move the weights in the state_dict to CPU
     if offload_module_to_cpu:
-        deallocate_module(exported_program.module(), delete_module=False)
+        deallocate_module(gm, delete_module=False)
+        # deallocate_module(exported_program.module(), delete_module=False)
         logger.info(
             "The PyTorch model was moved to the CPU to allocate all GPU memory to TensorRT. To retain the model on the GPU, set offload_module_to_cpu=False"
         )

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -494,8 +494,7 @@ def _save_weight_mapping(self) -> None:
         _LOGGER.info("Building weight name mapping...")
         # Stage 1: Name mapping
         torch_device = to_torch_device(self.compilation_settings.device)
-        self.module.to(torch_device)
-        sd = self.module.state_dict()
+        sd = {k: v.to(torch_device) for k, v in self.module.state_dict().items()}
         weight_name_map: dict[str, Any] = {}
         weight_refit_map = self.ctx.weight_refit_map
         constant_mapping = {k: v for k, v in weight_refit_map.items() if v.size == 1}

py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py

@@ -37,7 +37,9 @@ def constant_fold(
     # For TRT INetwork construction the constants are moved to CPU in get_attr call.
     for node, constant in cf.node_replacements.items():
         replace_node_with_constant(
-            gm, node, torch.nn.Parameter(constant, requires_grad=False)
+            gm,
+            node,
+            torch.nn.Parameter(constant.cpu().contiguous(), requires_grad=False),
         )
 
     erased_params = []

tools/perf/Flux/flux_perf.py

@@ -4,7 +4,7 @@
 from time import time
 
 sys.path.append(os.path.join(os.path.dirname(__file__), "../../../examples/apps"))
-from flux_demo import compile_model, parse_args
+from flux_demo import compile_model
 
 
 def benchmark(pipe, prompt, inference_step, batch_size=1, iterations=1):