chore: bug fixes for refit tests, restructure CI tests

Author: peri044

.github/workflows/build-test-linux.yml

@@ -172,8 +172,12 @@ jobs:
         cd tests/py
         python -m pip install -r requirements.txt
         cd dynamo
-        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dyn_models_export.xml --ir dynamo models/
-        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dyn_models_export.xml test_modelopt_models.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_models.xml --ir dynamo models/test_models.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_models_dynamic.xml --ir dynamo models/test_dyn_models.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/engine_cache.xml --ir dynamo models/test_engine_cache.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/model_refit.xml --ir dynamo models/test_model_refit.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/modelopt_models.xml --ir dynamo models/test_modelopt_models.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/weight_stripped_engine.xml --ir dynamo models/test_weight_stripped_engine.py
         popd
 
   tests-py-dynamo-serde:
@@ -206,6 +210,7 @@ jobs:
         cd dynamo
         python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/export_serde_test_results.xml --ir dynamo models/test_export_serde.py
         python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/reexport_test_results.xml --ir dynamo models/test_reexport.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/export_kwargs_serde_test_results.xml --ir dynamo models/test_export_kwargs_serde.py
         popd
 
   tests-py-torch-compile-be:

py/torch_tensorrt/dynamo/_refit.py

@@ -48,7 +48,7 @@
 logger = logging.getLogger(__name__)
 
 
-@needs_refit
+@needs_refit  # type: ignore
 def construct_refit_mapping(
     module: torch.fx.GraphModule,
     inputs: Sequence[Input],
@@ -110,7 +110,7 @@ def construct_refit_mapping(
     return weight_map
 
 
-@needs_refit
+@needs_refit  # type: ignore
 def construct_refit_mapping_from_weight_name_map(
     weight_name_map: dict[Any, Any],
     state_dict: dict[Any, Any],
@@ -141,7 +141,7 @@ def construct_refit_mapping_from_weight_name_map(
     return engine_weight_map
 
 
-@needs_refit
+@needs_refit  # type: ignore
 def _refit_single_trt_engine_with_gm(
     new_gm: torch.fx.GraphModule,
     old_engine: trt.ICudaEngine,
@@ -153,12 +153,12 @@ def _refit_single_trt_engine_with_gm(
     Refit a TensorRT Engine in place
     """
 
-    with unset_fake_temporarily():
-        refitted = set()
-        torch_device = get_model_device(new_gm)
-        refitter = trt.Refitter(old_engine, TRT_LOGGER)
-        weight_list = refitter.get_all_weights()
+    refitted = set()
+    torch_device = get_model_device(new_gm)
+    refitter = trt.Refitter(old_engine, TRT_LOGGER)
+    weight_list = refitter.get_all_weights()
 
+    with unset_fake_temporarily():
         if weight_name_map:
             # Get the refitting mapping
             trt_wt_location = (
@@ -185,41 +185,21 @@ def _refit_single_trt_engine_with_gm(
                     trt_dtype,
                 )
 
-                constant_mapping: dict[str, Any] = weight_name_map.pop(
-                    "constant_mapping", {}
-                )  # type: ignore
-                mapping = construct_refit_mapping_from_weight_name_map(
-                    weight_name_map, new_gm.state_dict()
-                )
-                constant_mapping_with_type = {}
-
-                for constant_name, val in constant_mapping.items():
-                    np_weight_type = val.dtype
-                    val_tensor = torch.from_numpy(val).cuda()
-                    trt_dtype = dtype.try_from(np_weight_type).to(trt.DataType)
-                    torch_dtype = dtype.try_from(np_weight_type).to(torch.dtype)
-                    constant_mapping_with_type[constant_name] = (
-                        val_tensor.clone().reshape(-1).contiguous().to(torch_dtype),
-                        trt_dtype,
-                    )
+            mapping.update(constant_mapping_with_type)
 
-                mapping.update(constant_mapping_with_type)
-
-                for layer_name in weight_list:
-                    if layer_name not in mapping:
-                        logger.warning(f"{layer_name} is not found in weight mapping.")
-                        continue
-                    # Use Numpy to create weights
-                    weight, weight_dtype = mapping[layer_name]
-                    trt_wt_tensor = trt.Weights(
-                        weight_dtype, weight.data_ptr(), torch.numel(weight)
-                    )
-                    refitter.set_named_weights(
-                        layer_name, trt_wt_tensor, trt_wt_location
-                    )
-                assert (
-                    len(refitter.get_missing_weights()) == 0
-                ), "Fast refitting failed due to incomplete mapping"
+            for layer_name in weight_list:
+                if layer_name not in mapping:
+                    logger.warning(f"{layer_name} is not found in weight mapping.")
+                    continue
+                # Use Numpy to create weights
+                weight, weight_dtype = mapping[layer_name]
+                trt_wt_tensor = trt.Weights(
+                    weight_dtype, weight.data_ptr(), torch.numel(weight)
+                )
+                refitter.set_named_weights(layer_name, trt_wt_tensor, trt_wt_location)
+            assert (
+                len(refitter.get_missing_weights()) == 0
+            ), "Fast refitting failed due to incomplete mapping"
 
         else:
             mapping = construct_refit_mapping(new_gm, input_list, settings)
@@ -241,7 +221,7 @@ def _refit_single_trt_engine_with_gm(
             raise AssertionError("Refitting failed.")
 
 
-@needs_refit
+@needs_refit  # type: ignore
 def refit_module_weights(
     compiled_module: torch.fx.GraphModule | ExportedProgram,
     new_weight_module: ExportedProgram,