removes fused_rmsnorm

pytorch · Feb 6, 2025 · aa02c15 · aa02c15
1 parent 690f299
commit aa02c15
Show file tree

Hide file tree

Showing 10 changed files with 7 additions and 369 deletions.
diff --git a/scripts/estimate/estimation.py b/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"

diff --git a/tests/integration_tests.py b/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(
             [
                 [

diff --git a/tests/unit_tests/test_fused_rms_norm_dtensor.py b/tests/unit_tests/test_fused_rms_norm_dtensor.py
diff --git a/torchtitan/config_manager.py b/torchtitan/config_manager.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",