Removes fused_rmsnorm (#826)

Context in #825.

Author: allenwang28

scripts/estimate/estimation.py

@@ -33,17 +33,6 @@ def estimate_memory(job_config: JobConfig):
     # Get the world size
     world_size = int(os.environ["WORLD_SIZE"])
 
-    # fake tensor doesn't work with fused rmsnorm
-    if (
-        job_config.model.norm_type == "fused_rmsnorm"
-        and not job_config.memory_estimation.disable_fake_mode
-    ):
-        logger.info(
-            "Fused RMSNorm is not supported yet under fake estimation mode. "
-            "Switching to rmsnorm."
-        )
-        job_config.model.norm_type = "rmsnorm"
-
     if job_config.model.norm_type == "compiled_rmsnorm":
         logger.info("Compiled RMSNorm is not supported yet. Switching to RMSNorm.")
         job_config.model.norm_type = "rmsnorm"

tests/integration_tests.py

@@ -94,16 +94,6 @@ def build_test_list():
             "2D compile",
             "2d_compile",
         ),
-        OverrideDefinitions(
-            [
-                [
-                    "--training.tensor_parallel_degree 2",
-                    "--model.norm_type=fused_rmsnorm",
-                ],
-            ],
-            "2D eager with fused_rmsnorm",
-            "2d_eager_fused_rmsnorm",
-        ),
         OverrideDefinitions(
             [
                 [

tests/unit_tests/test_fused_rms_norm_dtensor.py

@@ -174,7 +174,8 @@ def __init__(self):
             "--model.norm_type",
             type=str,
             default="rmsnorm",
-            help="Type of layer normalization to use [layernorm, np_layernorm, rmsnorm, fused_rmsnorm]",
+            choices=["layernorm", "np_layernorm", "rmsnorm"],
+            help="Type of layer normalization to use [layernorm, np_layernorm, rmsnorm]",
         )
         self.parser.add_argument(
             "--model.tokenizer_path",