fix norm outputs in dynamic and static mode

Author: Wangzheee

paddlenlp/experimental/transformers/fused_transformer_layers.py

@@ -741,7 +741,7 @@ def get_weight_create_dype(self):
 
     def compute_layernorm_before_qkv(self, src, i):
         if i == 0:
-            ln_out = self.norm_func(src, self.ln_scales[i], self.ln_biases[i], self._epsilon, begin_norm_axis=1)
+            ln_out = self.norm_func(src, self.ln_scales[i], self.ln_biases[i], self._epsilon, begin_norm_axis=1)[0]
         else:
             ln_out = src
 
@@ -1918,7 +1918,7 @@ def compute_layernorm_before_qkv(self, src, i):
                 quant_round_type=self.quant_round_type,
                 quant_max_bound=self.quant_max_bound,
                 quant_min_bound=self.quant_min_bound,
-            )
+            )[0]
         else:
             ln_out = src
 
@@ -2617,9 +2617,7 @@ def compute_layernorm_before_qkv(self, src, i):
                 quant_round_type=1,
                 quant_max_bound=self.config.quant_max_bound,
                 quant_min_bound=self.config.quant_min_bound,
-            )
-            if in_dynamic_mode():
-                ln_out = ln_out[0]
+            )[0]
         else:
             ln_out = src
 

paddlenlp/experimental/transformers/llama/modeling.py

@@ -99,12 +99,9 @@ def __init__(self, config):
         self.config = config
 
     def forward(self, hidden_states):
-        result = paddle.incubate.nn.functional.fused_rms_norm(
+        return paddle.incubate.nn.functional.fused_rms_norm(
             hidden_states, self.weight, None, self.variance_epsilon, begin_norm_axis=1
-        )
-        if isinstance(result, tuple):
-            return result[0]
-        return result
+        )[0]
 
 
 class LLamaAvxLMHead(nn.Layer):

paddlenlp/experimental/transformers/mixtral/modeling.py

@@ -77,12 +77,9 @@ def __init__(self, config):
         self.config = config
 
     def forward(self, hidden_states):
-        result = paddle.incubate.nn.functional.fused_rms_norm(
+        return paddle.incubate.nn.functional.fused_rms_norm(
             hidden_states, self.weight, None, self.variance_epsilon, begin_norm_axis=1
-        )
-        if isinstance(result, tuple):
-            return result[0]
-        return result
+        )[0]
 
 
 @register_base_model

paddlenlp/experimental/transformers/qwen/modeling.py

@@ -53,10 +53,7 @@ def __init__(self, config):
         )
 
     def forward(self, x):
-        result = paddle.incubate.nn.functional.fused_rms_norm(x, self.weight, None, self.eps, begin_norm_axis=1)
-        if isinstance(result, tuple):
-            return result[0]
-        return result
+        return paddle.incubate.nn.functional.fused_rms_norm(x, self.weight, None, self.eps, begin_norm_axis=1)[0]
 
 
 @register_base_model

paddlenlp/experimental/transformers/qwen2/modeling.py

@@ -78,10 +78,7 @@ def __init__(self, config):
         )
 
     def forward(self, x):
-        result = paddle.incubate.nn.functional.fused_rms_norm(x, self.weight, None, self.eps, begin_norm_axis=1)
-        if isinstance(result, tuple):
-            return result[0]
-        return result
+        return paddle.incubate.nn.functional.fused_rms_norm(x, self.weight, None, self.eps, begin_norm_axis=1)[0]
 
 
 @register_base_model

paddlenlp/experimental/transformers/qwen2_moe/modeling.py

@@ -63,12 +63,8 @@ def __init__(self, config):
             dtype=paddle.get_default_dtype(),
             default_initializer=nn.initializer.Constant(1.0),
         )
-
     def forward(self, x):
-        result = paddle.incubate.nn.functional.fused_rms_norm(x, self.weight, None, self.eps, begin_norm_axis=1)
-        if isinstance(result, tuple):
-            return result[0]
-        return result
+        return paddle.incubate.nn.functional.fused_rms_norm(x, self.weight, None, self.eps, begin_norm_axis=1)[0]
 
 
 @register_base_model