[SOT] Mark fleet func unified && fix circular import (#73278)

DrRyanHuang · web-flow · commit 16ed67835625 · 2025-06-14T19:36:45.000+08:00
diff --git a/python/paddle/distributed/fleet/layers/mpu/mp_ops.py b/python/paddle/distributed/fleet/layers/mpu/mp_ops.py
@@ -31,6 +31,7 @@
 from paddle.nn import Layer
 from paddle.nn.utils import dygraph_utils
 
+from .....jit.marker import unified
 from ....communication.reduce import ReduceOp, _get_reduce_op
 
 if TYPE_CHECKING:
@@ -74,6 +75,7 @@ def backward(ctx, dy):
         return out
 
 
+@unified
 def _c_identity(tensor, group=None, skip_c_identity_dynamic=False):
     """
     Return a copy of the tensor, mainly used with model parallel.
@@ -119,6 +121,7 @@ def _c_identity(tensor, group=None, skip_c_identity_dynamic=False):
         return out
 
 
+@unified
 def _c_concat(tensor, group=None):
     """
     Return allgather of the tensor, mainly used with model parallel.
@@ -169,6 +172,7 @@ def _c_concat(tensor, group=None):
         return out
 
 
+@unified
 def _c_split(tensor, group=None):
     """
     Split tensor evenly among all members, mainly used with model parallel.
@@ -256,6 +260,7 @@ def backward(ctx, dy):
             return _C_ops.c_identity(dy, ctx.ring_id, True, True)
 
 
+@unified
 def _mp_allreduce(
     tensor,
     op=ReduceOp.SUM,
@@ -307,6 +312,7 @@ def _mp_allreduce(
         return out
 
 
+@unified
 def _c_lookup_table(table, index, start_index=0, vocab_size=-1, name=None):
     """
     Lookup table according to index.
@@ -382,6 +388,7 @@ def extra_repr(self):
         return f'in_features={self.weight.shape[0]}, out_features={self.weight.shape[1]}, dtype={self._dtype}{name_str}'
 
 
+@unified
 def _c_softmax_with_cross_entropy(
     logits,
     label,
@@ -449,6 +456,7 @@ def _c_softmax_with_cross_entropy(
         return loss
 
 
+@unified
 def _c_softmax_with_multi_label_cross_entropy(
     logits,
     label,
@@ -537,6 +545,7 @@ def _c_softmax_with_multi_label_cross_entropy(
         return loss
 
 
+@unified
 def _linear(x, weight, bias=None, name=None):
     """
     Function Linear
diff --git a/python/paddle/jit/dy2static/function_spec.py b/python/paddle/jit/dy2static/function_spec.py
@@ -23,9 +23,6 @@
 from paddle.base import core
 from paddle.base.data_feeder import convert_dtype
 from paddle.base.dygraph.base import switch_to_static_graph
-from paddle.distributed.auto_parallel.placement_type import (
-    to_placements,
-)
 from paddle.jit.pir_translated_layer import PirTranslatedLayer
 from paddle.jit.translated_layer import TranslatedLayer
 from paddle.nn.layer import layers
@@ -177,6 +174,10 @@ def pir_to_static_inputs_with_spec(self, input_with_spec, main_program):
             input_with_spec(tuple): input arguments by replacing argument with InputSpec.
             main_program(Program): main program for inserting feed layer.
         """
+        from paddle.distributed.auto_parallel.placement_type import (
+            to_placements,
+        )
+
         flat_input_spec = paddle.utils.flatten(input_with_spec)
 
         # NOTE(zhangbo): Why do we need function_args and program_inputs: The primary function of this module is to construct the corresponding DataOp based on the inputSpec. The output of DataOp serves as the input to the Program and will also become the arguments for subsequent Static functions to construct the static graph. When the input is a DistributedInputSpec, the shard_tensor operation will be performed on the output of DataOp to obtain the corresponding distributed Value. The input to the Program will still be the output of DataOp, but in this case, the arguments for the Static functions will be the output of shard_tensor.
@@ -366,6 +367,9 @@ def get_buffers(layer_instance, include_sublayer=True):
 
 
 def _replace_value_with_input_spec(args):
+    from paddle.distributed.auto_parallel.placement_type import (
+        to_placements,
+    )
     from paddle.distributed.auto_parallel.static.dist_input_spec import (
         DistributedInputSpec,
     )
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -52,7 +52,6 @@
 from ..base.framework import Parameter
 from ..base.layer_helper import LayerHelper, LayerHelperBase
 from ..base.log_helper import get_logger
-from .fusion_utils import FusionStorage
 from .lr import LambdaDecay, LRScheduler
 
 if TYPE_CHECKING:
@@ -359,6 +358,8 @@ def set_merged_model_params(self, merged_model_params):
 
     @imperative_base.no_grad()
     def _maybe_refuse(self):
+        from .fusion_utils import FusionStorage
+
         # only support dygraph mode
         if not framework.in_dygraph_mode():
             return
