revert rms_norm.py

Author: shewu-quic

examples/models/llama/TARGETS

@@ -113,6 +113,7 @@ runtime.python_library(
         "source_transformation/prune_vocab.py",
         "source_transformation/quantize.py",
         "source_transformation/custom_kv_cache.py",
+        "source_transformation/rms_norm.py",
         "source_transformation/rope.py",
         "source_transformation/sdpa.py",
         "source_transformation/spin_quant.py",

examples/models/llama/source_transformation/rms_norm.py

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.	
+# All rights reserved.	
+#	
+# This source code is licensed under the BSD-style license found in the	
+# LICENSE file in the root directory of this source tree.	
+
+import torch	
+from executorch.examples.models.llama.llama_transformer import RMSNorm	
+
+
+def replace_rms_norm_with_native_rms_norm(module: torch.nn.Module):	
+    for name, child in module.named_children():	
+        if isinstance(child, RMSNorm):	
+            rms_norm = torch.nn.RMSNorm(child.dim, eps=child.eps)	
+            rms_norm.weight = child.weight	
+            setattr(	
+                module,	
+                name,	
+                rms_norm,	
+            )	
+        else:	
+            replace_rms_norm_with_native_rms_norm(child)	
+    return module