Revert "[TSP] last_norm allgather move to model.py"

fastdeploy/model_executor/layers/normalization.py

@@ -105,14 +105,14 @@ def __init__(
         self.tp_rank = self.fd_config.parallel_config.tensor_parallel_rank
         self.tp_group = self.fd_config.parallel_config.tp_group
         is_input_norm = prefix.endswith(".input_layernorm")
-        self.is_last_norm = prefix.endswith(".norm")
+        is_last_norm = prefix.endswith(".norm")
         self.split_x = (
             self.fd_config.parallel_config.use_sequence_parallel_moe
             and self.layer_id == self.fd_config.model_config.moe_layer_start_index
             and is_input_norm
         )
         self.allgather_out = self.fd_config.parallel_config.use_sequence_parallel_moe and (
-            (self.layer_id > self.fd_config.model_config.moe_layer_start_index and is_input_norm)
+            (self.layer_id > self.fd_config.model_config.moe_layer_start_index and is_input_norm) or is_last_norm
         )
 
         self.init_weight()

fastdeploy/model_executor/models/deepseek_v3.py

@@ -591,9 +591,6 @@ def forward(
             )
         out = self.norm(hidden_states, residual, forward_meta=forward_meta)[0]
 
-        if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe:
-            out = self.norm.allgather(out, forward_meta.ids_remove_padding.shape[0])
-
         return out
 
 

fastdeploy/model_executor/models/ernie4_5_moe.py

@@ -449,9 +449,6 @@ def forward(
 
         out = self.norm(hidden_states, residual, forward_meta=forward_meta)[0]
 
-        if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe:
-            out = self.norm.allgather(out, forward_meta.ids_remove_padding.shape[0])
-
         if current_platform.is_iluvatar() and forward_meta.attn_backend.mixed:
             out = forward_meta.attn_backend.reverse_transpose(out)
 

fastdeploy/model_executor/models/ernie4_5_mtp.py

@@ -326,10 +326,7 @@ def forward(
         for i in range(self.num_layers):
             hidden_states, residual = self.mtp_block[i](forward_meta, hidden_states, residual)
 
-        hidden_states = self.norm(hidden_states, residual, forward_meta=forward_meta)[0]
-
-        if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe:
-            hidden_states = self.norm.allgather(hidden_states, forward_meta.ids_remove_padding.shape[0])
+        hidden_states = self.norm(hidden_states, residual)[0]
 
         return hidden_states
 
@@ -429,7 +426,7 @@ def load_weights(self, weights_iterator) -> None:
             )
             process_weights_after_loading_fn(model_sublayer_name, param)
 
-    def compute_logits(self, hidden_states: paddle.Tensor, forward_meta: ForwardMeta):
+    def compute_logits(self, hidden_states: paddle.Tensor):
         """
         compute logits
         """

fastdeploy/model_executor/models/ernie4_5_vl/ernie4_5_vl_moe.py

@@ -543,9 +543,6 @@ def forward(
 
         out = self.norm(hidden_states, residual, forward_meta=forward_meta)[0]
 
-        if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe:
-            out = self.norm.allgather(out, forward_meta.ids_remove_padding.shape[0])
-
         return out
 
 

fastdeploy/model_executor/models/glm4_moe.py

@@ -370,9 +370,6 @@ def forward(
 
         out = self.norm(hidden_states, residual, forward_meta=forward_meta)[0]
 
-        if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe:
-            out = self.norm.allgather(out, forward_meta.ids_remove_padding.shape[0])
-
         return out
 
 

fastdeploy/model_executor/models/gpt_oss.py

@@ -214,12 +214,8 @@ def forward(self, ids_remove_padding: paddle.Tensor, forward_meta: ForwardMeta):
         for i in range(self.num_layers):
             hidden_states, residual = self.layers[i](forward_meta, hidden_states, residual)
 
-        out = self.norm(hidden_states, residual, forward_meta=forward_meta)[0]
-
-        if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe:
-            out = self.norm.allgather(out, forward_meta.ids_remove_padding.shape[0])
-
-        return out
+        hidden_states = self.norm(hidden_states, residual)[0]
+        return hidden_states
 
 
 @ModelRegistry.register_model_class(

fastdeploy/model_executor/models/qwen3moe.py

@@ -284,9 +284,6 @@ def forward(
 
         out = self.norm(hidden_states, residual, forward_meta=forward_meta)[0]
 
-        if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe:
-            out = self.norm.allgather(out, forward_meta.ids_remove_padding.shape[0])
-
         return out
 
 

fastdeploy/spec_decode/mtp.py

@@ -886,7 +886,7 @@ def _propose(self, step_use_cudagraph: bool = False, is_dummy_run=False):
                 )
 
                 # 4. Compute logits, Sample
-                logits = self.model.compute_logits(hidden_states, forward_meta=self.forward_meta)
+                logits = self.model.compute_logits(hidden_states)
                 if self.enable_logprob and substep == 0:
                     first_token_logits = self.model.compute_logits(self.model_inputs["first_token_hidden_states"])