Add NVTE_USE_LIGHTOP env var to control lightop import

Signed-off-by: maxiao3 <maxiao3@sugon.com>

See merge request dcutoolkit/deeplearing/TransformerEngine!71

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -9,11 +9,13 @@ import os
 import torch
 import transformer_engine_torch as tex
 import warnings
-try:
-    import lightop
-    enable_lightop = True
-except ImportError:
-    enable_lightop = False
+enable_lightop = os.getenv("NVTE_USE_LIGHTOP", "false").strip().lower() in ["true", "1"]
+if enable_lightop:
+    try:
+        import lightop
+        enable_lightop = True
+    except ImportError:
+        enable_lightop = False
 from ..constants import TE_DType, TE_DType_To_Torch
 from ..utils import get_sm_count, _empty_tensor
 from transformer_engine.pytorch.triton.blockwise_int8_gemm_nt import w8a8_block_int8_matmul, w8a8_block_int8_matmul_batched

transformer_engine/pytorch/module/_common.py

@@ -16,12 +16,15 @@ from ..export import is_in_onnx_export_mode
 from ..utils import get_default_init_method
 
 import warnings
-try:
-    from lightop import rmsnorm_forward,rmsnorm_backward
-    enable_lightop = True
-except ImportError:
-    enable_lightop = False
-    warnings.warn("Failed to import lightop module. Falling back to alternative implementation.", UserWarning)
+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:
+        enable_lightop = False
+        warnings.warn("Failed to import lightop module. Falling back to alternative implementation.", UserWarning)
 
 def _get_normalization_func(normalization: str, forward: bool):
     fwd_normalization_funcs = {

transformer_engine/pytorch/module/layernorm_linear.py

@@ -76,12 +76,16 @@ from ..cpp_extensions import (
     general_gemm,
 )
 import warnings
-try:
-    from lightop import rmsnorm_forward,rmsnorm_backward
-    enable_lightop = True
-except ImportError:
-    enable_lightop = False
-    warnings.warn("Failed to import lightop module. Falling back to alternative implementation.", UserWarning)
+
+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:
+        enable_lightop = False
+        warnings.warn("Failed to import lightop module. Falling back to alternative implementation.", UserWarning)
 
 __all__ = ["LayerNormLinear"]
 

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -83,11 +83,13 @@ from ..cpp_extensions import (
 from ..export import is_in_onnx_export_mode, assert_warmed_up
 from ...debug.pytorch.debug_state import TEDebugState
 
-try:
-    from lightop import rmsnorm_forward, rmsnorm_backward
-    enable_lightop = True
-except ImportError:
-    enable_lightop = False
+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:
+        enable_lightop = False
 
 __all__ = ["LayerNormMLP"]