TEv2 as default TE executor (#2510)

Author: riccardofelluga

thunder/__init__.py

@@ -574,8 +574,8 @@ def apply_transforms_and_build_cache_entry(cd, cs, cache_info, prologue_trc, com
             if requires_grad:
                 from thunder.transforms.autodiff import split_into_forward_and_backward
 
-                if "transformer_engine_v2" in {ex.name for ex in cd.executors_list}:
-                    from thunder.executors.transformer_engine_v2ex import _te_activation_checkpointing_transform
+                if "transformer_engine" in {ex.name for ex in cd.executors_list}:
+                    from thunder.executors.transformer_engineex import _te_activation_checkpointing_transform
 
                     computation_trc = _te_activation_checkpointing_transform(computation_trc)
 

thunder/benchmarks/__init__.py

@@ -27,7 +27,7 @@
 from thunder.executors.apexex import apex_ex, apex_entropy_available
 from thunder.executors.cudnn_layernormex import cudnn_layernorm_ex
 from thunder.executors.cudnnex import cudnn_ex, cudnn_available
-from thunder.executors.transformer_engineex import transformer_engine_ex, TE_AVAILABLE
+from thunder.executors.transformer_engineex import transformer_engine_ex, TransformerEngineTransform
 from thunder.executors.sdpaex import sdpa_ex
 from thunder.executors.torch_compile import torch_compile_cat_ex, torch_compile_ex
 from thunder.transforms.cudagraph import CUDAGraphTransform
@@ -764,10 +764,16 @@ def thunder_cudnn_layer_norm_nvfuser_executor(fn: Callable) -> Callable:
 
 thunder_transformerengine_executor: None | Callable = None
 
-if TE_AVAILABLE:
+if transformer_engine_ex is not None:
 
     def thunder_transformerengine_executor(fn: Callable):
-        return thunder.jit(fn, executors=(transformer_engine_ex,) + thunder.get_default_executors())
+        return thunder.jit(
+            fn,
+            executors=(transformer_engine_ex,) + thunder.get_default_executors(),
+            transforms=[
+                TransformerEngineTransform(),
+            ],
+        )
 
 
 def thunder_sdpa_executor(fn: Callable) -> Callable:

thunder/benchmarks/benchmark_litgpt.py

@@ -625,19 +625,19 @@ def setup_compile(self, model):
 
                 executors.insert(0, torch_compile_ex)
 
-            if "transformerengine_v2" in self.compile:
-                from thunder.executors.transformer_engine_v2ex import (
-                    transformer_engine_v2_ex,
-                    TransformerEngineTransformV2,
-                )
+            if "transformerengine_v1" in self.compile:
+                from thunder.executors.transformer_engine_v1ex import transformer_engine_v1_ex
 
-                executors.insert(0, transformer_engine_v2_ex)
-                transforms.insert(0, TransformerEngineTransformV2())
+                executors.insert(0, transformer_engine_v1_ex)
 
             elif "transformerengine" in self.compile:
-                from thunder.executors.transformer_engineex import transformer_engine_ex
+                from thunder.executors.transformer_engineex import (
+                    transformer_engine_ex,
+                    TransformerEngineTransform,
+                )
 
                 executors.insert(0, transformer_engine_ex)
+                transforms.insert(0, TransformerEngineTransform())
 
             if "dynamo" in self.compile:
                 if self.distributed_mode == "fsdp2":

thunder/core/trace.py

@@ -413,7 +413,7 @@ def keyfn(class_or_module: type | ModuleType) -> str:
                 # NOTE: For TE v1.6 onwards, `fp8_autocast` checks if `torch.is_grad_enabled` for updating
                 # the FP8 scales/inverses. So this decorator should be applied before `torch.no_grad` (so that
                 # it is in grad enabled part).
-                from thunder.executors.transformer_engineex import _is_te_linear_enabled, _get_te_wrapper_string
+                from thunder.executors.transformer_engine_v1ex import _is_te_linear_enabled, _get_te_wrapper_string
 
                 if TraceTag.AUGMENTED_FORWARD and _is_te_linear_enabled(import_ctx, object_ctx):
                     program.append(_get_te_wrapper_string())