update refit usage

Author: zewenli98

py/torch_tensorrt/dynamo/_compiler.py

@@ -236,6 +236,9 @@ def compile(
     logger.debug("Lowered Input graph: " + str(gm.graph))
 
     if cache_built_engines or reuse_cached_engines:
+        assert (
+            make_refitable
+        ), "Engine caching requires make_refitable to be set to True"
         if custom_engine_cache is None:
             custom_engine_cache = DiskEngineCache(engine_cache_dir, engine_cache_size)
 

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -501,6 +501,24 @@ def run(
                     _LOGGER.info(
                         "Found the cached engine that corresponds to this graph. It is directly loaded."
                     )
+
+                    from torch_tensorrt.dynamo._refit import (
+                        _refit_single_trt_engine_with_gm,
+                    )
+
+                    runtime = trt.Runtime(TRT_LOGGER)
+                    engine = runtime.deserialize_cuda_engine(serialized_engine)
+
+                    _refit_single_trt_engine_with_gm(
+                        new_gm=self.module,
+                        old_engine=engine,
+                        input_list=self.input_specs,
+                        settings=self.compilation_settings,
+                        weight_name_map=weight_name_map,
+                    )
+
+                    serialized_engine = bytes(engine.serialize())
+
                     return TRTInterpreterResult(
                         serialized_engine,
                         self._input_names,

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -107,36 +107,6 @@ def interpret_module_to_result(
         compilation_settings=settings,
     )
     interpreter_result = interpreter.run()
-
-    if settings.make_refitable:
-        # Run fast refit even if it's the first compilation.
-        # This is to ensure that the weight name map is correct for future refits.
-        # If the fast refit fails, remove the weight name map.
-        from torch_tensorrt.dynamo._refit import _refit_single_trt_engine_with_gm
-        from torch_tensorrt.logging import TRT_LOGGER
-
-        runtime = trt.Runtime(TRT_LOGGER)
-        refit_test_engine = runtime.deserialize_cuda_engine(
-            interpreter_result.serialized_engine
-        )
-        try:
-            _refit_single_trt_engine_with_gm(
-                new_gm=module,
-                old_engine=refit_test_engine,
-                input_list=inputs,
-                settings=settings,
-                weight_name_map=interpreter_result.weight_name_map,
-            )
-        except AssertionError:
-            # TRTInterpreterResult is a tuple, so we need to create a new one
-            interpreter_result = TRTInterpreterResult(
-                interpreter_result.serialized_engine,
-                interpreter_result.input_names,
-                interpreter_result.output_names,
-                None,
-            )
-            logger.warning("Fast refit test failed. Removing the weight map caching.")
-
     return interpreter_result
 
 

py/torch_tensorrt/dynamo/utils.py

@@ -359,6 +359,9 @@ def parse_dynamo_kwargs(kwargs: Any) -> CompilationSettings:
     # If cache_built_engines and reuse_cached_engines are True but custom_engine_cache is not provided,
     # then create a default disk engine cache
     if kwargs.get("cache_built_engines") or kwargs.get("reuse_cached_engines"):
+        assert kwargs.get(
+            "make_refitable"
+        ), "Engine caching requires make_refitable to be set to True"
         if settings.custom_engine_cache is None:
             from torch_tensorrt.dynamo._engine_caching import DiskEngineCache