fix conflict (#44891)

python/paddle/distributed/auto_parallel/dist_op.py

@@ -26,6 +26,7 @@ from .dist_attribute import get_op_dist_attr_field_keys
 
 
 class DistributedOperator:
+
     def __init__(self, serial_op, dist_attr=None):
         self._serial_op = serial_op
         self._serial_inputs = {}
@@ -248,6 +249,7 @@ class DistributedOperator:
 
 
 class DistributedModule:
+
     def __init__(self, serial_module, dist_attr=None):
         self._serial_module = serial_module
         self._dist_attr = dist_attr
@@ -265,6 +267,4 @@ class DistributedModule:
             dist_op = DistributedOperator(op, self._dist_attr)
             dist_op.dist_attr.mark_annotated_as(self._dist_attr)
             default_dist_ctx.add_dist_op_for_program(dist_op)
-        if isinstance(output, Variable):
-            output = [output]
-        return list(output)
+        return output

python/paddle/fluid/tests/unittests/auto_parallel/engine_api.py

@@ -47,6 +47,7 @@ paddle.seed(44)
 
 
 class MyDataset(Dataset):
+
     def __init__(self, num_samples):
         super(MyDataset, self).__init__()
         self.num_samples = num_samples
@@ -61,6 +62,7 @@ class MyDataset(Dataset):
 
 
 class MLPLayer(nn.Layer):
+
     def __init__(self,
                  hidden_size=1024,
                  intermediate_size=4 * 1024,
@@ -69,43 +71,45 @@ class MLPLayer(nn.Layer):
         super(MLPLayer, self).__init__()
         d_model = hidden_size
         dim_feedforward = intermediate_size
-        weight_attr = paddle.ParamAttr(initializer=nn.initializer.Normal(
-            mean=0.0, std=initializer_range))
+        weight_attr = paddle.ParamAttr(
+            initializer=nn.initializer.Normal(mean=0.0, std=initializer_range))
         bias_attr = None
 
-        self.linear0 = nn.Linear(
-            d_model, dim_feedforward, weight_attr, bias_attr=bias_attr)
-        self.linear1 = nn.Linear(
-            dim_feedforward, d_model, weight_attr, bias_attr=bias_attr)
+        self.linear0 = nn.Linear(d_model,
+                                 dim_feedforward,
+                                 weight_attr,
+                                 bias_attr=bias_attr)
+        self.linear1 = nn.Linear(dim_feedforward,
+                                 d_model,
+                                 weight_attr,
+                                 bias_attr=bias_attr)
         self.linear2 = nn.Linear(d_model, 1, weight_attr, bias_attr=bias_attr)
         self.norm = nn.LayerNorm(d_model, epsilon=1e-5)
         self.dropout = nn.Dropout(dropout_ratio, mode="upscale_in_train")
 
     def forward(self, input):
-        out = auto.shard_op(
-            self.norm, dist_attr={"process_mesh": PP_MESH_0})(input)[0]
-        out = self.linear0(input)
+        out = auto.shard_op(self.norm, dist_attr={"process_mesh":
+                                                  PP_MESH_0})(input)
+        out = self.linear0(out)
         out = F.gelu(out, approximate=True)
-        out = auto.shard_op(
-            self.linear1, dist_attr={"process_mesh": PP_MESH_1})(out)[0]
+        out = auto.shard_op(self.linear1, dist_attr={"process_mesh":
+                                                     PP_MESH_1})(out)
         out = self.dropout(out)
         out = self.linear2(out)
         return out
 
 
 def train():
-    mlp = MLPLayer(
-        hidden_size=hidden_size,
-        intermediate_size=4 * hidden_size,
-        dropout_ratio=0.1,
-        initializer_range=0.02)
+    mlp = MLPLayer(hidden_size=hidden_size,
+                   intermediate_size=4 * hidden_size,
+                   dropout_ratio=0.1,
+                   initializer_range=0.02)
     loss = paddle.nn.CrossEntropyLoss()
-    optimizer = paddle.fluid.optimizer.AdamOptimizer(
-        learning_rate=0.00001,
-        beta1=0.9,
-        beta2=0.999,
-        epsilon=1e-08,
-        grad_clip=None)
+    optimizer = paddle.fluid.optimizer.AdamOptimizer(learning_rate=0.00001,
+                                                     beta1=0.9,
+                                                     beta2=0.999,
+                                                     epsilon=1e-08,
+                                                     grad_clip=None)
 
     dataset = MyDataset(batch_num * batch_size)
     inputs_spec = InputSpec([batch_size, hidden_size], 'float32', 'x')
@@ -119,11 +123,10 @@ def train():
     dist_strategy.semi_auto = True
     fleet.init(is_collective=True, strategy=dist_strategy)
 
-    engine = Engine(
-        mlp,
-        inputs_spec=inputs_spec,
-        labels_spec=labels_spec,
-        strategy=dist_strategy)
+    engine = Engine(mlp,
+                    inputs_spec=inputs_spec,
+                    labels_spec=labels_spec,
+                    strategy=dist_strategy)
     engine.prepare(optimizer, loss)
     engine.fit(dataset,
                batch_size=batch_size,

python/paddle/fluid/tests/unittests/test_auto_parallel_reshard_mppp.py

@@ -38,6 +38,7 @@ PP_MESH_1 = auto.ProcessMesh([2, 3])
 
 
 class MLPLayer(nn.Layer):
+
     def __init__(self,
                  hidden_size=1024,
                  intermediate_size=4 * 1024,
@@ -45,42 +46,51 @@ class MLPLayer(nn.Layer):
         super(MLPLayer, self).__init__()
         d_model = hidden_size
         dim_feedforward = intermediate_size
-        weight_attr = paddle.ParamAttr(initializer=nn.initializer.Normal(
-            mean=0.0, std=initializer_range))
+        weight_attr = paddle.ParamAttr(
+            initializer=nn.initializer.Normal(mean=0.0, std=initializer_range))
         bias_attr = None
 
         self.word_embeddings = nn.Embedding(
             hidden_size,
             hidden_size,
-            weight_attr=paddle.ParamAttr(
-                name="word_embeddings",
-                initializer=nn.initializer.Normal(
-                    mean=0.0, std=initializer_range)))
-
-        self.linear0 = nn.Linear(
-            d_model, dim_feedforward, weight_attr, bias_attr=bias_attr)
-        self.linear1 = nn.Linear(
-            dim_feedforward, d_model, weight_attr, bias_attr=bias_attr)
-        self.linear2 = nn.Linear(
-            dim_feedforward, d_model, weight_attr, bias_attr=bias_attr)
+            weight_attr=paddle.ParamAttr(name="word_embeddings",
+                                         initializer=nn.initializer.Normal(
+                                             mean=0.0, std=initializer_range)))
+
+        self.linear0 = nn.Linear(d_model,
+                                 dim_feedforward,
+                                 weight_attr,
+                                 bias_attr=bias_attr)
+        self.linear1 = nn.Linear(dim_feedforward,
+                                 d_model,
+                                 weight_attr,
+                                 bias_attr=bias_attr)
+        self.linear2 = nn.Linear(dim_feedforward,
+                                 d_model,
+                                 weight_attr,
+                                 bias_attr=bias_attr)
 
     def forward(self, input):
-        auto.shard_tensor(
-            self.word_embeddings.weight,
-            dist_attr={"process_mesh": PP_MESH_0,
-                       "dims_mapping": [0, -1]})
-        auto.shard_tensor(
-            self.linear0.weight,
-            dist_attr={"process_mesh": PP_MESH_0,
-                       "dims_mapping": [-1, 0]})
-        auto.shard_tensor(
-            self.linear1.weight,
-            dist_attr={"process_mesh": PP_MESH_1,
-                       "dims_mapping": [0, -1]})
-        auto.shard_tensor(
-            self.linear2.weight,
-            dist_attr={"process_mesh": PP_MESH_1,
-                       "dims_mapping": [0, -1]})
+        auto.shard_tensor(self.word_embeddings.weight,
+                          dist_attr={
+                              "process_mesh": PP_MESH_0,
+                              "dims_mapping": [0, -1]
+                          })
+        auto.shard_tensor(self.linear0.weight,
+                          dist_attr={
+                              "process_mesh": PP_MESH_0,
+                              "dims_mapping": [-1, 0]
+                          })
+        auto.shard_tensor(self.linear1.weight,
+                          dist_attr={
+                              "process_mesh": PP_MESH_1,
+                              "dims_mapping": [0, -1]
+                          })
+        auto.shard_tensor(self.linear2.weight,
+                          dist_attr={
+                              "process_mesh": PP_MESH_1,
+                              "dims_mapping": [0, -1]
+                          })
         w_out = self.word_embeddings(input)
         out = self.linear0(w_out)
         gelu_out = F.gelu(out, approximate=True)
@@ -98,21 +108,24 @@ def mlp_forward(train_program, start_program):
         hidden_size = 1024
         sequence_len = 512
         input = static.data(name="input", shape=[batch_size], dtype='int32')
-        label = static.data(
-            name="label", shape=[batch_size, 1], dtype='float32')
-
-        auto.shard_tensor(
-            input, dist_attr={"process_mesh": PP_MESH_0,
-                              "dims_mapping": [-1]})
-        auto.shard_tensor(
-            label,
-            dist_attr={"process_mesh": PP_MESH_1,
-                       "dims_mapping": [-1, -1]})
-
-        mlp = MLPLayer(
-            hidden_size=hidden_size,
-            intermediate_size=4 * hidden_size,
-            initializer_range=0.02)
+        label = static.data(name="label",
+                            shape=[batch_size, 1],
+                            dtype='float32')
+
+        auto.shard_tensor(input,
+                          dist_attr={
+                              "process_mesh": PP_MESH_0,
+                              "dims_mapping": [-1]
+                          })
+        auto.shard_tensor(label,
+                          dist_attr={
+                              "process_mesh": PP_MESH_1,
+                              "dims_mapping": [-1, -1]
+                          })
+
+        mlp = MLPLayer(hidden_size=hidden_size,
+                       intermediate_size=4 * hidden_size,
+                       initializer_range=0.02)
 
         predict = mlp(input)
         error_cost = paddle.nn.functional.square_error_cost(predict, label)
@@ -137,13 +150,12 @@ def get_dist_prog(train_program, startup_program, dist_context, rank_id):
     complete_train_program = completer.complete_forward_annotation(
         train_program)
     dist_context.block_state.parse_forward_blocks(complete_train_program)
-    params_grads = parallelizer._generate_backward(
-        complete_train_program,
-        startup_program,
-        loss,
-        parameter_list=None,
-        no_grad_set=None,
-        callbacks=None)
+    params_grads = parallelizer._generate_backward(complete_train_program,
+                                                   startup_program,
+                                                   loss,
+                                                   parameter_list=None,
+                                                   no_grad_set=None,
+                                                   callbacks=None)
 
     # logical partition
     partitioner = Partitioner(dist_context, rank_id)
@@ -171,8 +183,7 @@ def check_send_recv_result(dist_main_prog, rank_id):
             if op.type == "send_v2" and "gelu_0.tmp_0@GRAD" in op.input_arg_names[
                     0]:
                 send_result = True
-            if op.type == "recv_v2" and "gelu_0.tmp_0" in op.output_arg_names[
-                    0]:
+            if op.type == "recv_v2" and "gelu_0.tmp_0" in op.output_arg_names[0]:
                 recv_result = True
 
     return send_result and recv_result
@@ -206,6 +217,7 @@ def check_allgather(dist_main_program):
 
 
 class TestMLPReshard(unittest.TestCase):
+
     def test_mlp_mppp(self):
         train_program = paddle.static.Program()
         startup_program = paddle.static.Program()
@@ -230,38 +242,29 @@ class TestMLPReshard(unittest.TestCase):
         process_mesh = auto.ProcessMesh(mesh=[0, 3])
         with static.program_guard(train_program, startup_program):
             x = paddle.static.data(name="x", shape=[4, 4], dtype='float32')
-            x = auto.shard_tensor(
-                x,
-                dist_attr={
-                    "process_mesh": process_mesh,
-                    "dims_mapping": [0, -1]
-                })
+            x = auto.shard_tensor(x,
+                                  dist_attr={
+                                      "process_mesh": process_mesh,
+                                      "dims_mapping": [0, -1]
+                                  })
 
             w = paddle.static.data(name="w", shape=[4, 4], dtype='float32')
-            w = auto.shard_tensor(
-                w,
-                dist_attr={
-                    "process_mesh": process_mesh,
-                    "dims_mapping": [-1, -1]
-                })
-
-            # y = paddle.distributed.shard_op(paddle.matmul, process_mesh, {
-            #     x.name: [-1, -1],
-            #     w.name: [-1, -1]
-            # }, **{"x": x,
-            #       "y": w})[0]
-
-            y = paddle.distributed.shard_op(
-                paddle.matmul,
-                dist_attr={
-                    "process_mesh": process_mesh,
-                    x: {
-                        "dims_mapping": [-1, -1]
-                    },
-                    w: {
-                        "dims_mapping": [-1, -1]
-                    }
-                })(x, w)[0]
+            w = auto.shard_tensor(w,
+                                  dist_attr={
+                                      "process_mesh": process_mesh,
+                                      "dims_mapping": [-1, -1]
+                                  })
+
+            y = paddle.distributed.shard_op(paddle.matmul,
+                                            dist_attr={
+                                                "process_mesh": process_mesh,
+                                                x: {
+                                                    "dims_mapping": [-1, -1]
+                                                },
+                                                w: {
+                                                    "dims_mapping": [-1, -1]
+                                                }
+                                            })(x, w)
 
         rank_id = 0
         dist_context = DistributedContext()