diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt b/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt index 8411669a3fe5fbec96b41cbc49106d9a32c5054b..a4575a35fbfe1517d8445f0aa3c265b240d547b1 100644 --- a/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt @@ -1,4 +1,4 @@ cc_library( spmd_rule - SRCS common.cc dist_tensor_spec.cc matmul_spmd_rule.cc + SRCS common.cc dist_tensor_spec.cc matmul_spmd_rule.cc replicated_spmd_rule.cc DEPS phi) diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/common.cc b/paddle/fluid/distributed/auto_parallel/spmd_rules/common.cc index c948acd715bcfe14398295449f8bb65bca8d8f6b..f957412fc69388448bd7ea27591bc735ca2cb35c 100644 --- a/paddle/fluid/distributed/auto_parallel/spmd_rules/common.cc +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/common.cc @@ -160,6 +160,14 @@ std::string GetBroadcastAxes(const int64_t& tenosr_ndim, return alphabet.substr(broadcast_ndim - tenosr_ndim, tenosr_ndim); } +TensorDistAttr ReplicatedOnMesh(const TensorDistAttr& src_dist_attr) { + TensorDistAttr replicated_dist_attr = src_dist_attr; + replicated_dist_attr.clear_annotated(); + size_t tensor_ndim = replicated_dist_attr.dims_mapping().size(); + replicated_dist_attr.set_dims_mapping(std::vector(tensor_ndim, -1)); + return replicated_dist_attr; +} + // SPMDRuleMap SPMDRuleMap& SPMDRuleMap::Instance() { static SPMDRuleMap g_spmd_rule_map; diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/common.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/common.h index 9d7c7086d91d1b9280719b40d8ee52fada01e186..a3dcbe0cce7cee7c73b49b696ca649c4dc835a1b 100644 --- a/paddle/fluid/distributed/auto_parallel/spmd_rules/common.h +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/common.h @@ -17,6 +17,7 @@ limitations under the License. */ #include #include #include +#include #include #include "paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.h" @@ -106,6 +107,10 @@ int64_t ShardingMergeForAxis(const std::string& axis, const int64_t& mesh_dim1, const int64_t& mesh_dim2); +// Intend to use for generating the TensorDistAttr of output based on the input +// activation TensorDistAttr. The process_mesh, batch_dim, dynamic_dim are +// copied with annotated is forced to False, and dims_mapping is leave to be +// null. TensorDistAttr CopyTensorDistAttrForOutput(const TensorDistAttr& src_dist_attr); // Resolute the partial mesh dimension of a output tensor, giving the @@ -124,6 +129,10 @@ std::string GetBroadcastAxes(const int64_t& tenosr_ndim, const int64_t& broadcast_ndim, const std::string& alphabet); +// Return a NEW TensorDistAttr whose dims mapping is consist of "-1" +// (unsharded). +TensorDistAttr ReplicatedOnMesh(const TensorDistAttr& src_dist_attr); + // The static map that stores and initializes all the registered SPMD rules. class SPMDRuleMap { public: diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.cc b/paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.cc new file mode 100644 index 0000000000000000000000000000000000000000..1b072b1003d3b81f2f73ff470aad198ba4f19667 --- /dev/null +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.cc @@ -0,0 +1,51 @@ +/* Copyright (c) 2023 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 + +http://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. */ + +#include "paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.h" + +namespace paddle { +namespace distributed { +namespace auto_parallel { + +using phi::distributed::auto_parallel::str_join; + +std::pair, std::vector> +ReplicatedSPMDRule::InferForward(const std::vector& input_specs, + const paddle::framework::AttributeMap& attrs) { + std::vector intput_dist_attrs; + std::vector output_dist_attrs; + intput_dist_attrs.reserve(input_specs.size()); + + for (auto& input_spec : input_specs) { + intput_dist_attrs.push_back(ReplicatedOnMesh(input_spec.dist_attr())); + } + + // TODO(ljz): we need to know num of output and size of each output before + // generate the excat replicasted dist tensor attr for the current op. + // here we just assume that only one output tensor and has the same size as + // the first input tensor. + return {intput_dist_attrs, {ReplicatedOnMesh(input_specs[0].dist_attr())}}; +} + +std::pair, std::vector> +ReplicatedSPMDRule::InferBackward( + const std::vector& input_specs, + const paddle::framework::AttributeMap& attrs) { + PADDLE_THROW(phi::errors::Unimplemented( + "InferBackward of ReplicatedSPMDRule is NOT implemented yet.")); +} + +} // namespace auto_parallel +} // namespace distributed +} // namespace paddle diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.h new file mode 100644 index 0000000000000000000000000000000000000000..bcca646d351d534982be51ad88bd46ff5bc7111f --- /dev/null +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.h @@ -0,0 +1,41 @@ +/* Copyright (c) 2023 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 + +http://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. */ + +#pragma once + +#include "paddle/fluid/distributed/auto_parallel/spmd_rules/common.h" + +namespace paddle { +namespace distributed { +namespace auto_parallel { + +// A Bottom Line Rule that enforces input(s) and output(s) of the Op to be +// replicated among the given mesh. +class ReplicatedSPMDRule : public SPMDRuleBase { + public: + // The dims_mapping of ALL TensorDistAttrs would be repeat of "-1" + // (unsharded). + std::pair, std::vector> + InferForward(const std::vector& input_specs, + const paddle::framework::AttributeMap& attrs) override; + + // The dims_mapping of ALL TensorDistAttrs would be repeat of "-1" + // (unsharded). + std::pair, std::vector> + InferBackward(const std::vector& output_specs, + const paddle::framework::AttributeMap& attrs) override; +}; +} // namespace auto_parallel +} // namespace distributed +} // namespace paddle diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h index 334723059411b112939d25d0637d532d783115e3..329d613776174b101daf68a194e9431125a3edc4 100644 --- a/paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h @@ -16,6 +16,7 @@ #include "paddle/fluid/distributed/auto_parallel/spmd_rules/common.h" #include "paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h" +#include "paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.h" // TODO(ljz) Automatic this process in cmake file. namespace paddle { @@ -25,6 +26,9 @@ namespace auto_parallel { // matmul rule REGISTER_SPMD_RULE(matmul, MatmulSPMDRule); +// replicated rule +REGISTER_SPMD_RULE(replicated, ReplicatedSPMDRule); + } // namespace auto_parallel } // namespace distributed } // namespace paddle diff --git a/test/auto_parallel/spmd_rules/test_replicated_rule.py b/test/auto_parallel/spmd_rules/test_replicated_rule.py new file mode 100644 index 0000000000000000000000000000000000000000..cd91e8a8fb765d6a0e874c36053ecfb2d1b009fe --- /dev/null +++ b/test/auto_parallel/spmd_rules/test_replicated_rule.py @@ -0,0 +1,61 @@ +# Copyright (c) 2023 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 +# +# http://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 + +from paddle.distributed.auto_parallel.static.completion import get_spmd_rule +from paddle.distributed.auto_parallel.static.dist_attribute import ( + DistTensorSpec, + TensorDistAttr, +) +from paddle.distributed.fleet import auto + + +class TestMatmulSPMDRule(unittest.TestCase): + def setUp(self): + self.rule = get_spmd_rule("replicated") + + x_shape = [64, 32, 10, 10] + y_shape = [32, 48] + process_mesh = auto.ProcessMesh(mesh=[[0, 1, 2], [3, 4, 5]]) + + x_tensor_dist_attr = TensorDistAttr() + x_tensor_dist_attr.dims_mapping = [-1, 1, 0, -1] + x_tensor_dist_attr.process_mesh = process_mesh + self.x_dist_tensor_spec = DistTensorSpec(x_shape, x_tensor_dist_attr) + + y_tensor_dist_attr = TensorDistAttr() + y_tensor_dist_attr.dims_mapping = [0, -1] + y_tensor_dist_attr.process_mesh = process_mesh + self.y_dist_tensor_spec = DistTensorSpec(y_shape, y_tensor_dist_attr) + + def test_replicated_infer_forward(self): + # return all -1 + result_tensor_specs = self.rule.infer_forward( + [self.x_dist_tensor_spec, self.y_dist_tensor_spec], {} + ) + self.assertEqual(len(result_tensor_specs), 2) + self.assertEqual(len(result_tensor_specs[0]), 2) + self.assertEqual(len(result_tensor_specs[1]), 1) + self.assertEqual( + result_tensor_specs[0][0].dims_mapping, [-1, -1, -1, -1] + ) + self.assertEqual(result_tensor_specs[0][1].dims_mapping, [-1, -1]) + self.assertEqual( + result_tensor_specs[1][0].dims_mapping, [-1, -1, -1, -1] + ) + + +if __name__ == "__main__": + unittest.main()