From e14eecb4f7cb5025a0a218956105c9804b49674e Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 14 May 2025 20:30:05 +0800 Subject: [PATCH 01/10] Config spmd_rule yaml for log_softmax --- paddle/phi/ops/yaml/backward.yaml | 1 + paddle/phi/ops/yaml/ops.yaml | 1 + 2 files changed, 2 insertions(+) diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml index 6b1849aa4f8a40..254ea3c6633166 100644 --- a/paddle/phi/ops/yaml/backward.yaml +++ b/paddle/phi/ops/yaml/backward.yaml @@ -1978,6 +1978,7 @@ infer_meta : func : UnchangedInferMeta param: [out] + spmd_rule : SoftmaxGradInferSpmd kernel : func : log_softmax_grad data_type : out_grad diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml index cc9e0da3a9bcea..1adee28f2c2ac0 100755 --- a/paddle/phi/ops/yaml/ops.yaml +++ b/paddle/phi/ops/yaml/ops.yaml @@ -3113,6 +3113,7 @@ output : Tensor(out) infer_meta : func : UnchangedInferMetaCheckAxis + spmd_rule : SoftmaxInferSpmd kernel : func : log_softmax data_type : x From f678ec6514d2c9b518e3fa38ad00b71f1843b0d8 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 14 May 2025 21:48:07 +0800 Subject: [PATCH 02/10] split the general func for topk, cummax and cummin;add spmd_rule for cummax and cummin. --- paddle/phi/infermeta/spmd_rules/cummax.cc | 34 +++++++++++++++++++++++ paddle/phi/infermeta/spmd_rules/cummax.h | 32 +++++++++++++++++++++ paddle/phi/infermeta/spmd_rules/cummin.cc | 34 +++++++++++++++++++++++ paddle/phi/infermeta/spmd_rules/cummin.h | 32 +++++++++++++++++++++ paddle/phi/infermeta/spmd_rules/rules.cc | 11 ++++++++ paddle/phi/infermeta/spmd_rules/topk.cc | 30 ++++++++++++++------ paddle/phi/infermeta/spmd_rules/topk.h | 6 +++- paddle/phi/ops/yaml/backward.yaml | 2 ++ paddle/phi/ops/yaml/ops.yaml | 2 ++ 9 files changed, 173 insertions(+), 10 deletions(-) create mode 100644 paddle/phi/infermeta/spmd_rules/cummax.cc create mode 100644 paddle/phi/infermeta/spmd_rules/cummax.h create mode 100644 paddle/phi/infermeta/spmd_rules/cummin.cc create mode 100644 paddle/phi/infermeta/spmd_rules/cummin.h diff --git a/paddle/phi/infermeta/spmd_rules/cummax.cc b/paddle/phi/infermeta/spmd_rules/cummax.cc new file mode 100644 index 00000000000000..e5e6e73c3289a0 --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/cummax.cc @@ -0,0 +1,34 @@ +/* Copyright (c) 2025 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/phi/infermeta/spmd_rules/cummax.h" +#include "paddle/phi/infermeta/spmd_rules/topk.h" + +namespace phi { +namespace distributed { + +SpmdInfo CummaxInferSpmd(const DistMetaTensor& x, int axis, DataType dtype) { + return TopkInferSpmdBase(x, axis); +} + +SpmdInfo CummaxGradInferSpmd(const DistMetaTensor& x, + const DistMetaTensor& indices, + const DistMetaTensor& out_grad, + int axis, + DataType dtype) { + return TopkGradInferSpmdBase(x, indices, out_grad, axis); +} + +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/cummax.h b/paddle/phi/infermeta/spmd_rules/cummax.h new file mode 100644 index 00000000000000..84059419325303 --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/cummax.h @@ -0,0 +1,32 @@ +/* Copyright (c) 2025 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/phi/core/distributed/auto_parallel/dist_meta_tensor.h" +#include "paddle/phi/core/distributed/type_defs.h" + +namespace phi { +namespace distributed { + +SpmdInfo CummaxInferSpmd(const DistMetaTensor& x, int axis, DataType dtype); + +SpmdInfo CummaxGradInferSpmd(const DistMetaTensor& x, + const DistMetaTensor& indices, + const DistMetaTensor& out_grad, + int axis, + DataType dtype); + +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/cummin.cc b/paddle/phi/infermeta/spmd_rules/cummin.cc new file mode 100644 index 00000000000000..8ce157e5c09c07 --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/cummin.cc @@ -0,0 +1,34 @@ +/* Copyright (c) 2025 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/phi/infermeta/spmd_rules/cummin.h" +#include "paddle/phi/infermeta/spmd_rules/topk.h" + +namespace phi { +namespace distributed { + +SpmdInfo CumminInferSpmd(const DistMetaTensor& x, int axis, DataType dtype) { + return TopkInferSpmdBase(x, axis); +} + +SpmdInfo CumminGradInferSpmd(const DistMetaTensor& x, + const DistMetaTensor& indices, + const DistMetaTensor& out_grad, + int axis, + DataType dtype) { + return TopkGradInferSpmdBase(x, indices, out_grad, axis); +} + +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/cummin.h b/paddle/phi/infermeta/spmd_rules/cummin.h new file mode 100644 index 00000000000000..5698c33f09d202 --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/cummin.h @@ -0,0 +1,32 @@ +/* Copyright (c) 2025 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/phi/core/distributed/auto_parallel/dist_meta_tensor.h" +#include "paddle/phi/core/distributed/type_defs.h" + +namespace phi { +namespace distributed { + +SpmdInfo CumminInferSpmd(const DistMetaTensor& x, int axis, DataType dtype); + +SpmdInfo CumminGradInferSpmd(const DistMetaTensor& x, + const DistMetaTensor& indices, + const DistMetaTensor& out_grad, + int axis, + DataType dtype); + +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc index 35b582811505f0..baca0f189f4e71 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.cc +++ b/paddle/phi/infermeta/spmd_rules/rules.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/phi/infermeta/spmd_rules/rules.h" +#include "paddle/phi/infermeta/spmd_rules/topk.h" /** * Design Notes: @@ -752,4 +753,14 @@ PD_REGISTER_SPMD_RULE(nonzero, // add_n PD_REGISTER_SPMD_RULE(add_n, PD_INFER_SPMD(phi::distributed::AddNInferSpmd)); + +// cummax +PD_REGISTER_SPMD_RULE(cummax, + PD_INFER_SPMD(phi::distributed::TopkInferSpmdBase), + PD_INFER_SPMD(phi::distributed::TopkGradInferSpmdBase)); +// cummin +PD_REGISTER_SPMD_RULE(cummin, + PD_INFER_SPMD(phi::distributed::TopkInferSpmdBase), + PD_INFER_SPMD(phi::distributed::TopkGradInferSpmdBase)); + } // namespace phi::distributed diff --git a/paddle/phi/infermeta/spmd_rules/topk.cc b/paddle/phi/infermeta/spmd_rules/topk.cc index c48365419baa74..bc0481d424775c 100644 --- a/paddle/phi/infermeta/spmd_rules/topk.cc +++ b/paddle/phi/infermeta/spmd_rules/topk.cc @@ -20,8 +20,7 @@ limitations under the License. */ namespace phi { namespace distributed { -SpmdInfo TopkInferSpmd( - const DistMetaTensor& x, int k, int axis, bool largest, bool sorted) { +SpmdInfo TopkInferSpmdBase(const DistMetaTensor& x, int axis) { // Verify input args EXTRACT_SHAPE_AND_DIST_ATTR(x); axis = axis < 0 ? x_ndim + axis : axis; @@ -60,13 +59,10 @@ SpmdInfo TopkInferSpmd( return {{x_dist_attr_dst}, {out_dist_attr_dst, indices_dist_attr_dst}}; } -SpmdInfo TopkGradInferSpmd(const DistMetaTensor& x, - const DistMetaTensor& indices, - const DistMetaTensor& out_grad, - int k, - int axis, - bool largest, - bool sorted) { +SpmdInfo TopkGradInferSpmdBase(const DistMetaTensor& x, + const DistMetaTensor& indices, + const DistMetaTensor& out_grad, + int axis) { // Verify input args EXTRACT_SHAPE_AND_DIST_ATTR(x); EXTRACT_SHAPE_AND_DIST_ATTR(indices); @@ -141,6 +137,22 @@ SpmdInfo TopkGradInferSpmd(const DistMetaTensor& x, return {{x_dist_attr_dst, indices_dist_attr_dst, out_grad_dist_attr_dst}, {x_grad_dist_attr_dst}}; } + +SpmdInfo TopkInferSpmd( + const DistMetaTensor& x, int k, int axis, bool largest, bool sorted) { + return TopkInferSpmdBase(x, axis); +} + +SpmdInfo TopkGradInferSpmd(const DistMetaTensor& x, + const DistMetaTensor& indices, + const DistMetaTensor& out_grad, + int k, + int axis, + bool largest, + bool sorted) { + return TopkGradInferSpmdBase(x, indices, out_grad, axis); +} + SpmdInfo TopkInferSpmdDynamic(const DistMetaTensor& x, const Scalar& k, int axis, diff --git a/paddle/phi/infermeta/spmd_rules/topk.h b/paddle/phi/infermeta/spmd_rules/topk.h index 3f69caeb2be6fb..d0d34c8ae69b94 100644 --- a/paddle/phi/infermeta/spmd_rules/topk.h +++ b/paddle/phi/infermeta/spmd_rules/topk.h @@ -20,7 +20,11 @@ limitations under the License. */ namespace phi { namespace distributed { - +SpmdInfo TopkInferSpmdBase(const DistMetaTensor& x, int axis); +SpmdInfo TopkGradInferSpmdBase(const DistMetaTensor& x, + const DistMetaTensor& indices, + const DistMetaTensor& out_grad, + int axis); SpmdInfo TopkInferSpmd( const DistMetaTensor& x, int k, int axis, bool largest, bool sorted); diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml index 254ea3c6633166..e74552aaa7fdae 100644 --- a/paddle/phi/ops/yaml/backward.yaml +++ b/paddle/phi/ops/yaml/backward.yaml @@ -697,6 +697,7 @@ infer_meta : func : UnchangedInferMeta param: [x] + spmd_rule : CummaxGradInferSpmd kernel : func : cummax_grad data_type : out_grad @@ -708,6 +709,7 @@ infer_meta : func : UnchangedInferMeta param: [x] + spmd_rule : CumminGradInferSpmd kernel : func : cummin_grad data_type : out_grad diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml index 1adee28f2c2ac0..d28da521bbaffa 100755 --- a/paddle/phi/ops/yaml/ops.yaml +++ b/paddle/phi/ops/yaml/ops.yaml @@ -1275,6 +1275,7 @@ output : Tensor(out), Tensor(indices) infer_meta : func : CumWithIndicesInferMeta + spmd_rule : CummaxInferSpmd kernel : func : cummax data_type : x @@ -1286,6 +1287,7 @@ output : Tensor(out), Tensor(indices) infer_meta : func : CumWithIndicesInferMeta + spmd_rule : CumminInferSpmd kernel : func : cummin data_type : x From 8747e1be41223c478f5c3c868fe957febd1ea068 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 14 May 2025 22:01:05 +0800 Subject: [PATCH 03/10] fix rules --- paddle/phi/infermeta/spmd_rules/rules.cc | 1 - paddle/phi/infermeta/spmd_rules/rules.h | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc index baca0f189f4e71..34a018540bd916 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.cc +++ b/paddle/phi/infermeta/spmd_rules/rules.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/phi/infermeta/spmd_rules/rules.h" -#include "paddle/phi/infermeta/spmd_rules/topk.h" /** * Design Notes: diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h index 45023b11c4422c..f9c6df6aceb433 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.h +++ b/paddle/phi/infermeta/spmd_rules/rules.h @@ -27,6 +27,8 @@ limitations under the License. */ #include "paddle/phi/infermeta/spmd_rules/concat.h" #include "paddle/phi/infermeta/spmd_rules/conv2d.h" #include "paddle/phi/infermeta/spmd_rules/cross_entropy_with_softmax.h" +#include "paddle/phi/infermeta/spmd_rules/cummax.h" +#include "paddle/phi/infermeta/spmd_rules/cummin.h" #include "paddle/phi/infermeta/spmd_rules/cumsum.h" #include "paddle/phi/infermeta/spmd_rules/default_data_parallel.h" #include "paddle/phi/infermeta/spmd_rules/dropout.h" From d8b307b54a87facd7a74be5696a5062332427cb8 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Thu, 15 May 2025 21:51:32 +0800 Subject: [PATCH 04/10] fix topk spmd_rule; add unittests --- paddle/phi/infermeta/spmd_rules/topk.cc | 4 +- .../spmd_rules/test_cummax_rule.py | 96 +++++++++++++++++++ .../spmd_rules/test_cummin_rule.py | 96 +++++++++++++++++++ test/cpp/auto_parallel/spmd_rule_test.cc | 96 +++++++++++++++++++ 4 files changed, 290 insertions(+), 2 deletions(-) create mode 100644 test/auto_parallel/spmd_rules/test_cummax_rule.py create mode 100644 test/auto_parallel/spmd_rules/test_cummin_rule.py diff --git a/paddle/phi/infermeta/spmd_rules/topk.cc b/paddle/phi/infermeta/spmd_rules/topk.cc index bc0481d424775c..fc4d3dd15818c1 100644 --- a/paddle/phi/infermeta/spmd_rules/topk.cc +++ b/paddle/phi/infermeta/spmd_rules/topk.cc @@ -91,7 +91,8 @@ SpmdInfo TopkGradInferSpmdBase(const DistMetaTensor& x, axis)); // Build einsum notation std::string alphabet = "abcdefghijlopqrstuvwxyz"; - std::string x_axes = alphabet.substr(0, x_ndim - 1); + std::string x_axes = alphabet.substr(0, x_ndim); + x_axes[axis] = '1'; std::string indices_axes = x_axes; std::string out_grad_axes = x_axes; @@ -108,7 +109,6 @@ SpmdInfo TopkGradInferSpmdBase(const DistMetaTensor& x, // Infer dims mapping std::vector x_grad_dims_mapping_dst = GetDimsMappingForAxes(x_axes, axis_to_dim_map); - x_grad_dims_mapping_dst.insert(x_grad_dims_mapping_dst.begin() + axis, -1); std::vector x_dims_mapping_dst = x_grad_dims_mapping_dst; std::vector indices_dims_mapping_dst = x_grad_dims_mapping_dst; std::vector out_grad_dims_mapping_dst = x_grad_dims_mapping_dst; diff --git a/test/auto_parallel/spmd_rules/test_cummax_rule.py b/test/auto_parallel/spmd_rules/test_cummax_rule.py new file mode 100644 index 00000000000000..8a3166e7feb05d --- /dev/null +++ b/test/auto_parallel/spmd_rules/test_cummax_rule.py @@ -0,0 +1,96 @@ +# Copyright (c) 2025 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 collections import OrderedDict + +from paddle.distributed.auto_parallel.static.dist_attribute import ( + DistTensorSpec, + TensorDistAttr, +) +from paddle.distributed.fleet import auto +from paddle.framework import core + + +class TestTopkSPMDRule(unittest.TestCase): + def setUp(self): + x_shape = [16, 16, 16] + out_shape = [16, 2, 16] + process_mesh = auto.ProcessMesh(mesh=[[0, 1], [2, 3]]) + + x_tensor_dist_attr = TensorDistAttr() + x_tensor_dist_attr.dims_mapping = [-1, -1, -1] + x_tensor_dist_attr.process_mesh = process_mesh + self.x_dist_tensor_spec = DistTensorSpec(x_shape, x_tensor_dist_attr) + out_tensor_dist_attr = TensorDistAttr() + out_tensor_dist_attr.dims_mapping = [-1, -1, -1] + out_tensor_dist_attr.process_mesh = process_mesh + self.out_dist_tensor_spec = DistTensorSpec( + out_shape, x_tensor_dist_attr + ) + + self.rule = core.get_phi_spmd_rule("cummax") + self.attrs = OrderedDict() + self.attrs['axis'] = 1 + + def test_cummax_forward(self): + # axis = 1 + # [0, 1, -1] --> [0, -1, -1], [0, -1, -1] + self.attrs['axis'] = 1 + self.x_dist_tensor_spec.set_dims_mapping([0, 1, -1]) + result_dist_attrs = self.rule.infer_forward( + self.x_dist_tensor_spec, + self.attrs['k'], + ) + + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + + self.assertEqual(len(inferred_input_dist_attrs), 1) + self.assertEqual(len(inferred_output_dist_attrs), 2) + + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) + self.assertEqual( + inferred_output_dist_attrs[0].dims_mapping, [0, -1, -1] + ) + + def test_cummax_backward(self): + # axis = 1 + # [0, -1, 1], [0, -1, 1], [-1, 1, -1] --> [0, -1, 1], [0, -1, 1], [0, -1, 1], [0, -1, 1] + self.attrs['axis'] = 1 + self.x_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + self.out_dist_tensor_spec.shape = [16, 2, 16] + self.out_dist_tensor_spec.set_dims_mapping([-1, 1, -1]) + result_dist_attrs = self.rule.infer_backward( + self.x_dist_tensor_spec, + self.out_dist_tensor_spec, + self.out_dist_tensor_spec, + self.attrs['axis'], + ) + + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 3) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/auto_parallel/spmd_rules/test_cummin_rule.py b/test/auto_parallel/spmd_rules/test_cummin_rule.py new file mode 100644 index 00000000000000..377aec9f216080 --- /dev/null +++ b/test/auto_parallel/spmd_rules/test_cummin_rule.py @@ -0,0 +1,96 @@ +# Copyright (c) 2025 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 collections import OrderedDict + +from paddle.distributed.auto_parallel.static.dist_attribute import ( + DistTensorSpec, + TensorDistAttr, +) +from paddle.distributed.fleet import auto +from paddle.framework import core + + +class TestTopkSPMDRule(unittest.TestCase): + def setUp(self): + x_shape = [16, 16, 16] + out_shape = [16, 2, 16] + process_mesh = auto.ProcessMesh(mesh=[[0, 1], [2, 3]]) + + x_tensor_dist_attr = TensorDistAttr() + x_tensor_dist_attr.dims_mapping = [-1, -1, -1] + x_tensor_dist_attr.process_mesh = process_mesh + self.x_dist_tensor_spec = DistTensorSpec(x_shape, x_tensor_dist_attr) + out_tensor_dist_attr = TensorDistAttr() + out_tensor_dist_attr.dims_mapping = [-1, -1, -1] + out_tensor_dist_attr.process_mesh = process_mesh + self.out_dist_tensor_spec = DistTensorSpec( + out_shape, x_tensor_dist_attr + ) + + self.rule = core.get_phi_spmd_rule("cummin") + self.attrs = OrderedDict() + self.attrs['axis'] = 1 + + def test_cummin_forward(self): + # axis = 1 + # [0, 1, -1] --> [0, -1, -1], [0, -1, -1] + self.attrs['axis'] = 1 + self.x_dist_tensor_spec.set_dims_mapping([0, 1, -1]) + result_dist_attrs = self.rule.infer_forward( + self.x_dist_tensor_spec, + self.attrs['k'], + ) + + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + + self.assertEqual(len(inferred_input_dist_attrs), 1) + self.assertEqual(len(inferred_output_dist_attrs), 2) + + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) + self.assertEqual( + inferred_output_dist_attrs[0].dims_mapping, [0, -1, -1] + ) + + def test_cummin_backward(self): + # axis = 1 + # [0, -1, 1], [0, -1, 1], [-1, 1, -1] --> [0, -1, 1], [0, -1, 1], [0, -1, 1], [0, -1, 1] + self.attrs['axis'] = 1 + self.x_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + self.out_dist_tensor_spec.shape = [16, 2, 16] + self.out_dist_tensor_spec.set_dims_mapping([-1, 1, -1]) + result_dist_attrs = self.rule.infer_backward( + self.x_dist_tensor_spec, + self.out_dist_tensor_spec, + self.out_dist_tensor_spec, + self.attrs['axis'], + ) + + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 3) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc index f41695748fdcb8..1cb20682eee9e8 100644 --- a/test/cpp/auto_parallel/spmd_rule_test.cc +++ b/test/cpp/auto_parallel/spmd_rule_test.cc @@ -2516,6 +2516,102 @@ TEST(Topk, Ctor) { check_dim_mapping(backward_info.second[0], {0, -1, 1}); } +TEST(Cummax, Ctor) { + std::vector mesh_shape = {2, 2}; + std::vector process_ids = {0, 1, 2, 3}; + std::vector dim_names = {"x", "y"}; + ProcessMesh process_mesh(mesh_shape, process_ids, dim_names); + int axis = 1; + + // test forward + // axis = 1 + // [0, 1, -1] -> [0, -1, -1], [0, -1, -1] + auto x_dist_attr = TensorDistAttr(); + x_dist_attr.set_process_mesh(process_mesh); + x_dist_attr.set_dims_mapping({0, 1, -1}); + x_dist_attr.set_dynamic_dims({false, false, false}); + + phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor( + common::make_ddim({16, 16, 16}), x_dist_attr); + phi::distributed::SpmdInfo forward_info = + phi::distributed::CummaxInferSpmd(x, axis, phi::DataType::INT64); + + EXPECT_EQ(forward_info.first.size(), 1UL); + EXPECT_EQ(forward_info.second.size(), 2UL); + check_dim_mapping(forward_info.first[0], {0, -1, -1}); + check_dim_mapping(forward_info.second[0], {0, -1, -1}); + check_dim_mapping(forward_info.second[1], {0, -1, -1}); + + // test backward + // axis = 1 + // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1, + // 1], [0, -1, 1] + x_dist_attr.set_dims_mapping({0, -1, 1}); + auto out_grad_dist_attr = TensorDistAttr(); + out_grad_dist_attr.set_process_mesh(process_mesh); + out_grad_dist_attr.set_dims_mapping({-1, 1, -1}); + out_grad_dist_attr.set_dynamic_dims({false, false, false}); + phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor( + common::make_ddim({16, 2, 16}), out_grad_dist_attr); + phi::distributed::SpmdInfo backward_info = + phi::distributed::CummaxGradInferSpmd( + x, x, out_grad, axis, phi::DataType::INT64); + EXPECT_EQ(backward_info.first.size(), 3UL); + EXPECT_EQ(backward_info.second.size(), 1UL); + check_dim_mapping(backward_info.first[0], {0, -1, 1}); + check_dim_mapping(backward_info.first[1], {0, -1, 1}); + check_dim_mapping(backward_info.first[2], {0, -1, 1}); + check_dim_mapping(backward_info.second[0], {0, -1, 1}); +} + +TEST(Cummin, Ctor) { + std::vector mesh_shape = {2, 2}; + std::vector process_ids = {0, 1, 2, 3}; + std::vector dim_names = {"x", "y"}; + ProcessMesh process_mesh(mesh_shape, process_ids, dim_names); + int axis = 1; + + // test forward + // axis = 1 + // [0, 1, -1] -> [0, -1, -1], [0, -1, -1] + auto x_dist_attr = TensorDistAttr(); + x_dist_attr.set_process_mesh(process_mesh); + x_dist_attr.set_dims_mapping({0, 1, -1}); + x_dist_attr.set_dynamic_dims({false, false, false}); + + phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor( + common::make_ddim({16, 16, 16}), x_dist_attr); + phi::distributed::SpmdInfo forward_info = + phi::distributed::CumminInferSpmd(x, axis, phi::DataType::INT64); + + EXPECT_EQ(forward_info.first.size(), 1UL); + EXPECT_EQ(forward_info.second.size(), 2UL); + check_dim_mapping(forward_info.first[0], {0, -1, -1}); + check_dim_mapping(forward_info.second[0], {0, -1, -1}); + check_dim_mapping(forward_info.second[1], {0, -1, -1}); + + // test backward + // axis = 1 + // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1, + // 1], [0, -1, 1] + x_dist_attr.set_dims_mapping({0, -1, 1}); + auto out_grad_dist_attr = TensorDistAttr(); + out_grad_dist_attr.set_process_mesh(process_mesh); + out_grad_dist_attr.set_dims_mapping({-1, 1, -1}); + out_grad_dist_attr.set_dynamic_dims({false, false, false}); + phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor( + common::make_ddim({16, 2, 16}), out_grad_dist_attr); + phi::distributed::SpmdInfo backward_info = + phi::distributed::CumminGradInferSpmd( + x, x, out_grad, axis, phi::DataType::INT64); + EXPECT_EQ(backward_info.first.size(), 3UL); + EXPECT_EQ(backward_info.second.size(), 1UL); + check_dim_mapping(backward_info.first[0], {0, -1, 1}); + check_dim_mapping(backward_info.first[1], {0, -1, 1}); + check_dim_mapping(backward_info.first[2], {0, -1, 1}); + check_dim_mapping(backward_info.second[0], {0, -1, 1}); +} + } // namespace auto_parallel } // namespace distributed } // namespace paddle From c8b54d04ed2dab34b4b21ba5964d6603f8dfc5a7 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Thu, 15 May 2025 23:07:32 +0800 Subject: [PATCH 05/10] Cooperate with GetDimsMappingForAxes to deal 1 for replicated.Because it can easily use to have more chance to reshared --- paddle/phi/infermeta/spmd_rules/utils.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/paddle/phi/infermeta/spmd_rules/utils.cc b/paddle/phi/infermeta/spmd_rules/utils.cc index 6d3e354b30e7cd..3f773c6184eccd 100644 --- a/paddle/phi/infermeta/spmd_rules/utils.cc +++ b/paddle/phi/infermeta/spmd_rules/utils.cc @@ -86,6 +86,10 @@ std::unordered_map ShardingMergeForTensors( for (auto& pair : tensor_axes_to_dim_pairs) { for (size_t i = 0; i < pair.second.size(); ++i) { auto tensor_axis = pair.first.substr(i, 1); + // Cooperate with GetDimsMappingForAxes to deal "1" for replicated. + if (tensor_axis == "1") { + continue; + } auto mesh_dim = pair.second[i]; if (axis_to_dim_map.count(tensor_axis) == 0) { From 644e782aabe65cd41ec61f15af38847cde66d101 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Fri, 16 May 2025 01:19:03 +0800 Subject: [PATCH 06/10] try to ci --- test/auto_parallel/spmd_rules/CMakeLists.txt | 2 ++ test/auto_parallel/spmd_rules/test_cummax_rule.py | 2 +- test/auto_parallel/spmd_rules/test_cummin_rule.py | 2 +- test/auto_parallel/spmd_rules/test_topk_rule.py | 2 +- test/cpp/auto_parallel/spmd_rule_test.cc | 4 ++-- 5 files changed, 7 insertions(+), 5 deletions(-) diff --git a/test/auto_parallel/spmd_rules/CMakeLists.txt b/test/auto_parallel/spmd_rules/CMakeLists.txt index 879d540a619c60..2fef779305c357 100644 --- a/test/auto_parallel/spmd_rules/CMakeLists.txt +++ b/test/auto_parallel/spmd_rules/CMakeLists.txt @@ -48,6 +48,8 @@ if(WITH_DISTRIBUTE) py_test_modules(test_add_n_rule MODULES test_add_n_rule) py_test_modules(test_mean_all_rule MODULES test_mean_all_rule) py_test_modules(test_argmin_rule MODULES test_argmin_rule) + py_test_modules(test_cummax_rule MODULES test_cummax_rule) + py_test_modules(test_cummin_rule MODULES test_cummin_rule) endif() # End of unittests WITH single card WITHOUT timeout diff --git a/test/auto_parallel/spmd_rules/test_cummax_rule.py b/test/auto_parallel/spmd_rules/test_cummax_rule.py index 8a3166e7feb05d..c35f9eced96662 100644 --- a/test/auto_parallel/spmd_rules/test_cummax_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummax_rule.py @@ -76,7 +76,7 @@ def test_cummax_backward(self): self.out_dist_tensor_spec.set_dims_mapping([-1, 1, -1]) result_dist_attrs = self.rule.infer_backward( self.x_dist_tensor_spec, - self.out_dist_tensor_spec, + self.x_dist_tensor_spec, self.out_dist_tensor_spec, self.attrs['axis'], ) diff --git a/test/auto_parallel/spmd_rules/test_cummin_rule.py b/test/auto_parallel/spmd_rules/test_cummin_rule.py index 377aec9f216080..0977a2f67330f9 100644 --- a/test/auto_parallel/spmd_rules/test_cummin_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummin_rule.py @@ -76,7 +76,7 @@ def test_cummin_backward(self): self.out_dist_tensor_spec.set_dims_mapping([-1, 1, -1]) result_dist_attrs = self.rule.infer_backward( self.x_dist_tensor_spec, - self.out_dist_tensor_spec, + self.x_dist_tensor_spec, self.out_dist_tensor_spec, self.attrs['axis'], ) diff --git a/test/auto_parallel/spmd_rules/test_topk_rule.py b/test/auto_parallel/spmd_rules/test_topk_rule.py index fe16055734cf3c..a0d57ec5d6bfde 100644 --- a/test/auto_parallel/spmd_rules/test_topk_rule.py +++ b/test/auto_parallel/spmd_rules/test_topk_rule.py @@ -82,7 +82,7 @@ def test_topk_backward(self): self.out_dist_tensor_spec.set_dims_mapping([-1, 1, -1]) result_dist_attrs = self.rule.infer_backward( self.x_dist_tensor_spec, - self.out_dist_tensor_spec, + self.x_dist_tensor_spec, self.out_dist_tensor_spec, self.attrs['k'], self.attrs['axis'], diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc index 1cb20682eee9e8..0db89bdeaa84a2 100644 --- a/test/cpp/auto_parallel/spmd_rule_test.cc +++ b/test/cpp/auto_parallel/spmd_rule_test.cc @@ -2552,7 +2552,7 @@ TEST(Cummax, Ctor) { out_grad_dist_attr.set_dims_mapping({-1, 1, -1}); out_grad_dist_attr.set_dynamic_dims({false, false, false}); phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor( - common::make_ddim({16, 2, 16}), out_grad_dist_attr); + common::make_ddim({16, 16, 16}), out_grad_dist_attr); phi::distributed::SpmdInfo backward_info = phi::distributed::CummaxGradInferSpmd( x, x, out_grad, axis, phi::DataType::INT64); @@ -2600,7 +2600,7 @@ TEST(Cummin, Ctor) { out_grad_dist_attr.set_dims_mapping({-1, 1, -1}); out_grad_dist_attr.set_dynamic_dims({false, false, false}); phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor( - common::make_ddim({16, 2, 16}), out_grad_dist_attr); + common::make_ddim({16, 16, 16}), out_grad_dist_attr); phi::distributed::SpmdInfo backward_info = phi::distributed::CumminGradInferSpmd( x, x, out_grad, axis, phi::DataType::INT64); From 92b2ae54f9ce1d0719dabe9e5e874ebc7724278a Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Fri, 16 May 2025 10:34:19 +0800 Subject: [PATCH 07/10] fix typos;fix bugs in Topk unittest --- test/auto_parallel/spmd_rules/test_cummax_rule.py | 2 +- test/auto_parallel/spmd_rules/test_cummin_rule.py | 2 +- test/cpp/auto_parallel/spmd_rule_test.cc | 6 ++++++ 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/test/auto_parallel/spmd_rules/test_cummax_rule.py b/test/auto_parallel/spmd_rules/test_cummax_rule.py index c35f9eced96662..eec68b2d8d8860 100644 --- a/test/auto_parallel/spmd_rules/test_cummax_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummax_rule.py @@ -51,7 +51,7 @@ def test_cummax_forward(self): self.x_dist_tensor_spec.set_dims_mapping([0, 1, -1]) result_dist_attrs = self.rule.infer_forward( self.x_dist_tensor_spec, - self.attrs['k'], + self.attrs['axis'], ) self.assertEqual(len(result_dist_attrs), 2) diff --git a/test/auto_parallel/spmd_rules/test_cummin_rule.py b/test/auto_parallel/spmd_rules/test_cummin_rule.py index 0977a2f67330f9..82f1086543a38e 100644 --- a/test/auto_parallel/spmd_rules/test_cummin_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummin_rule.py @@ -51,7 +51,7 @@ def test_cummin_forward(self): self.x_dist_tensor_spec.set_dims_mapping([0, 1, -1]) result_dist_attrs = self.rule.infer_forward( self.x_dist_tensor_spec, - self.attrs['k'], + self.attrs['axis'], ) self.assertEqual(len(result_dist_attrs), 2) diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc index 0db89bdeaa84a2..9659fd48b05221 100644 --- a/test/cpp/auto_parallel/spmd_rule_test.cc +++ b/test/cpp/auto_parallel/spmd_rule_test.cc @@ -2499,6 +2499,8 @@ TEST(Topk, Ctor) { // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1, // 1], [0, -1, 1] x_dist_attr.set_dims_mapping({0, -1, 1}); + x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16}), + x_dist_attr); auto out_grad_dist_attr = TensorDistAttr(); out_grad_dist_attr.set_process_mesh(process_mesh); out_grad_dist_attr.set_dims_mapping({-1, 1, -1}); @@ -2547,6 +2549,8 @@ TEST(Cummax, Ctor) { // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1, // 1], [0, -1, 1] x_dist_attr.set_dims_mapping({0, -1, 1}); + x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16}), + x_dist_attr); auto out_grad_dist_attr = TensorDistAttr(); out_grad_dist_attr.set_process_mesh(process_mesh); out_grad_dist_attr.set_dims_mapping({-1, 1, -1}); @@ -2595,6 +2599,8 @@ TEST(Cummin, Ctor) { // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1, // 1], [0, -1, 1] x_dist_attr.set_dims_mapping({0, -1, 1}); + x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16}), + x_dist_attr); auto out_grad_dist_attr = TensorDistAttr(); out_grad_dist_attr.set_process_mesh(process_mesh); out_grad_dist_attr.set_dims_mapping({-1, 1, -1}); From 63e28d7bdfc5ac199cc0d0433388d09e5d700e87 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Fri, 16 May 2025 11:13:26 +0800 Subject: [PATCH 08/10] fix bugs in unittest --- test/auto_parallel/spmd_rules/test_cummax_rule.py | 10 ++++++---- test/auto_parallel/spmd_rules/test_cummin_rule.py | 10 ++++++---- test/auto_parallel/spmd_rules/test_topk_rule.py | 8 +++++--- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/test/auto_parallel/spmd_rules/test_cummax_rule.py b/test/auto_parallel/spmd_rules/test_cummax_rule.py index eec68b2d8d8860..f7506888cc0096 100644 --- a/test/auto_parallel/spmd_rules/test_cummax_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummax_rule.py @@ -23,7 +23,7 @@ from paddle.framework import core -class TestTopkSPMDRule(unittest.TestCase): +class TestCummaxSPMDRule(unittest.TestCase): def setUp(self): x_shape = [16, 16, 16] out_shape = [16, 2, 16] @@ -61,11 +61,13 @@ def test_cummax_forward(self): self.assertEqual(len(inferred_input_dist_attrs), 1) self.assertEqual(len(inferred_output_dist_attrs), 2) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) self.assertEqual( inferred_output_dist_attrs[0].dims_mapping, [0, -1, -1] ) + self.assertEqual( + inferred_output_dist_attrs[1].dims_mapping, [0, -1, -1] + ) def test_cummax_backward(self): # axis = 1 @@ -87,8 +89,8 @@ def test_cummax_backward(self): self.assertEqual(len(inferred_input_dist_attrs), 3) self.assertEqual(len(inferred_output_dist_attrs), 1) self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [0, -1, 1]) self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) diff --git a/test/auto_parallel/spmd_rules/test_cummin_rule.py b/test/auto_parallel/spmd_rules/test_cummin_rule.py index 82f1086543a38e..441f5857ae204f 100644 --- a/test/auto_parallel/spmd_rules/test_cummin_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummin_rule.py @@ -23,7 +23,7 @@ from paddle.framework import core -class TestTopkSPMDRule(unittest.TestCase): +class TestCumminSPMDRule(unittest.TestCase): def setUp(self): x_shape = [16, 16, 16] out_shape = [16, 2, 16] @@ -61,11 +61,13 @@ def test_cummin_forward(self): self.assertEqual(len(inferred_input_dist_attrs), 1) self.assertEqual(len(inferred_output_dist_attrs), 2) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) self.assertEqual( inferred_output_dist_attrs[0].dims_mapping, [0, -1, -1] ) + self.assertEqual( + inferred_output_dist_attrs[1].dims_mapping, [0, -1, -1] + ) def test_cummin_backward(self): # axis = 1 @@ -87,8 +89,8 @@ def test_cummin_backward(self): self.assertEqual(len(inferred_input_dist_attrs), 3) self.assertEqual(len(inferred_output_dist_attrs), 1) self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [0, -1, 1]) self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) diff --git a/test/auto_parallel/spmd_rules/test_topk_rule.py b/test/auto_parallel/spmd_rules/test_topk_rule.py index a0d57ec5d6bfde..0adb2c87d140f5 100644 --- a/test/auto_parallel/spmd_rules/test_topk_rule.py +++ b/test/auto_parallel/spmd_rules/test_topk_rule.py @@ -67,11 +67,13 @@ def test_topk_forward(self): self.assertEqual(len(inferred_input_dist_attrs), 1) self.assertEqual(len(inferred_output_dist_attrs), 2) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, -1]) self.assertEqual( inferred_output_dist_attrs[0].dims_mapping, [0, -1, -1] ) + self.assertEqual( + inferred_output_dist_attrs[1].dims_mapping, [0, -1, -1] + ) def test_topk_backward(self): # axis = 1 @@ -96,8 +98,8 @@ def test_topk_backward(self): self.assertEqual(len(inferred_input_dist_attrs), 3) self.assertEqual(len(inferred_output_dist_attrs), 1) self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) - self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [0, -1, 1]) self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) From 96320b4425fbb9fdf2724d7864e02e94bd6e1787 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Tue, 20 May 2025 16:23:13 +0800 Subject: [PATCH 09/10] apply review, wow~ --- paddle/phi/infermeta/spmd_rules/topk.cc | 13 +++++++++---- paddle/phi/infermeta/spmd_rules/utils.cc | 4 ---- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/paddle/phi/infermeta/spmd_rules/topk.cc b/paddle/phi/infermeta/spmd_rules/topk.cc index fc4d3dd15818c1..1df317bba541bc 100644 --- a/paddle/phi/infermeta/spmd_rules/topk.cc +++ b/paddle/phi/infermeta/spmd_rules/topk.cc @@ -95,14 +95,19 @@ SpmdInfo TopkGradInferSpmdBase(const DistMetaTensor& x, x_axes[axis] = '1'; std::string indices_axes = x_axes; std::string out_grad_axes = x_axes; + std::vector x_dims_mapping(x_dims_mapping_src); + std::vector indices_dims_mapping(indices_dims_mapping_src); + std::vector out_grad_dims_mapping(out_grad_dims_mapping_src); + x_dims_mapping[axis] = -1; + indices_dims_mapping[axis] = -1; + out_grad_dims_mapping[axis] = -1; // Merge sharding std::pair> indices_pair( - indices_axes, indices_dims_mapping_src); + indices_axes, indices_dims_mapping); std::pair> out_grad_pair( - out_grad_axes, out_grad_dims_mapping_src); - std::pair> x_pair(x_axes, - x_dims_mapping_src); + out_grad_axes, out_grad_dims_mapping); + std::pair> x_pair(x_axes, x_dims_mapping); auto axis_to_dim_map = ShardingMergeForTensors({x_pair, indices_pair, out_grad_pair}); diff --git a/paddle/phi/infermeta/spmd_rules/utils.cc b/paddle/phi/infermeta/spmd_rules/utils.cc index 3f773c6184eccd..6d3e354b30e7cd 100644 --- a/paddle/phi/infermeta/spmd_rules/utils.cc +++ b/paddle/phi/infermeta/spmd_rules/utils.cc @@ -86,10 +86,6 @@ std::unordered_map ShardingMergeForTensors( for (auto& pair : tensor_axes_to_dim_pairs) { for (size_t i = 0; i < pair.second.size(); ++i) { auto tensor_axis = pair.first.substr(i, 1); - // Cooperate with GetDimsMappingForAxes to deal "1" for replicated. - if (tensor_axis == "1") { - continue; - } auto mesh_dim = pair.second[i]; if (axis_to_dim_map.count(tensor_axis) == 0) { From 8564a9feeb9a5af51dbbea69e32aa2679de1b964 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 21 May 2025 14:22:40 +0800 Subject: [PATCH 10/10] apply review;refine pybind function --- paddle/fluid/pybind/auto_parallel_py.cc | 4 ++++ paddle/fluid/pybind/op_function_common.h | 2 ++ .../phi/core/distributed/auto_parallel/inferspmd_utils.cc | 1 + .../phi/core/distributed/auto_parallel/inferspmd_utils.h | 1 + paddle/phi/infermeta/spmd_rules/rules.cc | 8 ++++---- paddle/phi/infermeta/spmd_rules/topk.cc | 1 - test/auto_parallel/spmd_rules/test_cummax_rule.py | 5 ++++- test/auto_parallel/spmd_rules/test_cummin_rule.py | 5 ++++- 8 files changed, 20 insertions(+), 7 deletions(-) diff --git a/paddle/fluid/pybind/auto_parallel_py.cc b/paddle/fluid/pybind/auto_parallel_py.cc index 0d10e540c81791..03db67ad8acf20 100644 --- a/paddle/fluid/pybind/auto_parallel_py.cc +++ b/paddle/fluid/pybind/auto_parallel_py.cc @@ -926,6 +926,10 @@ static void parse_attr(PyObject *obj, auto attr = CastPyArg2Float(obj, infer_spmd_string, static_cast(arg_pos)); ctx->EmplaceBackAttr(attr); + } else if (PyObject_CheckDataType(obj)) { + auto attr = CastPyArg2DataType( + obj, infer_spmd_string, static_cast(arg_pos)); + ctx->EmplaceBackAttr(attr); } else { // TODO(ljz) support other types PADDLE_THROW(common::errors::InvalidArgument( "%s(): argument (position %d) must be " diff --git a/paddle/fluid/pybind/op_function_common.h b/paddle/fluid/pybind/op_function_common.h index 527df7bcc6fa70..89b2a14a2f5001 100644 --- a/paddle/fluid/pybind/op_function_common.h +++ b/paddle/fluid/pybind/op_function_common.h @@ -44,6 +44,8 @@ bool PyObject_CheckBool(PyObject** obj); bool PyObject_CheckLong(PyObject* obj); +bool PyObject_CheckDataType(PyObject* obj); + int32_t PyObject_ToInt32(PyObject* obj); uint32_t PyObject_ToUInt32(PyObject* obj); diff --git a/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc b/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc index 73ff5a34f1b3a6..a8e7bc827822a3 100644 --- a/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc +++ b/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc @@ -66,6 +66,7 @@ AttrType InferSpmdContext::AttrAt(size_t idx) const { template float InferSpmdContext::AttrAt(size_t idx) const; template int InferSpmdContext::AttrAt(size_t idx) const; template int64_t InferSpmdContext::AttrAt(size_t idx) const; +template DataType InferSpmdContext::AttrAt(size_t idx) const; template <> bool InferSpmdContext::AttrAt(size_t idx) const { diff --git a/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h b/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h index 5679cf512991a4..a74e52a39a24a8 100644 --- a/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h +++ b/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h @@ -178,6 +178,7 @@ struct InferSpmdFnImpl { PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(int); PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(float); PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(int64_t); + PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(DataType); PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF(std::vector); PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF( std::vector); diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc index b93adae13d94df..9b012ee3b2a44e 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.cc +++ b/paddle/phi/infermeta/spmd_rules/rules.cc @@ -755,12 +755,12 @@ PD_REGISTER_SPMD_RULE(add_n, PD_INFER_SPMD(phi::distributed::AddNInferSpmd)); // cummax PD_REGISTER_SPMD_RULE(cummax, - PD_INFER_SPMD(phi::distributed::TopkInferSpmdBase), - PD_INFER_SPMD(phi::distributed::TopkGradInferSpmdBase)); + PD_INFER_SPMD(phi::distributed::CummaxInferSpmd), + PD_INFER_SPMD(phi::distributed::CummaxGradInferSpmd)); // cummin PD_REGISTER_SPMD_RULE(cummin, - PD_INFER_SPMD(phi::distributed::TopkInferSpmdBase), - PD_INFER_SPMD(phi::distributed::TopkGradInferSpmdBase)); + PD_INFER_SPMD(phi::distributed::CumminInferSpmd), + PD_INFER_SPMD(phi::distributed::CumminGradInferSpmd)); // argsort PD_REGISTER_SPMD_RULE(argsort, diff --git a/paddle/phi/infermeta/spmd_rules/topk.cc b/paddle/phi/infermeta/spmd_rules/topk.cc index 1df317bba541bc..6b4123437b656c 100644 --- a/paddle/phi/infermeta/spmd_rules/topk.cc +++ b/paddle/phi/infermeta/spmd_rules/topk.cc @@ -92,7 +92,6 @@ SpmdInfo TopkGradInferSpmdBase(const DistMetaTensor& x, // Build einsum notation std::string alphabet = "abcdefghijlopqrstuvwxyz"; std::string x_axes = alphabet.substr(0, x_ndim); - x_axes[axis] = '1'; std::string indices_axes = x_axes; std::string out_grad_axes = x_axes; std::vector x_dims_mapping(x_dims_mapping_src); diff --git a/test/auto_parallel/spmd_rules/test_cummax_rule.py b/test/auto_parallel/spmd_rules/test_cummax_rule.py index f7506888cc0096..96338c5451d62c 100644 --- a/test/auto_parallel/spmd_rules/test_cummax_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummax_rule.py @@ -20,7 +20,7 @@ TensorDistAttr, ) from paddle.distributed.fleet import auto -from paddle.framework import core +from paddle.framework import convert_np_dtype_to_dtype_, core class TestCummaxSPMDRule(unittest.TestCase): @@ -43,6 +43,7 @@ def setUp(self): self.rule = core.get_phi_spmd_rule("cummax") self.attrs = OrderedDict() self.attrs['axis'] = 1 + self.attrs['dtype'] = convert_np_dtype_to_dtype_("int64") def test_cummax_forward(self): # axis = 1 @@ -52,6 +53,7 @@ def test_cummax_forward(self): result_dist_attrs = self.rule.infer_forward( self.x_dist_tensor_spec, self.attrs['axis'], + self.attrs['dtype'], ) self.assertEqual(len(result_dist_attrs), 2) @@ -81,6 +83,7 @@ def test_cummax_backward(self): self.x_dist_tensor_spec, self.out_dist_tensor_spec, self.attrs['axis'], + self.attrs['dtype'], ) self.assertEqual(len(result_dist_attrs), 2) diff --git a/test/auto_parallel/spmd_rules/test_cummin_rule.py b/test/auto_parallel/spmd_rules/test_cummin_rule.py index 441f5857ae204f..b15ebfad7bfc4b 100644 --- a/test/auto_parallel/spmd_rules/test_cummin_rule.py +++ b/test/auto_parallel/spmd_rules/test_cummin_rule.py @@ -20,7 +20,7 @@ TensorDistAttr, ) from paddle.distributed.fleet import auto -from paddle.framework import core +from paddle.framework import convert_np_dtype_to_dtype_, core class TestCumminSPMDRule(unittest.TestCase): @@ -43,6 +43,7 @@ def setUp(self): self.rule = core.get_phi_spmd_rule("cummin") self.attrs = OrderedDict() self.attrs['axis'] = 1 + self.attrs['dtype'] = convert_np_dtype_to_dtype_("int64") def test_cummin_forward(self): # axis = 1 @@ -52,6 +53,7 @@ def test_cummin_forward(self): result_dist_attrs = self.rule.infer_forward( self.x_dist_tensor_spec, self.attrs['axis'], + self.attrs['dtype'], ) self.assertEqual(len(result_dist_attrs), 2) @@ -81,6 +83,7 @@ def test_cummin_backward(self): self.x_dist_tensor_spec, self.out_dist_tensor_spec, self.attrs['axis'], + self.attrs['dtype'], ) self.assertEqual(len(result_dist_attrs), 2)