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,16 +88,18 @@ 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)
-
-    return normalization_func(
-        *inputs,
-        eps,
-        ln_out,
-        output_quantizer,
-        TE_DType[output_dtype] if output_dtype in TE_DType else output_dtype,
-        fwd_ln_sm_margin,
-        zero_centered_gamma,
-    )
+    if enable_lightop and (ln_bias is None):
+        return rmsnorm_forward(inputmat, ln_weight,ln_out,eps,True)
+    else:
+        return normalization_func(
+            *inputs,
+            eps,
+            ln_out,
+            output_quantizer,
+            TE_DType[output_dtype] if output_dtype in TE_DType else output_dtype,
+            fwd_ln_sm_margin,
+            zero_centered_gamma,
+        )
 
 
 class _NoopCatFunc(torch.autograd.Function):

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,14 +764,18 @@ class _LayerNormLinear(torch.autograd.Function):
                 )
                 dgrad = dgrad.reshape(inputmat.size())
             elif ctx.normalization == "RMSNorm":
-                dgrad, dgamma = tex.rmsnorm_bwd(
-                    dgrad,
-                    inputmat,
-                    rsigma,
-                    ln_weight,
-                    ctx.bwd_ln_sm_margin,
-                    ctx.zero_centered_gamma,
-                )
+                if enable_lightop:
+                    dgrad, dgamma =rmsnorm_backward(dgrad,inputmat,rsigma,ln_weight)
+                else:
+
+                    dgrad, dgamma = tex.rmsnorm_bwd(
+                        dgrad,
+                        inputmat,
+                        rsigma,
+                        ln_weight,
+                        ctx.bwd_ln_sm_margin,
+                        ctx.zero_centered_gamma,
+                    )
                 dgrad = dgrad.reshape(inputmat.size())
                 dbeta = None
             nvtx_range_pop(f"{nvtx_label}.norm")