refactor tests to call modified LayerNormMLP

Signed-off-by: Jaime Cardenas <[email protected]>

tests/pytorch/selective_layernorm_mlp/compare.py

@@ -1,6 +1,6 @@
 import time
 import torch
-from transformer_engine.pytorch import SelectiveLayerNormMLP
+from transformer_engine.pytorch import LayerNormMLP
 from collections import defaultdict
 
 torch.manual_seed(1234)
@@ -32,10 +32,10 @@ def build(self):
 
         ln_list, sln_list = [], []
         for _ in range(self._layers):
-            ln = SelectiveLayerNormMLP(
+            ln = LayerNormMLP(
                 self._hidden_size, self._ffn_hidden_size, checkpoint=False
             ).to(device)
-            sln = SelectiveLayerNormMLP(
+            sln = LayerNormMLP(
                 self._hidden_size, self._ffn_hidden_size, checkpoint=True
             ).to(device)
             with torch.no_grad():
@@ -180,7 +180,7 @@ def _run_bwd(model, out):
             self.stats[desc]["diff"][key] = self._max_diff(ln_grads[key], sln_grads[key])
 
     def summarize(self):
-        _modules = [("ln_stats", "LayerNormMLP"), ("sln_stats", "SelectiveLayerNormMLP")]
+        _modules = [("ln_stats", "No Checkpointing"), ("sln_stats", "Checkpointing")]
         _metric_map = {"time": (1, "ms"), "mem": (1e-6, "MB")}
 
         left_w = 18  # "fwd time" / "bwd mem" label

tests/pytorch/selective_layernorm_mlp/distributed/run_numerics.py

@@ -162,7 +162,7 @@ def main(argv=None, namespace=None):
         HIDDEN_SIZE = 512
 
     test_dict = [
-        test_selective_layernorm_mlp,
+        test_layernorm_mlp,
     ]
 
     for test in test_dict:
@@ -378,13 +378,13 @@ def _alloc_main_grad(model_single_node, model_distributed):
 
 
 ############################################
-#          SelectiveLayerNormMLP          #
+#          LayerNormMLP          #
 ############################################
 
 
 @run_distributed_test()
-def _test_selective_layernorm_mlp(set_parallel_mode=None, sequence_parallel=False, **kwargs):
-    """Test the SelectiveLayerNormMLP with specified parallel mode and sequence parallelization.
+def _test_layernorm_mlp(set_parallel_mode=None, sequence_parallel=False, **kwargs):
+    """Test the LayerNormMLP with specified parallel mode and sequence parallelization.
 
     Args:
         set_parallel_mode (bool): Enable parallel mode.
@@ -396,8 +396,8 @@ def _test_selective_layernorm_mlp(set_parallel_mode=None, sequence_parallel=Fals
     FFN_HIDDEN_SIZE = 32 if QUANTIZATION is None else 128
 
     # Create models
-    model_single_node = te.SelectiveLayerNormMLP(HIDDEN_SIZE, FFN_HIDDEN_SIZE, **kwargs)
-    model_distributed = te.SelectiveLayerNormMLP(
+    model_single_node = te.LayerNormMLP(HIDDEN_SIZE, FFN_HIDDEN_SIZE, **kwargs)
+    model_distributed = te.LayerNormMLP(
         HIDDEN_SIZE,
         FFN_HIDDEN_SIZE,
         tp_size=WORLD_SIZE,
@@ -464,7 +464,7 @@ def _test_selective_layernorm_mlp(set_parallel_mode=None, sequence_parallel=Fals
         )
 
 
-def test_selective_layernorm_mlp():
+def test_layernorm_mlp():
     kwargs_list = [
         {},
         {"init_method": _constant},
@@ -485,4 +485,4 @@ def test_selective_layernorm_mlp():
     for kwargs in kwargs_list:
         for set_parallel_mode in [True]:
             for sequence_parallel in [False, True]:
-                _test_selective_layernorm_mlp(set_parallel_mode, sequence_parallel, **kwargs)
+                _test_layernorm_mlp(set_parallel_mode, sequence_parallel, **kwargs)

tests/pytorch/selective_layernorm_mlp/test_cuda_graphs.py

@@ -7,7 +7,7 @@
 
 import torch
 from transformer_engine.pytorch import (
-    SelectiveLayerNormMLP,
+    LayerNormMLP,
     autocast,
     quantized_model_init,
     make_graphed_callables,
@@ -165,7 +165,7 @@ def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor:
 
 
 # Supported modules
-_test_cuda_graphs_modules: List[str] = ["selective_layernorm_mlp"]
+_test_cuda_graphs_modules: List[str] = ["layernorm_mlp"]
 
 
 def _test_cuda_graphs(
@@ -192,9 +192,9 @@ def _test_cuda_graphs(
     # Create modules.
     with quantized_model_init(enabled=fp8_params, recipe=fp8_recipe):
 
-        if module == "selective_layernorm_mlp":
+        if module == "layernorm_mlp":
             modules = [
-                SelectiveLayerNormMLP(
+                LayerNormMLP(
                     model_config.hidden_size,
                     model_config.hidden_size,
                     params_dtype=dtype,
@@ -322,7 +322,7 @@ def test_make_graphed_callables(
 
 
 _test_make_graphed_callables_with_fp8_weight_caching_modules = [
-    "selective_layernorm_mlp",
+    "layernorm_mlp",
 ]
 
 

tests/pytorch/selective_layernorm_mlp/test_deferred_init.py

@@ -8,7 +8,7 @@
 import transformer_engine.pytorch as te
 
 _core_modules = [
-    te.SelectiveLayerNormMLP,
+    te.LayerNormMLP,
 ]
 _composed_modules = []
 
@@ -26,7 +26,7 @@ def get_module_args(module, checkpoint):
         hidden_size = num_heads * head_dim
         args = (hidden_size,)
         kwargs = {"params_dtype": dtype, "device": "meta"}
-        if module == te.SelectiveLayerNormMLP:
+        if module == te.LayerNormMLP:
             ffn_hidden_size = 2 * hidden_size
             args += (ffn_hidden_size,)
             kwargs["bias"] = True

tests/pytorch/selective_layernorm_mlp/test_numerics.py

@@ -22,7 +22,7 @@
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 from transformer_engine.pytorch import (
     autocast,
-    SelectiveLayerNormMLP,
+    LayerNormMLP,
     get_device_compute_capability,
     is_fp8_available,
     is_mxfp8_available,
@@ -388,13 +388,13 @@ def reset_global_fp8_state():
 @pytest.mark.parametrize("return_bias", all_boolean)
 @pytest.mark.parametrize("bias", all_boolean)
 @pytest.mark.parametrize("checkpoint", all_boolean)
-def test_selective_layernorm_mlp_accuracy(
+def test_layernorm_mlp_accuracy(
     dtype, bs, model, activation, normalization, return_bias, bias, checkpoint
 ):
     config = model_configs[model]
 
     te_ln_mlp = TestReturnBiasModule(
-        SelectiveLayerNormMLP,
+        LayerNormMLP,
         hidden_size=config.hidden_size,
         ffn_hidden_size=4 * config.hidden_size,
         activation=activation,
@@ -466,12 +466,12 @@ def test_selective_layernorm_mlp_accuracy(
 @pytest.mark.parametrize("bias", all_boolean)
 @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
 @pytest.mark.parametrize("checkpoint", all_boolean)
-def test_selective_layernorm_mlp_accuracy_delay_wgrad_compute(
+def test_layernorm_mlp_accuracy_delay_wgrad_compute(
     dtype, bs, model, bias, fuse_wgrad_accumulation, checkpoint
 ):
     config = model_configs[model]
 
-    ln_mlp = SelectiveLayerNormMLP(
+    ln_mlp = LayerNormMLP(
         hidden_size=config.hidden_size,
         ffn_hidden_size=4 * config.hidden_size,
         eps=config.eps,
@@ -483,7 +483,7 @@ def test_selective_layernorm_mlp_accuracy_delay_wgrad_compute(
         fuse_wgrad_accumulation=fuse_wgrad_accumulation,
     ).eval()
 
-    ln_mlp_ref = SelectiveLayerNormMLP(
+    ln_mlp_ref = LayerNormMLP(
         hidden_size=config.hidden_size,
         ffn_hidden_size=4 * config.hidden_size,
         eps=config.eps,

tests/pytorch/selective_layernorm_mlp/test_recipe.py

@@ -8,7 +8,7 @@
 import transformer_engine.pytorch as te
 from transformer_engine.pytorch import (
     quantized_model_init,
-    SelectiveLayerNormMLP,
+    LayerNormMLP,
 )
 
 from transformer_engine.common.recipe import DelayedScaling
@@ -35,7 +35,7 @@ def setup_class(cls) -> None:
 
     @pytest.mark.parametrize(
         "module_class",
-        [SelectiveLayerNormMLP],
+        [LayerNormMLP],
     )
     @pytest.mark.parametrize("checkpoint", (True, False))
     def test_quantizer_update(self, module_class, checkpoint):

tests/pytorch/selective_layernorm_mlp/test_sanity.py

@@ -15,7 +15,7 @@
 )
 from transformer_engine.pytorch import (
     autocast,
-    SelectiveLayerNormMLP,
+    LayerNormMLP,
     is_bf16_available,
 )
 from transformer_engine.common import recipe
@@ -158,7 +158,7 @@ def _test_sanity_common(
 @pytest.mark.parametrize("normalization", all_normalizations)
 @pytest.mark.parametrize("microbatching", all_boolean)
 @pytest.mark.parametrize("checkpoint", all_boolean)
-def test_sanity_selective_layernorm_mlp(
+def test_sanity_layernorm_mlp(
     dtype,
     fp8_recipe,
     model,
@@ -182,7 +182,7 @@ def test_sanity_selective_layernorm_mlp(
     init_method = init_method_normal(sigma)
     output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
 
-    block = SelectiveLayerNormMLP(
+    block = LayerNormMLP(
         config.hidden_size,
         4 * config.hidden_size,
         init_method=init_method,