Added fuse_rms_norm lowering

Author: cehongwang

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

@@ -18,13 +18,15 @@
 from .remove_input_alias_fixing_clones import remove_input_alias_fixing_clones
 from .remove_num_users_is_0_nodes import remove_num_users_is_0_nodes
 from .repair_input_as_output import repair_input_as_output
+from .replace_fused_rms_norm import replace_fused_rms_norm
 from .replace_max_pool_with_indices import replace_max_pool_with_indices
 from .rule_based_autocast import rule_based_autocast
 
 pre_lowering_pass_list = [
     remove_detach,
     remove_assert_nodes,
     rule_based_autocast,
+    replace_fused_rms_norm,
 ]
 
 post_lowering_pass_list = [

py/torch_tensorrt/dynamo/lowering/passes/replace_fused_rms_norm.py

@@ -0,0 +1,27 @@
+import logging
+
+import torch
+from torch_tensorrt.dynamo._settings import CompilationSettings
+
+logger = logging.getLogger(__name__)
+
+
+def replace_fused_rms_norm(
+    gm: torch.fx.GraphModule, settings: CompilationSettings
+) -> torch.fx.GraphModule:
+    """Replace fused rms norm ops in the graph"""
+    count = 0
+    for node in gm.graph.nodes:
+        if node.target == torch.ops.aten._fused_rms_norm.default:
+            # Replace fused rms norm with standard rms norm
+            new_node = gm.graph.call_function(
+                torch.nn.functional.rms_norm.default,
+                args=node.args,
+            )
+            gm.graph.replace_node_with_new_node(node, new_node)
+            gm.graph.erase_node(node)
+            count += 1
+
+    logger.debug(f"Replaced {count} fused rms norm nodes:\n{gm.graph}")
+
+    return gm