dist.to_static support pir program (#62560)

zhiqiu · web-flow · commit bc56513ce46c · 2024-03-09T22:55:20.000+08:00
* auto_parallel engine build pir program

* skip prepare_op_amp_options in build_program

* add ut

* fix cmake

* remove print
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
@@ -13,11 +13,13 @@
 // limitations under the License.
 
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_dialect.h"
+
 #include "paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_op.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_type.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/type_storage.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
 
 REGISTER_FILE_SYMBOLS(dist_dialect);
 namespace paddle {
@@ -39,18 +41,45 @@ void DistDialect::initialize() {
 void DistDialect::PrintType(pir::Type type, std::ostream &os) const {
   if (auto dist_dense_tensor_type = type.dyn_cast<DistDenseTensorType>()) {
     // Todo: Design the dist dense tensor type print format.
-    os << dist_dense_tensor_type.dense_tensor_type();
+    os << type.dialect().name();
+    os << '.';
+    if (auto tensor_type = type.dyn_cast<pir::DenseTensorType>()) {
+      os << "tensor<";
+      for (auto d : common::vectorize(tensor_type.dims())) {
+        os << d;
+        os << "x";
+      }
+      tensor_type.dtype().Print(os);
+      os << ", ";
+      PrintAttribute(dist_dense_tensor_type.tensor_dist_attr(), os);
+      os << ">";
+    }
   } else {
     os << "error_type!";
   }
 }
 
 void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
   if (auto process_mesh_attr = attr.dyn_cast<ProcessMeshAttribute>()) {
-    os << process_mesh_attr.process_mesh();
+    os << "mesh: " << process_mesh_attr.process_mesh();
   } else if (auto tensor_dist_attr = attr.dyn_cast<TensorDistAttribute>()) {
     // Todo: Design the tensor dist attr print format.
-    os << tensor_dist_attr.process_mesh_attr().process_mesh();
+    os << "mesh: " << tensor_dist_attr.process_mesh_attr().process_mesh();
+    os << ", dims_mappings: [" +
+              phi::distributed::auto_parallel::str_join(
+                  tensor_dist_attr.dims_mapping()) +
+              "]";
+    if (tensor_dist_attr.partial_status().size() > 0) {
+      std::vector<std::string> partial_status_strs;
+      for (auto &itr : tensor_dist_attr.partial_status()) {
+        std::string s = "partial(" + std::to_string(itr.first) + "," +
+                        phi::ReduceTypeStrings[static_cast<int>(itr.second)] +
+                        ")";
+        partial_status_strs.emplace_back(s);
+      }
+      os << ", "
+         << phi::distributed::auto_parallel::str_join(partial_status_strs);
+    }
   } else {
     os << "error_attribute_type";
   }
diff --git a/python/paddle/distributed/auto_parallel/static/dist_input_spec.py b/python/paddle/distributed/auto_parallel/static/dist_input_spec.py
@@ -29,11 +29,13 @@ def __init__(
         stop_gradient=False,
         mesh=None,
         placements=None,
+        local_shape=None,
     ):
         super().__init__(shape, dtype, name, stop_gradient)
         self.mesh = copy.deepcopy(mesh)
         sharding_specs = get_shard_spec(mesh, placements, len(self.shape))
         self.dims_mapping = convert_to_dims_mapping(sharding_specs, mesh)
+        self.local_shape = local_shape
 
     @classmethod
     def from_dtensor(cls, dtensor, name=None):
@@ -53,6 +55,7 @@ def from_dtensor(cls, dtensor, name=None):
             stop_gradient=dtensor.stop_gradient,
             mesh=dtensor.process_mesh,
             placements=dtensor.placements,
+            local_shape=dtensor._local_value().shape,
         )
 
     def __repr__(self):
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -239,6 +239,9 @@ def __init__(
         self._dygraph_mode = False
         self._tuning = self._strategy.tuning
         self._acc_steps = 1
+        self._in_pir_mode = paddle.base.framework.get_flags(
+            "FLAGS_enable_pir_api"
+        )["FLAGS_enable_pir_api"]
         if self._strategy.gradient_merge.enable:
             self._acc_steps = self._strategy.gradient_merge.k_steps
         elif self._strategy.pipeline.enable:
@@ -618,6 +621,9 @@ def _prepare_logger(
     def _prepare_program(self, mode, init_parameters=True):
         # Do the build process
         self._build(mode)
+        # TODO(zhiqiu): fit the processes below for pir
+        if self._in_pir_mode:
+            return
         # Do the planning process
         self._plan(mode)
         # Do the parallel process
@@ -676,7 +682,7 @@ def _build(self, mode):
 
             self._inputs = self.program_helper.input_vars
             self._labels = self.program_helper.label_vars
-            self._process_dist_input_specs()
+            # self._process_dist_input_specs()
             outputs = self.program_helper.output_vars
             self._losses = self.program_helper.loss_vars
             metrics = self.program_helper.metric_vars
@@ -729,6 +735,17 @@ def _build(self, mode):
                 ), "the type of `loss` of the Engine arguments should be Variable."
                 self._losses = auto_utils.to_list(self._loss)
 
+        # TODO(zhiqiu): distributed_context is no longer used in pir_program
+        # so, just return here and need to reimplement the logics below
+        if self._in_pir_mode:
+            if mode != "train":
+                self._fwd_main_progs[mode] = serial_main_prog.clone(
+                    for_test=True
+                )
+            else:
+                self._fwd_main_progs[mode] = serial_main_prog
+            return
+
         default_ctx = get_default_distributed_context()
         if not default_ctx.has_annotation:
             # We build the world process group because the data parallel
diff --git a/python/paddle/jit/dy2static/function_spec.py b/python/paddle/jit/dy2static/function_spec.py
@@ -194,6 +194,20 @@ def pir_to_static_inputs_with_spec(self, input_with_spec, main_program):
                         dtype=convert_dtype(var_spec.dtype),
                     )
                     feed_value.stop_gradient = stop_gradient
+
+                    # warp dist tensor
+                    from paddle.distributed.auto_parallel.static.dist_input_spec import (
+                        DistributedInputSpec,
+                    )
+
+                    if isinstance(var_spec, DistributedInputSpec):
+                        dist_dense_tensor_type = paddle.base.libpaddle.pir.create_dist_dense_tensor_type_by_dense_tensor(
+                            feed_value.type(),
+                            var_spec.local_shape,
+                            var_spec.mesh,
+                            var_spec.dims_mapping,
+                        )
+                        feed_value.set_type(dist_dense_tensor_type)
                 else:
                     feed_value = var_spec
                 inputs.append(feed_value)
@@ -225,8 +239,29 @@ def to_static_inputs_with_spec(self, input_with_spec, main_program):
                     need_check_feed=False,
                     stop_gradient=stop_gradient,
                 )
+                # warp dist tensor
+                from paddle.distributed.auto_parallel.static.dist_input_spec import (
+                    DistributedInputSpec,
+                )
+                from paddle.distributed.auto_parallel.static.dist_tensor import (
+                    DistributedTensor,
+                )
+
+                if isinstance(var_spec, DistributedInputSpec):
+                    from paddle.distributed.auto_parallel.static.dist_context import (
+                        get_default_distributed_context,
+                    )
+
+                    default_dist_ctx = get_default_distributed_context()
+                    dist_tensor = DistributedTensor(feed_layer)
+                    dist_tensor.dist_attr.process_mesh = var_spec.mesh
+                    dist_tensor.dist_attr.dims_mapping = var_spec.dims_mapping
+                    dist_tensor.dist_attr.mark_annotated("process_mesh")
+                    dist_tensor.dist_attr.mark_annotated("dims_mapping")
+                    default_dist_ctx.add_dist_tensor_for_program(dist_tensor)
             else:
                 feed_layer = var_spec
+
             inputs.append(feed_layer)
 
         return paddle.utils.pack_sequence_as(input_with_spec, inputs)
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -4,6 +4,7 @@
 add_subdirectory(spmd_rules)
 add_subdirectory(hybrid_strategy)
 add_subdirectory(custom_op)
+add_subdirectory(pir)
 
 if(WITH_DISTRIBUTE AND WITH_GPU)
 
diff --git a/test/auto_parallel/pir/CMakeLists.txt b/test/auto_parallel/pir/CMakeLists.txt
@@ -0,0 +1,5 @@
+if(WITH_DISTRIBUTE AND WITH_GPU)
+  py_test_modules(test_to_static_pir_program MODULES test_to_static_pir_program)
+  set_tests_properties(test_to_static_pir_program
+                       PROPERTIES ENVIRONMENT "FLAGS_enable_pir_api=1")
+endif()
diff --git a/test/auto_parallel/pir/test_to_static_pir_program.py b/test/auto_parallel/pir/test_to_static_pir_program.py
@@ -0,0 +1,115 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import Shard
+from paddle.io import DataLoader
+
+BATCH_SIZE = 4
+BATCH_NUM = 4
+IMAGE_SIZE = 16
+CLASS_NUM = 8
+np.random.seed(2024)
+
+
+class RandomDataset(paddle.io.Dataset):
+    def __init__(self, images, labels, num_samples):
+        self.images = images
+        self.labels = labels
+        self.num_samples = num_samples
+
+    def __getitem__(self, idx):
+        return self.images[idx], self.labels[idx]
+
+    def __len__(self):
+        return self.num_samples
+
+
+class DemoNet(nn.Layer):
+    def __init__(self, mesh):
+        super().__init__()
+        self._mesh = mesh
+        self.linear_0 = nn.Linear(IMAGE_SIZE, IMAGE_SIZE)
+        self.linear_1 = nn.Linear(IMAGE_SIZE, CLASS_NUM)
+        self.relu = nn.ReLU()
+        # shard the weights of this layer
+        self.linear_0.weight = dist.shard_tensor(
+            self.linear_0.weight,
+            self._mesh,
+            [Shard(1)],
+            stop_gradient=False,
+        )
+        self.linear_1.weight = dist.shard_tensor(
+            self.linear_1.weight,
+            self._mesh,
+            [Shard(0)],
+            stop_gradient=False,
+        )
+
+    def forward(self, x):
+        out = self.linear_0(x)
+        out = self.relu(out)
+        out = self.linear_1(out)
+        return out
+
+
+def create_data_loader():
+    images = np.random.rand(BATCH_SIZE, IMAGE_SIZE).astype('float32')
+    labels = np.random.rand(BATCH_SIZE, CLASS_NUM).astype('float32')
+    dataset = RandomDataset(images, labels, BATCH_SIZE)
+    loader = DataLoader(dataset, batch_size=BATCH_SIZE)
+    return loader
+
+
+class TestToStaticPirProgram(unittest.TestCase):
+    def test_to_static_program(self):
+        paddle.base.set_flags({'FLAGS_enable_pir_api': 1})
+        mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        layer = DemoNet(mesh)
+        opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=layer.parameters()
+        )
+        loss_fn = nn.MSELoss()
+        loader = create_data_loader()
+        dist_loader = dist.shard_dataloader(loader, meshes=[mesh])
+        dist_model = dist.to_static(layer, dist_loader, loss_fn, opt)
+
+        dist_model.train()
+        main_program = dist_model._engine._fwd_main_progs["train"]
+        for op in main_program.global_block().ops:
+            tensor = op.result(0)
+            if op.name() == 'pd_op.data':
+                self.assertTrue(tensor.is_dist_dense_tensor_type())
+                self.assertEqual(tensor.process_mesh.shape, [2])
+                self.assertEqual(tensor.process_mesh.process_ids, [0, 1])
+                self.assertEqual(tensor.dims_mapping, [-1, -1])
+                self.assertEqual(tensor.partial_dims, set())
+            else:
+                self.assertTrue(tensor.is_dense_tensor_type())
+                self.assertFalse(tensor.is_dist_dense_tensor_type())
+
+        # training
+        # dist_model.train()
+        # for batch_id, (image, label) in enumerate(dist_loader()):
+        #     loss = dist_model(image, label)
+
+
+if __name__ == "__main__":
+    unittest.main()
