Merge branch 'main' of http://10.6.10.68/dcutoolkit/deeplearing/TransformerEngine

transformer_engine/pytorch/module/_common.py

@@ -88,8 +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)
+    if enable_lightop and (ln_bias is None) and normalization == "RMSNorm":
+        out, rsigma = rmsnorm_forward(inputmat, ln_weight,ln_out,eps,True)
+        return out, None, rsigma
     else:
         return normalization_func(
             *inputs,