add lightop rmsnorm

transformer_engine/pytorch/module/_common.py

@@ -15,6 +15,13 @@ from .. import cpp_extensions as tex
 from ..constants import TE_DType
 from ..utils import get_default_init_method
 from ..tensor.float8_tensor import Float8Tensor
+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)
 
 
 def _get_normalization_func(normalization: str, forward: bool):
@@ -81,7 +88,9 @@ def apply_normalization(
     normalization_func = _get_normalization_func(normalization, True)
 
     inputs = (inputmat, ln_weight) if ln_bias is None else (inputmat, ln_weight, ln_bias)
-
+    if enable_lightop and (ln_bias is None):
+        return rmsnorm_forward(inputmat, ln_weight,ln_out,eps,True)
+    else:
         return normalization_func(
             *inputs,
             eps,

transformer_engine/pytorch/module/layernorm_linear.py

@@ -61,6 +61,13 @@ from ..cpu_offload import is_cpu_offload_enabled, set_offloading_param
 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)
 
 __all__ = ["LayerNormLinear"]
 
@@ -757,6 +764,10 @@ class _LayerNormLinear(torch.autograd.Function):
                 )
                 dgrad = dgrad.reshape(inputmat.size())
             elif ctx.normalization == "RMSNorm":
+                if enable_lightop:
+                    dgrad, dgamma =rmsnorm_backward(dgrad,inputmat,rsigma,ln_weight)
+                else:
+
                     dgrad, dgamma = tex.rmsnorm_bwd(
                         dgrad,
                         inputmat,