[CP][BugFix]Dev fix custom ar unstable result (PaddlePaddle#4437)

Author: ckl117

fastdeploy/distributed/communication.py

@@ -53,7 +53,7 @@ def tensor_model_parallel_all_reduce(
         global _TP_AR
         if _TP_AR is not None and _TP_AR.should_custom_ar(input_):
             # TODO: supports different_group custom allreduce
-            _TP_AR.custom_all_reduce(input_)
+            input_ = _TP_AR.custom_all_reduce(input_)
         elif paddle.in_dynamic_mode():
             if group_ is not None:
                 dist.all_reduce(input_, group=group_)
@@ -63,6 +63,7 @@ def tensor_model_parallel_all_reduce(
                 dist.all_reduce(input_, group=mp_group)
         else:
             dist.all_reduce(input_)
+        return input_
 
 except:
     tensor_model_parallel_all_reduce = None

fastdeploy/distributed/custom_all_reduce/custom_all_reduce.py

@@ -212,13 +212,13 @@ def custom_all_reduce(self, input: paddle.Tensor) -> Optional[paddle.Tensor]:
             stream_capturing = lib.cudaStreamIsCapturing(stream)
             if stream_capturing.value == 1:
                 # 1 is cudaStreamCaptureStatusActive: The stream is capturing.
-                return self.all_reduce(input, input, registered=True)
+                return self.all_reduce(input, registered=True)
             else:
                 # If warm up, mimic the allocation pattern since custom
                 # allreduce is out-of-place.
                 return paddle.empty_like(input)
         else:
-            return self.all_reduce(input, input, registered=False)
+            return self.all_reduce(input, registered=False)
 
     def close(self):
         if self._ptr:

fastdeploy/model_executor/layers/backends/dcu/fused_moe_triton_backends.py

@@ -243,5 +243,5 @@ def apply(
         out = intermediate_cache3.sum(axis=1)
 
         if layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(out)
+            out = tensor_model_parallel_all_reduce(out)
         return out

fastdeploy/model_executor/layers/backends/gcu/moe/fused_moe_method_gcu_backend.py

@@ -180,7 +180,7 @@ def compute_ffn(
                 tensor_model_parallel_all_reduce,
             )
 
-            tensor_model_parallel_all_reduce(fused_moe_out)
+            fused_moe_out = tensor_model_parallel_all_reduce(fused_moe_out)
 
         return fused_moe_out
 

fastdeploy/model_executor/layers/backends/metax/moe/fused_moe_triton_metax_backend.py

@@ -778,7 +778,7 @@ def forward_cuda(self, x: paddle.Tensor) -> paddle.Tensor:
             out = paddle.matmul(x, self.weight)
 
         if self.reduce_results and self.nranks > 1:
-            tensor_model_parallel_all_reduce(out, self.tp_group)
+            out = tensor_model_parallel_all_reduce(out, self.tp_group)
         if not self.fd_config.quant_config and self.add_bias:
             out = paddle.add(out, self.bias)
         return out

fastdeploy/model_executor/layers/linear.py

@@ -298,7 +298,7 @@ def apply_tp(
         )
 
         if layer.reduce_results and layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(fused_moe_out)
+            fused_moe_out = tensor_model_parallel_all_reduce(fused_moe_out, layer.fd_config.parallel_config.tp_group)
 
         return fused_moe_out
 

fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py

@@ -575,6 +575,6 @@ def apply_tp(
             1.0,
         )[0]
         if layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(tmp_ffn_out)
+            tmp_ffn_out = tensor_model_parallel_all_reduce(tmp_ffn_out)
 
         return tmp_ffn_out

fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py

@@ -353,6 +353,6 @@ def apply(
         ffn_out = ffn_out.sum(axis=1)
 
         if layer.reduce_results and layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(ffn_out)
+            ffn_out = tensor_model_parallel_all_reduce(ffn_out)
 
         return ffn_out

fastdeploy/model_executor/layers/moe/fused_moe_marlin_backend.py

@@ -599,7 +599,7 @@ def apply(
         out = down_proj_out.sum(axis=1)
 
         if layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(out)
+            out = tensor_model_parallel_all_reduce(out)
 
         return out
 
@@ -997,6 +997,6 @@ def apply(
         out = intermediate_cache3.sum(axis=1)
 
         if layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(out)
+            out = tensor_model_parallel_all_reduce(out)
 
         return out