clean up tests, remove unused imports

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

Author: jaimec00

tests/pytorch/selective_layernorm_mlp/distributed/run_numerics.py

@@ -9,14 +9,12 @@
 import os
 import sys
 from functools import wraps
-import math
 
 import torch
 from torch import nn
 import torch.distributed as dist
 import transformer_engine.pytorch as te
 
-import transformer_engine_torch as tex
 from transformer_engine.common.recipe import (
     MXFP8BlockScaling,
     DelayedScaling,
@@ -27,9 +25,6 @@
     Recipe,
     QParams,
 )
-from transformer_engine.pytorch import Float8CurrentScalingQuantizer, NVFP4Quantizer
-from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE
-from transformer_engine.pytorch.distributed import gather_along_first_dim
 
 
 def _compare_tensors(name, test, ref, rtol, atol):

tests/pytorch/selective_layernorm_mlp/test_cuda_graphs.py

@@ -166,12 +166,7 @@ def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor:
 
 
 # Supported modules
-_test_cuda_graphs_modules: List[str] = [
-    # Put linear first to test the case where the cuda context might not be set in
-    # creating TMA descriptor for MXFP8 quantization.
-    "selective_layernorm_mlp",
-]
-
+_test_cuda_graphs_modules: List[str] = ["selective_layernorm_mlp"]
 
 def _test_cuda_graphs(
     *,
@@ -195,20 +190,8 @@ def _test_cuda_graphs(
 
     # Create modules.
     with quantized_model_init(enabled=fp8_params, recipe=fp8_recipe):
-        if module == "transformer":
-            modules = [
-                TransformerLayer(
-                    model_config.hidden_size,
-                    model_config.hidden_size,
-                    model_config.num_heads,
-                    hidden_dropout=0.0,
-                    attention_dropout=0.0,
-                    fuse_qkv_params=True,
-                    params_dtype=dtype,
-                )
-                for _ in range(num_layers)
-            ]
-        elif module == "selective_layernorm_mlp":
+
+        if module == "selective_layernorm_mlp":
             modules = [
                 SelectiveLayerNormMLP(
                     model_config.hidden_size,
@@ -217,47 +200,7 @@ def _test_cuda_graphs(
                 )
                 for _ in range(num_layers)
             ]
-        elif module == "layernorm_linear":
-            modules = [
-                LayerNormLinear(
-                    model_config.hidden_size,
-                    model_config.hidden_size,
-                    params_dtype=dtype,
-                )
-                for _ in range(num_layers)
-            ]
-        elif module == "mha":
-            modules = [
-                MultiheadAttention(
-                    model_config.hidden_size,
-                    model_config.num_heads,
-                    attention_dropout=0.0,
-                    params_dtype=dtype,
-                    fuse_qkv_params=True,
-                )
-                for _ in range(num_layers)
-            ]
-        elif module == "linear":
-            modules = [
-                Linear(
-                    model_config.hidden_size,
-                    model_config.hidden_size,
-                    device="cuda",
-                    params_dtype=dtype,
-                )
-                for _ in range(num_layers)
-            ]
-        elif module == "linear_op":
-            modules = [
-                te_ops.Sequential(
-                    te_ops.Linear(
-                        model_config.hidden_size,
-                        model_config.hidden_size,
-                        dtype=dtype,
-                    ),
-                )
-                for _ in range(num_layers)
-            ]
+        
         else:
             raise ValueError(f"Unknown module type ({module})")
 

tests/pytorch/selective_layernorm_mlp/test_deferred_init.py

@@ -7,10 +7,7 @@
 
 import transformer_engine.pytorch as te
 
-_core_modules = [
-    te.SelectiveLayerNormMLP,
-]
-
+_core_modules = [te.SelectiveLayerNormMLP,]
 _composed_modules = []
 
 batch_size = 32
@@ -19,7 +16,6 @@
 head_dim = 64
 dtype = torch.bfloat16
 
-
 class TestDeferredInit:
 
     @staticmethod

tests/pytorch/selective_layernorm_mlp/test_numerics.py

@@ -2,7 +2,6 @@
 #
 # See LICENSE for license information.
 
-import math
 import os
 
 os.environ.setdefault("NVIDIA_TF32_OVERRIDE", "0")
@@ -12,41 +11,25 @@
 os.environ.setdefault("PYTORCH_CUDNN_ALLOW_TF32", "0")
 os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
 
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Tuple
 import pytest
-import random
 
 import torch
 import torch.nn as nn
 from torch.nn import Parameter
 
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager
-from transformer_engine.pytorch.utils import (
-    init_method_normal,
-    scaled_init_method_normal,
-    attention_mask_func,
-)
 from transformer_engine.pytorch import (
     autocast,
-    quantized_model_init,
     SelectiveLayerNormMLP,
-    Fp8Padding,
-    Fp8Unpadding,
-    Float8Quantizer,
-    Float8CurrentScalingQuantizer,
-    MXFP8Quantizer,
     get_device_compute_capability,
     is_fp8_available,
     is_mxfp8_available,
     is_fp8_block_scaling_available,
     is_bf16_available,
     is_nvfp4_available,
 )
-from transformer_engine.pytorch import checkpoint as te_checkpoint
-from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm
-from transformer_engine.pytorch.module.base import get_multi_stream_cublas_workspace, get_workspace
 from transformer_engine.common import recipe
-import transformer_engine_torch as tex
 from utils import ModelConfig, reset_rng_states
 
 

tests/pytorch/selective_layernorm_mlp/test_recipe.py

@@ -2,30 +2,16 @@
 #
 # See LICENSE for license information.
 
-from typing import Optional
-
 import pytest
 import torch
-import warnings
 
-import transformer_engine.common.recipe
 import transformer_engine.pytorch as te
 from transformer_engine.pytorch import (
-    Float8BlockQuantizer,
-    MXFP8Quantizer,
-    Float8Quantizer,
-    NVFP4Quantizer,
     quantized_model_init,
     SelectiveLayerNormMLP,
 )
 
-import transformer_engine_torch as tex
-from transformer_engine.pytorch.quantization import (
-    FP8GlobalStateManager,
-    _amax_and_scale_update,
-)
-import transformer_engine.pytorch.ops as te_ops
-from transformer_engine.common.recipe import DelayedScaling, Float8BlockScaling, MXFP8BlockScaling
+from transformer_engine.common.recipe import DelayedScaling
 
 # Check if FP8 is supported
 fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)

tests/pytorch/selective_layernorm_mlp/test_sanity.py

@@ -2,8 +2,6 @@
 #
 # See LICENSE for license information.
 
-from typing import Optional
-
 import torch
 import pytest
 import os
@@ -17,21 +15,10 @@
 )
 from transformer_engine.pytorch import (
     autocast,
-    quantized_model_init,
     SelectiveLayerNormMLP,
-    Float8CurrentScalingQuantizer,
-    Float8Quantizer,
-    Float8Tensor,
-    MXFP8Tensor,
-    checkpoint,
-    QuantizedTensor,
     is_bf16_available,
 )
 from transformer_engine.common import recipe
-import transformer_engine_torch as tex
-from transformer_engine.pytorch.cpp_extensions import general_gemm
-from transformer_engine.pytorch.module.base import get_workspace
-from transformer_engine.pytorch.tensor.utils import replace_raw_data
 from utils import ModelConfig
 
 # Only run FP8 tests on supported devices.