[Cherry-Pick] change rms norm for glm #7269 (#7275)

* change rms_norm

* refine code

* refine code

* refine code

* refine code

Author: zhangbo9674

fastdeploy/envs.py

@@ -197,6 +197,8 @@
     "FD_USE_PHI_MOE_TOPK": lambda: bool(int(os.getenv("FD_USE_PHI_MOE_TOPK", "0"))),
     # Whether to use phi MOE permute,if 1,use paddle op.
     "FD_USE_PHI_MOE_PERMUTE": lambda: bool(int(os.getenv("FD_USE_PHI_MOE_PERMUTE", "0"))),
+    # Whether to use phi rms_norm,if 1,use paddle op.
+    "FD_USE_PHI_RMSNORM": lambda: bool(int(os.getenv("FD_USE_PHI_RMSNORM", "0"))),
     # Control class SiluAndMul to use swiglu or fusid_bias_act operator in the forward_cuda function
     "FD_SiluAndMul_USE_PHI_SWIGLU": lambda: bool(int(os.getenv("FD_SiluAndMul_USE_PHI_SWIGLU", "0"))),
     # Reserve output blocks for decoding requests when schedule new prefill requests

fastdeploy/model_executor/models/glm4_moe.py

@@ -24,6 +24,7 @@
 from paddleformers.transformers import PretrainedModel
 from paddleformers.utils.log import logger
 
+import fastdeploy
 from fastdeploy.config import FDConfig
 from fastdeploy.model_executor.forward_meta import ForwardMeta
 from fastdeploy.model_executor.graph_optimization.decorator import (
@@ -252,6 +253,14 @@ def forward(
         return output
 
 
+def rms_norm_func(x, weight, eps):
+    rms_norm_out = paddle.nn.functional.rms_norm(x, x.shape[-1:], weight, eps)
+    if isinstance(rms_norm_out, (tuple, list)):
+        return rms_norm_out[0].astype(weight.dtype)
+    else:
+        return rms_norm_out.astype(weight.dtype)
+
+
 class Glm4MoeDecoderLayer(nn.Layer):
     """ """
 
@@ -305,8 +314,9 @@ def forward(
         residual: paddle.Tensor = None,
     ):
         """ """
+        proxy_rmsnorm = rms_norm_func if fastdeploy.envs.FD_USE_PHI_RMSNORM else None
         hidden_states, residual = self.input_layernorm(
-            hidden_states, residual_input=residual, forward_meta=forward_meta
+            hidden_states, residual_input=residual, forward_meta=forward_meta, proxy_rmsnorm=proxy_rmsnorm
         )
 
         hidden_states = self.self_attn(
@@ -315,7 +325,7 @@ def forward(
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual, proxy_rmsnorm=proxy_rmsnorm)
 
         hidden_states = self.mlp(hidden_states, forward_meta)