feat(turbo): Add turbo RMSNorm patch (#263)

Author: ChengYao-amd

primus/backends/megatron/core/extensions/primus_turbo.py

@@ -1068,3 +1068,19 @@ def combine_postprocess(self, hidden_states: torch.Tensor):
         """
         hidden_states = self.deepep_dispatcher._post_combine(hidden_states)
         return hidden_states.view(self.hidden_shape)
+
+
+class PrimusTurboRMSNorm(te.pytorch.RMSNorm):
+    def __init__(self, *args, **kwargs):
+        assert "device" in kwargs
+        assert "dtype" in kwargs or "params_dtype" in kwargs, "device and dtype must be provided"
+        super().__init__(*args, **kwargs)
+        self.rms_norm_func = pt.modules.RMSNorm(
+            normalized_shape=kwargs["hidden_size"],
+            eps=self.eps,
+            device=kwargs["device"],
+            dtype=kwargs["dtype"] if "dtype" in kwargs else kwargs["params_dtype"],
+        )
+
+    def forward(self, x):
+        return self.rms_norm_func(x)

primus/configs/modules/megatron/primus_turbo.yaml

@@ -24,3 +24,6 @@ grouped_gemm_backend: "turbo-gg" # turbo-gg, lagacy-gg
 
 # use turbo fused activation_with_probs to optmize redundant computation
 use_turbo_fused_act_with_probs: false
+
+# layer norm
+use_turbo_rms_norm: false

primus/modules/trainer/megatron/trainer.py

@@ -192,6 +192,7 @@ def patch_pt_replace_te(self, args):
         from primus.backends.megatron.core.extensions.primus_turbo import (
             PrimusTurboColumnParallelLinearTorch,
             PrimusTurboDeepEPTokenDispatcher,
+            PrimusTurboRMSNorm,
         )
         from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import (
             PrimusTurboSpecProvider,
@@ -223,6 +224,11 @@ def patch_pt_replace_te(self, args):
             token_dispatcher.MoEFlexTokenDispatcher = PrimusTurboDeepEPTokenDispatcher
             moe_layer.MoEFlexTokenDispatcher = PrimusTurboDeepEPTokenDispatcher
 
+        if args.use_turbo_rms_norm:
+            import transformer_engine as te
+
+            te.pytorch.RMSNorm = PrimusTurboRMSNorm
+
     def patch_fp8_context(self):
         from megatron.core import fp8_utils
         from megatron.core.ssm import mamba_block