Merge branch 'develop_v2.10' into release_v2.10

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -10,10 +10,12 @@ import functools
 import torch
 import transformer_engine_torch as tex
 import warnings
-try:
+enable_lightop = os.getenv("NVTE_USE_LIGHTOP", "false").strip().lower() in ["true", "1"]
+if enable_lightop:
+    try:
         import lightop
         enable_lightop = True
-except ImportError:
+    except ImportError:
         enable_lightop = False
 from ..constants import TE_DType, TE_DType_To_Torch
 from ..utils import get_sm_count, _empty_tensor

transformer_engine/pytorch/module/_common.py

@@ -16,10 +16,13 @@ from ..export import is_in_onnx_export_mode
 from ..utils import get_default_init_method
 
 import warnings
-try:
+import os
+enable_lightop = os.getenv("NVTE_USE_LIGHTOP", "false").strip().lower() in ["true", "1"]
+if enable_lightop:
+    try:
         from lightop import rmsnorm_forward,rmsnorm_backward
         enable_lightop = True
-except ImportError:
+    except ImportError:
         enable_lightop = False
         warnings.warn("Failed to import lightop module. Falling back to alternative implementation.", UserWarning)
 

transformer_engine/pytorch/module/layernorm_linear.py

@@ -80,10 +80,14 @@ from ..cpp_extensions import (
     general_gemm,
 )
 import warnings
-try:
+
+enable_lightop = os.getenv("NVTE_USE_LIGHTOP", "false").strip().lower() in ["true", "1"]
+
+if enable_lightop:
+    try:
         from lightop import rmsnorm_forward,rmsnorm_backward
         enable_lightop = True
-except ImportError:
+    except ImportError:
         enable_lightop = False
         warnings.warn("Failed to import lightop module. Falling back to alternative implementation.", UserWarning)
 

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -88,10 +88,12 @@ from ..cpp_extensions import (
 from ..export import is_in_onnx_export_mode, assert_warmed_up
 from ...debug.pytorch.debug_state import TEDebugState
 
-try:
+enable_lightop = os.getenv("NVTE_USE_LIGHTOP", "false").strip().lower() in ["true", "1"]
+if enable_lightop:
+    try:
         from lightop import rmsnorm_forward, rmsnorm_backward
         enable_lightop = True
-except ImportError:
+    except ImportError:
         enable_lightop = False
 
 __all__ = ["LayerNormMLP"]