move assertion check cleanup back to stock export (#93)

Signed-off-by: Lucas Liebenwein <[email protected]>

Author: lucaslie

tensorrt_llm/_torch/auto_deploy/config/default.yaml

@@ -17,7 +17,5 @@ transforms:
     stage: post_export
   cleanup_noop_add:
     stage: post_export
-  cleanup_checks:
-    stage: post_export
   cleanup_input_constraints:
     stage: post_export

tensorrt_llm/_torch/auto_deploy/export/export.py

@@ -17,6 +17,7 @@
     tree_to,
 )
 from ..utils.logger import ad_logger
+from ..utils.node_utils import is_op
 from .interface import ExportPatchRegistry, apply_export_patches
 
 try:
@@ -176,6 +177,24 @@ def aliasing_load_pre_hook(state_dict: Dict[str, torch.Tensor], prefix: str, *ar
     gm._register_load_state_dict_pre_hook(aliasing_load_pre_hook)
 
 
+def _clean_up_assertions(gm: fx.GraphModule):
+    """This transformations removes shape checks and assertions from the graph."""
+    check_ops = {
+        torch.ops.aten._assert_scalar,
+        torch.ops.aten.sym_constrain_range,
+        torch.ops.aten.sym_constrain_range_for_size,
+        torch.ops.aten._assert_tensor_metadata,
+        # torch.ops.aten._functional_sym_constrain_range,
+        # torch.ops.aten._functional_sym_constrain_range_for_size
+    }
+    graph: fx.Graph = gm.graph
+    for node in reversed(graph.nodes):
+        if len(node.users) > 0 or not is_op(node, check_ops):
+            continue
+        graph.erase_node(node)
+    canonicalize_graph(gm)
+
+
 def torch_export_to_gm(
     model: nn.Module,
     args: Tuple[Any, ...],
@@ -196,6 +215,7 @@ def torch_export_to_gm(
         3. Automatically extract the GraphModule from the exported program.
         4. Retain load hooks for state_dict loading from the original module.
         5. Manage parameter aliasing in the model.
+        6. Remove assertions from the graph.
 
     Args:
         model: The model to export
@@ -255,6 +275,9 @@ def torch_export_to_gm(
     # This is a consequence of lifting to meta during export.
     _clean_up_device_info(egm)
 
+    # clean up checks --> generally the sanity checks are overly conservative and we can remove them
+    _clean_up_assertions(egm)
+
     # show exported graph
     ad_logger.debug("exported graph: " + str(egm))
 

tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_checks.py

@@ -10,6 +10,8 @@
 from ..interface import BaseTransform, TransformInfo, TransformRegistry
 
 
+# TODO (lucaslie): consider reconfiguring this transform to run before we switch to flattened
+# sequences which is done in update_in_out_nodes at the moment.
 @TransformRegistry.register("cleanup_input_constraints")
 class CleanupInputConstraints(BaseTransform):
     """Cleanup input constraints from the graph.