[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,14 +1270,17 @@ class _LayerNormMLP(torch.autograd.Function):
                     ctx.zero_centered_gamma,
                 )
             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 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,
+                        rsigma,
+                        ln_weight,
+                        ctx.bwd_ln_sm_margin,
+                        ctx.zero_centered_gamma,
+                    )
                 dbeta = None
         clear_tensor_data(mu, rsigma)