[DCU] fix

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -80,6 +80,12 @@ 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
+
 __all__ = ["LayerNormMLP"]
 
 
@@ -1264,6 +1270,9 @@ class _LayerNormMLP(torch.autograd.Function):
                     ctx.zero_centered_gamma,
                 )
             elif ctx.normalization == "RMSNorm":
+                if enable_lightop and (rsigma.dtype is torch.bfloat16 or rsigma.dtype is torch.float16):
+                    dgrad, dgamma = rmsnorm_backward(dgrad, inputmat, rsigma, ln_weight)
+                else:
                     dgrad, dgamma = tex.rmsnorm_bwd(
                         dgrad,
                         inputmat,