[fix] Remove TE v1 from low precision options (#2785)

Author: riccardofelluga

thunder/benchmarks/benchmark_litgpt.py

@@ -79,20 +79,18 @@ class LowPrecisionHandler:
     enabled: bool = False
 
     use_thunder_te: bool = True
-    use_legacy_thunder_te: bool = False
 
     @property
     def use_te_autocast(self) -> bool:
-        return self.enabled and not self.use_legacy_thunder_te
+        return self.enabled
 
     def check_and_add_compile_options(self, compile_options: str) -> str:
         if not self.enabled and "_transformerengine" in compile_options:
             raise ValueError("Low precision mode not specified but found transfomerengine in the compile options!")
 
         self.use_thunder_te = "thunder" in compile_options
-        self.use_legacy_thunder_te = "transformerengine_v1" in compile_options
 
-        if self.enabled and self.use_thunder_te and not self.use_legacy_thunder_te:
+        if self.enabled and self.use_thunder_te:
             compile_options += "_transformerengine"
         return compile_options
 
@@ -120,15 +118,13 @@ def __post_init__(self) -> None:
     def executor_str(self) -> str:
         if not self.enabled:
             return "low precision is not enabled"
-        elif self.use_legacy_thunder_te:
-            return "Thunder TE executor v1"
         elif self.use_thunder_te:
             return "Thunder TE executor"
         else:
             return "TransformerEngine without Thunder"
 
     def maybe_apply_te_autocast(self):
-        if self.enabled and not self.use_legacy_thunder_te:
+        if self.enabled:
             return te.fp8_autocast(fp8_recipe=self.fp8_recipe)
         else:
             return nullcontext()
@@ -160,7 +156,7 @@ def update_if_not_divisible(attr_name, divisor):
             print("No updates were necessary.")
 
     def swap_linear_layers_for_te(self, model: torch.nn.Module, device: Any) -> None:
-        if not self.enabled or self.use_thunder_te:
+        if not self.enabled:
             return
 
         swap_layernorm = self.mode == "fp8-default-te-wo_layernorm"