From 49fdb3604c5327264356b2ab92c82b155c9789f8 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 21 May 2025 17:39:02 +0800 Subject: [PATCH 1/9] add spmd_rule for label_smooth --- paddle/phi/infermeta/spmd_rules/expand_as.cc | 6 +- paddle/phi/infermeta/spmd_rules/expand_as.h | 6 +- .../phi/infermeta/spmd_rules/label_smooth.cc | 44 +++++++ .../phi/infermeta/spmd_rules/label_smooth.h | 31 +++++ paddle/phi/infermeta/spmd_rules/rules.cc | 4 +- paddle/phi/infermeta/spmd_rules/rules.h | 1 + paddle/phi/ops/yaml/backward.yaml | 2 + paddle/phi/ops/yaml/ops.yaml | 2 + test/auto_parallel/spmd_rules/CMakeLists.txt | 1 + .../spmd_rules/test_label_smooth_rule.py | 117 ++++++++++++++++++ .../auto_parallel/expand_as_spmd_rule_test.cc | 2 +- test/cpp/auto_parallel/spmd_rule_test.cc | 60 +++++++++ 12 files changed, 267 insertions(+), 9 deletions(-) create mode 100644 paddle/phi/infermeta/spmd_rules/label_smooth.cc create mode 100644 paddle/phi/infermeta/spmd_rules/label_smooth.h create mode 100644 test/auto_parallel/spmd_rules/test_label_smooth_rule.py diff --git a/paddle/phi/infermeta/spmd_rules/expand_as.cc b/paddle/phi/infermeta/spmd_rules/expand_as.cc index 9600ff0c0e57b4..744c89e40cf9e8 100644 --- a/paddle/phi/infermeta/spmd_rules/expand_as.cc +++ b/paddle/phi/infermeta/spmd_rules/expand_as.cc @@ -44,7 +44,7 @@ std::tuple AlignExpandAsDistAttrs( SpmdInfo ExpandAsInferSpmd(const DistMetaTensor& x, const DistMetaTensor& y, - const std::vector& target_shape) { + const std::vector& target_shape) { auto [x_dist_attr, y_dist_attr] = AlignExpandAsDistAttrs(x, y); return {{x_dist_attr, y_dist_attr}, {y_dist_attr}}; } @@ -52,14 +52,14 @@ SpmdInfo ExpandAsInferSpmd(const DistMetaTensor& x, SpmdInfo ExpandAsInferSpmdReverse(const DistMetaTensor& x, const DistMetaTensor& y, const DistMetaTensor& output, - const std::vector& target_shape) { + const std::vector& target_shape) { auto [x_dist_attr, y_dist_attr] = AlignExpandAsDistAttrs(x, output); return {{x_dist_attr, y_dist_attr}, {y_dist_attr}}; } SpmdInfo ExpandAsGradInferSpmd(const DistMetaTensor& x, const DistMetaTensor& out_grad, - const std::vector& target_shape) { + const std::vector& target_shape) { auto [x_dist_attr, y_dist_attr] = AlignExpandAsDistAttrs(x, out_grad); const auto& x_dims_mapping = x_dist_attr.dims_mapping(); const auto& y_dims_mapping = y_dist_attr.dims_mapping(); diff --git a/paddle/phi/infermeta/spmd_rules/expand_as.h b/paddle/phi/infermeta/spmd_rules/expand_as.h index 67cc6f3853dc16..96e5be0e1e4aa2 100644 --- a/paddle/phi/infermeta/spmd_rules/expand_as.h +++ b/paddle/phi/infermeta/spmd_rules/expand_as.h @@ -23,16 +23,16 @@ namespace phi { namespace distributed { SpmdInfo ExpandAsInferSpmd(const DistMetaTensor& x, const DistMetaTensor& y, - const std::vector& target_shape); + const std::vector& target_shape); SpmdInfo ExpandAsInferSpmdReverse(const DistMetaTensor& x, const DistMetaTensor& y, const DistMetaTensor& output, - const std::vector& target_shape); + const std::vector& target_shape); SpmdInfo ExpandAsGradInferSpmd(const DistMetaTensor& x, const DistMetaTensor& out_grad, - const std::vector& target_shape); + const std::vector& target_shape); } // namespace distributed } // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/label_smooth.cc b/paddle/phi/infermeta/spmd_rules/label_smooth.cc new file mode 100644 index 00000000000000..8c2161a4ebf8b6 --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/label_smooth.cc @@ -0,0 +1,44 @@ +/* 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/label_smooth.h" +#include "glog/logging.h" +#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h" +#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h" +#include "paddle/phi/core/distributed/type_defs.h" +#include "paddle/phi/infermeta/spmd_rules/elementwise.h" +#include "paddle/phi/infermeta/spmd_rules/spmd_rule_macro_define.h" +#include "paddle/phi/infermeta/spmd_rules/utils.h" + +namespace phi { +namespace distributed { +SpmdInfo LabelSmoothInferSpmd(const DistMetaTensor& label, + const DistMetaTensor& prior_dist, + float epsilon) { + if (prior_dist.initialized()) { + return ElementwiseBinaryInferSpmd(label, prior_dist); + } + SpmdInfo unary_spmd_info = ElementwiseUnaryInferSpmd(label); + unary_spmd_info.first.push_back(TensorDistAttr()); + VLOG(4) << "LabelSmoothInferSpmd:"; + return unary_spmd_info; +} + +SpmdInfo LabelSmoothGradInferSpmd(const DistMetaTensor& out_grad, + float epsilon) { + VLOG(4) << "LabelSmoothGradInferSpmd:"; + return ElementwiseUnaryGradInferSpmd(out_grad); +} +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/label_smooth.h b/paddle/phi/infermeta/spmd_rules/label_smooth.h new file mode 100644 index 00000000000000..ccb7195c125616 --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/label_smooth.h @@ -0,0 +1,31 @@ +/* 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 LabelSmoothInferSpmd(const DistMetaTensor& label, + const DistMetaTensor& prior_dist, + float epsilon); + +SpmdInfo LabelSmoothGradInferSpmd(const DistMetaTensor& out_grad, + float epsilon); + +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc index d2b3277bb9113f..e48f1cc1c4ae7d 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.cc +++ b/paddle/phi/infermeta/spmd_rules/rules.cc @@ -218,8 +218,8 @@ PD_REGISTER_SPMD_RULE( PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmdReverse)); PD_REGISTER_SPMD_RULE( label_smooth, - PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmd), - PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmdReverse)); + PD_INFER_SPMD(phi::distributed::LabelSmoothInferSpmd), + PD_INFER_SPMD(phi::distributed::LabelSmoothGradInferSpmd)); PD_REGISTER_SPMD_RULE( leaky_relu, PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmd), diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h index 24f37017967f0d..8de61f83389d82 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.h +++ b/paddle/phi/infermeta/spmd_rules/rules.h @@ -43,6 +43,7 @@ limitations under the License. */ #include "paddle/phi/infermeta/spmd_rules/fused_rope.h" #include "paddle/phi/infermeta/spmd_rules/gather.h" #include "paddle/phi/infermeta/spmd_rules/gather_nd.h" +#include "paddle/phi/infermeta/spmd_rules/label_smooth.h" #include "paddle/phi/infermeta/spmd_rules/layer_norm.h" #include "paddle/phi/infermeta/spmd_rules/logsumexp.h" #include "paddle/phi/infermeta/spmd_rules/matmul.h" diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml index 3a48214c8e1579..9d59aa5c084ca1 100644 --- a/paddle/phi/ops/yaml/backward.yaml +++ b/paddle/phi/ops/yaml/backward.yaml @@ -996,6 +996,7 @@ infer_meta : func : UnchangedInferMeta param : [x] + spmd_rule : ExpandAsGradInferSpmd kernel : func : expand_as_grad no_need_buffer : x @@ -1826,6 +1827,7 @@ infer_meta : func : UnchangedInferMeta param : [out_grad] + spmd_rule : LabelSmoothGradInferSpmd kernel : func : label_smooth_grad diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml index 5bd4eca9a850ac..2db4ebec5eb177 100755 --- a/paddle/phi/ops/yaml/ops.yaml +++ b/paddle/phi/ops/yaml/ops.yaml @@ -1784,6 +1784,7 @@ infer_meta : func : ExpandAsInferMeta local_shape: out + spmd_rule : ExpandAsInferSpmd kernel : func : expand_as data_type : x @@ -2928,6 +2929,7 @@ infer_meta : func : UnchangedInferMeta param : [label] + spmd_rule : LabelSmoothInferSpmd kernel : func : label_smooth data_type : label diff --git a/test/auto_parallel/spmd_rules/CMakeLists.txt b/test/auto_parallel/spmd_rules/CMakeLists.txt index 0bed8bba3a5573..3622c4b1c24bad 100644 --- a/test/auto_parallel/spmd_rules/CMakeLists.txt +++ b/test/auto_parallel/spmd_rules/CMakeLists.txt @@ -49,6 +49,7 @@ if(WITH_DISTRIBUTE) 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_argsort_rule MODULES test_argsort_rule) + py_test_modules(test_label_smooth_rule MODULES test_label_smooth_rule) endif() # End of unittests WITH single card WITHOUT timeout diff --git a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py new file mode 100644 index 00000000000000..4f09a9fc556476 --- /dev/null +++ b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py @@ -0,0 +1,117 @@ +# 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 TestLabelSmoothSPMDRule(unittest.TestCase): + def setUp(self): + label_shape = [16, 16, 16] + out_shape = [16, 16, 16] + prior_dist_shape = [1, 16] + process_mesh = auto.ProcessMesh(mesh=[[0, 1], [2, 3]]) + + label_tensor_dist_attr = TensorDistAttr() + label_tensor_dist_attr.dims_mapping = [-1, -1, -1] + label_tensor_dist_attr.process_mesh = process_mesh + + self.label_dist_tensor_spec = DistTensorSpec( + label_shape, label_tensor_dist_attr + ) + prior_dist_tensor_dist_attr = TensorDistAttr() + prior_dist_tensor_dist_attr.dims_mapping = [-1, -1] + prior_dist_tensor_dist_attr.process_mesh = process_mesh + + self.prior_dist_dist_tensor_spec = DistTensorSpec( + prior_dist_shape, prior_dist_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, out_tensor_dist_attr + ) + + self.rule = core.get_phi_spmd_rule("label_smooth") + self.attrs = OrderedDict() + self.attrs['epsilon'] = 0.1 + + def test_topk_forward(self): + # [0, 1, -1], [Fake] --> [0, 1, -1], [Fake], [0, 1, -1] + self.label_dist_tensor_spec.set_dims_mapping([0, 1, -1]) + + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + DistTensorSpec(), + self.attrs['epsilon'], + ) + + 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), 2) + 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[1], TensorDistAttr()) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, -1] --> [0, 1, -1], [-1, -1], [0, 1, -1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, -1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + + 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), 2) + 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[1].dims_mapping, [-1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + def test_topk_backward(self): + # [0, -1, 1] --> [0, -1, 1], [0, -1, 1] + self.out_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + result_dist_attrs = self.rule.infer_backward( + self.out_dist_tensor_spec, + self.attrs['epsilon'], + ) + + 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), 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/expand_as_spmd_rule_test.cc b/test/cpp/auto_parallel/expand_as_spmd_rule_test.cc index ca9daa84f99fdf..79fdd939b9d7d4 100644 --- a/test/cpp/auto_parallel/expand_as_spmd_rule_test.cc +++ b/test/cpp/auto_parallel/expand_as_spmd_rule_test.cc @@ -21,7 +21,7 @@ namespace auto_parallel { TEST(ExpandAsInferSpmd, Ctor) { // Sharding along axes besides softmax axis. std::vector x_shape = {1, 48}; - std::vector y_shape = {2, 32, 48}; + std::vector y_shape = {2, 32, 48}; std::vector mesh_shape = {2, 3}; std::vector process_ids = {0, 1, 2, 3, 4, 5}; diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc index 54e2f2a7c98cb7..6499b11fc28c78 100644 --- a/test/cpp/auto_parallel/spmd_rule_test.cc +++ b/test/cpp/auto_parallel/spmd_rule_test.cc @@ -2653,6 +2653,66 @@ TEST(Topk, Ctor) { check_dim_mapping(backward_info.second[0], {0, -1, 1}); } +TEST(LabelSmooth, 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); + float epsilon = 0.1; + + // test forward + // [0, 1, -1], [Fake] --> [0, 1, -1], [Fake], [0, 1, -1] + auto label_dist_attr = TensorDistAttr(); + label_dist_attr.set_process_mesh(process_mesh); + label_dist_attr.set_dims_mapping({0, 1, -1}); + label_dist_attr.set_dynamic_dims({false, false, false}); + + auto prior_dist_dist_attr = TensorDistAttr(); + prior_dist_dist_attr.set_process_mesh(process_mesh); + prior_dist_dist_attr.set_dims_mapping({-1, -1}); + prior_dist_dist_attr.set_dynamic_dims({false, false}); + + phi::distributed::DistMetaTensor label = phi::distributed::DistMetaTensor( + common::make_ddim({16, 16, 16}), label_dist_attr); + phi::distributed::DistMetaTensor prior_dist = + phi::distributed::DistMetaTensor(common::make_ddim({1, 16}), + prior_dist_dist_attr); + phi::distributed::SpmdInfo forward_info = + phi::distributed::LabelSmoothInferSpmd( + label, phi::distributed::DistMetaTensor(), epsilon); + + EXPECT_EQ(forward_info.first.size(), 2UL); + EXPECT_EQ(forward_info.second.size(), 1UL); + check_dim_mapping(forward_info.first[0], {0, 1, -1}); + EXPECT_EQ(forward_info.first[1], phi::distributed::DistMetaTensor()); + check_dim_mapping(forward_info.second[0], {0, 1, -1}); + + // shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, -1] --> [0, 1, -1], [-1, + // -1], [0, 1, -1] + forward_info = + phi::distributed::LabelSmoothInferSpmd(label, prior_dist, epsilon); + EXPECT_EQ(forward_info.first.size(), 2UL); + EXPECT_EQ(forward_info.second.size(), 1UL); + check_dim_mapping(forward_info.first[0], {0, 1, -1}); + check_dim_mapping(forward_info.first[1], {-1, -1}); + check_dim_mapping(forward_info.second[0], {0, 1, -1}); + + // test backward + // [0, -1, 1] --> [0, -1, 1], [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({0, -1, 1}); + out_grad_dist_attr.set_dynamic_dims({false, false, false}); + phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor( + common::make_ddim({16, 16, 16}), out_grad_dist_attr); + phi::distributed::SpmdInfo backward_info = + phi::distributed::LabelSmoothGradInferSpmd(out_grad, epsilon); + EXPECT_EQ(backward_info.first.size(), 1UL); + EXPECT_EQ(backward_info.second.size(), 1UL); + check_dim_mapping(backward_info.first[0], {0, -1, 1}); + check_dim_mapping(backward_info.second[0], {0, -1, 1}); +} + TEST(ArgSortGradInferSpmd, Ctor) { // Sharding along axes besides argsort axis. std::vector x_shape = {16, 32, 48}; From ed9f3aa0c64095610e157ed0158efcd9e28c7ea8 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Fri, 23 May 2025 01:27:40 +0800 Subject: [PATCH 2/9] update --- test/cpp/auto_parallel/spmd_rule_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc index 6499b11fc28c78..820b99409f452d 100644 --- a/test/cpp/auto_parallel/spmd_rule_test.cc +++ b/test/cpp/auto_parallel/spmd_rule_test.cc @@ -2684,7 +2684,7 @@ TEST(LabelSmooth, Ctor) { EXPECT_EQ(forward_info.first.size(), 2UL); EXPECT_EQ(forward_info.second.size(), 1UL); check_dim_mapping(forward_info.first[0], {0, 1, -1}); - EXPECT_EQ(forward_info.first[1], phi::distributed::DistMetaTensor()); + EXPECT_EQ(forward_info.first[1], phi::distributed::TensorDistAttr()); check_dim_mapping(forward_info.second[0], {0, 1, -1}); // shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, -1] --> [0, 1, -1], [-1, From e630d42c67cc0ce2fb6b8a5a3d8bd0b6770ac77d Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Sun, 25 May 2025 20:54:10 +0800 Subject: [PATCH 3/9] fix bugs --- test/cpp/auto_parallel/spmd_rule_test.cc | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc index e40420e4c6b365..63a01d4402cbbf 100644 --- a/test/cpp/auto_parallel/spmd_rule_test.cc +++ b/test/cpp/auto_parallel/spmd_rule_test.cc @@ -2971,7 +2971,13 @@ TEST(LabelSmooth, Ctor) { EXPECT_EQ(forward_info.first.size(), 2UL); EXPECT_EQ(forward_info.second.size(), 1UL); check_dim_mapping(forward_info.first[0], {0, 1, -1}); - EXPECT_EQ(forward_info.first[1], phi::distributed::TensorDistAttr()); + const phi::distributed::ArgDistAttr& attr = forward_info.first[1]; + if (paddle::holds_alternative(attr)) { + EXPECT_EQ(paddle::get(attr), + phi::distributed::TensorDistAttr()); + } else { + FAIL() << "forward_info.first[1] is not TensorDistAttr"; + } check_dim_mapping(forward_info.second[0], {0, 1, -1}); // shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, -1] --> [0, 1, -1], [-1, From bedf595ed329f4547f69de48e76e467c19cc3437 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 4 Jun 2025 20:58:41 +0800 Subject: [PATCH 4/9] apply review --- test/auto_parallel/spmd_rules/test_label_smooth_rule.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py index 4f09a9fc556476..24a69aebec5ddf 100644 --- a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py +++ b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py @@ -56,7 +56,7 @@ def setUp(self): self.attrs = OrderedDict() self.attrs['epsilon'] = 0.1 - def test_topk_forward(self): + def test_label_smooth_forward(self): # [0, 1, -1], [Fake] --> [0, 1, -1], [Fake], [0, 1, -1] self.label_dist_tensor_spec.set_dims_mapping([0, 1, -1]) @@ -96,7 +96,7 @@ def test_topk_forward(self): self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1]) self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) - def test_topk_backward(self): + def test_label_smooth_backward(self): # [0, -1, 1] --> [0, -1, 1], [0, -1, 1] self.out_dist_tensor_spec.set_dims_mapping([0, -1, 1]) result_dist_attrs = self.rule.infer_backward( From 0d542270e8990e699cb0daaa075edc12725b513c Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 4 Jun 2025 23:09:57 +0800 Subject: [PATCH 5/9] revert due to https://github.com/PaddlePaddle/Paddle/pull/72753 --- paddle/phi/infermeta/spmd_rules/expand_as.cc | 6 +++--- paddle/phi/infermeta/spmd_rules/expand_as.h | 6 +++--- test/cpp/auto_parallel/expand_as_spmd_rule_test.cc | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/paddle/phi/infermeta/spmd_rules/expand_as.cc b/paddle/phi/infermeta/spmd_rules/expand_as.cc index 744c89e40cf9e8..9600ff0c0e57b4 100644 --- a/paddle/phi/infermeta/spmd_rules/expand_as.cc +++ b/paddle/phi/infermeta/spmd_rules/expand_as.cc @@ -44,7 +44,7 @@ std::tuple AlignExpandAsDistAttrs( SpmdInfo ExpandAsInferSpmd(const DistMetaTensor& x, const DistMetaTensor& y, - const std::vector& target_shape) { + const std::vector& target_shape) { auto [x_dist_attr, y_dist_attr] = AlignExpandAsDistAttrs(x, y); return {{x_dist_attr, y_dist_attr}, {y_dist_attr}}; } @@ -52,14 +52,14 @@ SpmdInfo ExpandAsInferSpmd(const DistMetaTensor& x, SpmdInfo ExpandAsInferSpmdReverse(const DistMetaTensor& x, const DistMetaTensor& y, const DistMetaTensor& output, - const std::vector& target_shape) { + const std::vector& target_shape) { auto [x_dist_attr, y_dist_attr] = AlignExpandAsDistAttrs(x, output); return {{x_dist_attr, y_dist_attr}, {y_dist_attr}}; } SpmdInfo ExpandAsGradInferSpmd(const DistMetaTensor& x, const DistMetaTensor& out_grad, - const std::vector& target_shape) { + const std::vector& target_shape) { auto [x_dist_attr, y_dist_attr] = AlignExpandAsDistAttrs(x, out_grad); const auto& x_dims_mapping = x_dist_attr.dims_mapping(); const auto& y_dims_mapping = y_dist_attr.dims_mapping(); diff --git a/paddle/phi/infermeta/spmd_rules/expand_as.h b/paddle/phi/infermeta/spmd_rules/expand_as.h index 96e5be0e1e4aa2..67cc6f3853dc16 100644 --- a/paddle/phi/infermeta/spmd_rules/expand_as.h +++ b/paddle/phi/infermeta/spmd_rules/expand_as.h @@ -23,16 +23,16 @@ namespace phi { namespace distributed { SpmdInfo ExpandAsInferSpmd(const DistMetaTensor& x, const DistMetaTensor& y, - const std::vector& target_shape); + const std::vector& target_shape); SpmdInfo ExpandAsInferSpmdReverse(const DistMetaTensor& x, const DistMetaTensor& y, const DistMetaTensor& output, - const std::vector& target_shape); + const std::vector& target_shape); SpmdInfo ExpandAsGradInferSpmd(const DistMetaTensor& x, const DistMetaTensor& out_grad, - const std::vector& target_shape); + const std::vector& target_shape); } // namespace distributed } // namespace phi diff --git a/test/cpp/auto_parallel/expand_as_spmd_rule_test.cc b/test/cpp/auto_parallel/expand_as_spmd_rule_test.cc index 79fdd939b9d7d4..ca9daa84f99fdf 100644 --- a/test/cpp/auto_parallel/expand_as_spmd_rule_test.cc +++ b/test/cpp/auto_parallel/expand_as_spmd_rule_test.cc @@ -21,7 +21,7 @@ namespace auto_parallel { TEST(ExpandAsInferSpmd, Ctor) { // Sharding along axes besides softmax axis. std::vector x_shape = {1, 48}; - std::vector y_shape = {2, 32, 48}; + std::vector y_shape = {2, 32, 48}; std::vector mesh_shape = {2, 3}; std::vector process_ids = {0, 1, 2, 3, 4, 5}; From eaf3c5cc3670c46d41a915837234323577999383 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Thu, 5 Jun 2025 13:13:49 +0800 Subject: [PATCH 6/9] revert change for expand_as --- paddle/phi/ops/yaml/backward.yaml | 1 - paddle/phi/ops/yaml/ops.yaml | 1 - 2 files changed, 2 deletions(-) diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml index 55c78011d2a0f9..3064abcd614164 100644 --- a/paddle/phi/ops/yaml/backward.yaml +++ b/paddle/phi/ops/yaml/backward.yaml @@ -998,7 +998,6 @@ infer_meta : func : UnchangedInferMeta param : [x] - spmd_rule : ExpandAsGradInferSpmd kernel : func : expand_as_grad no_need_buffer : x diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml index e89ed24683a56b..805cfb6619f2e6 100755 --- a/paddle/phi/ops/yaml/ops.yaml +++ b/paddle/phi/ops/yaml/ops.yaml @@ -1786,7 +1786,6 @@ infer_meta : func : ExpandAsInferMeta local_shape: out - spmd_rule : ExpandAsInferSpmd kernel : func : expand_as data_type : x From 624db849bb5dd4702c2b43229778bf536bfeb25c Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Sat, 7 Jun 2025 15:29:59 +0800 Subject: [PATCH 7/9] Add more tests --- .../spmd_rules/test_label_smooth_rule.py | 51 ++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) diff --git a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py index 24a69aebec5ddf..e984d6038a77f6 100644 --- a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py +++ b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py @@ -96,8 +96,57 @@ def test_label_smooth_forward(self): self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1]) self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, 1] --> [0, 1, -1], [-1, -1], [0, 1, -1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + 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), 2) + 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[1].dims_mapping, [-1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, 0] --> [0, 1, -1], [-1, -1], [0, 1, -1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 0]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + 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), 2) + 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[1].dims_mapping, [-1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + # shape: [16, 16, 16], [16, 16, 16]. [0, 1, -1], [0, -1, 1] --> [0, 1, -1], [0, 1, -1], [0, 1, -1] + self.prior_dist_dist_tensor_spec.shape = [16, 16, 16] + self.prior_dist_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + 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), 2) + 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[1].dims_mapping, [0, 1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + def test_label_smooth_backward(self): - # [0, -1, 1] --> [0, -1, 1], [0, -1, 1] + # shape: [16, 16 ,16]. [0, -1, 1] --> [0, -1, 1], [0, -1, 1] self.out_dist_tensor_spec.set_dims_mapping([0, -1, 1]) result_dist_attrs = self.rule.infer_backward( self.out_dist_tensor_spec, From 6cb63c25f75cab5af6e7ccc90ef5f800824aea73 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Wed, 11 Jun 2025 23:49:28 +0800 Subject: [PATCH 8/9] apply review --- .../phi/infermeta/spmd_rules/label_smooth.cc | 5 +- .../spmd_rules/test_label_smooth_rule.py | 47 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/paddle/phi/infermeta/spmd_rules/label_smooth.cc b/paddle/phi/infermeta/spmd_rules/label_smooth.cc index 8c2161a4ebf8b6..07de29aff9fd0a 100644 --- a/paddle/phi/infermeta/spmd_rules/label_smooth.cc +++ b/paddle/phi/infermeta/spmd_rules/label_smooth.cc @@ -27,17 +27,18 @@ SpmdInfo LabelSmoothInferSpmd(const DistMetaTensor& label, const DistMetaTensor& prior_dist, float epsilon) { if (prior_dist.initialized()) { + VLOG(4) << "LabelSmoothInferSpmd call ElementwiseBinaryInferSpmd:"; return ElementwiseBinaryInferSpmd(label, prior_dist); } + VLOG(4) << "LabelSmoothInferSpmd call ElementwiseUnaryInferSpmd:"; SpmdInfo unary_spmd_info = ElementwiseUnaryInferSpmd(label); unary_spmd_info.first.push_back(TensorDistAttr()); - VLOG(4) << "LabelSmoothInferSpmd:"; return unary_spmd_info; } SpmdInfo LabelSmoothGradInferSpmd(const DistMetaTensor& out_grad, float epsilon) { - VLOG(4) << "LabelSmoothGradInferSpmd:"; + VLOG(4) << "LabelSmoothGradInferSpmd call ElementwiseUnaryGradInferSpmd:"; return ElementwiseUnaryGradInferSpmd(out_grad); } } // namespace distributed diff --git a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py index e984d6038a77f6..62abea990d0437 100644 --- a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py +++ b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py @@ -128,7 +128,54 @@ def test_label_smooth_forward(self): self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1]) self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 1] --> [0, -1, 1], [-1, 1], [0, -1, 1] + self.label_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + 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), 2) + 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[1].dims_mapping, [-1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + + # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 0]: + # Because reuse the ElementwiseBinaryInferSpmd, it will cause error now when meet 1 and 0 with same axes. + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 0]) + with self.assertRaises(NotImplementedError) as cm: + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + self.assertIn( + "is Sharded by two different mesh dimension", str(cm.exception) + ) + + # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, -1] --> [0, -1, 1], [-1, 1], [0, -1, 1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, -1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + 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), 2) + 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[1].dims_mapping, [-1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + # shape: [16, 16, 16], [16, 16, 16]. [0, 1, -1], [0, -1, 1] --> [0, 1, -1], [0, 1, -1], [0, 1, -1] + self.label_dist_tensor_spec.set_dims_mapping([0, 1, -1]) self.prior_dist_dist_tensor_spec.shape = [16, 16, 16] self.prior_dist_dist_tensor_spec.set_dims_mapping([0, -1, 1]) result_dist_attrs = self.rule.infer_forward( From 96dd4a4732d47d61a28702bc489507564f1c1318 Mon Sep 17 00:00:00 2001 From: ooooo <3164076421@qq.com> Date: Thu, 12 Jun 2025 20:45:18 +0800 Subject: [PATCH 9/9] refine code --- .../phi/infermeta/spmd_rules/label_smooth.cc | 40 +++++++++++++++++++ .../spmd_rules/test_label_smooth_rule.py | 23 ++++++----- 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/paddle/phi/infermeta/spmd_rules/label_smooth.cc b/paddle/phi/infermeta/spmd_rules/label_smooth.cc index 07de29aff9fd0a..3146e27b98012b 100644 --- a/paddle/phi/infermeta/spmd_rules/label_smooth.cc +++ b/paddle/phi/infermeta/spmd_rules/label_smooth.cc @@ -23,10 +23,50 @@ limitations under the License. */ namespace phi { namespace distributed { + +std::vector ReplicateWithExcludeAxes( + const std::vector& x, const std::vector& exclude_axes) { + std::vector res(x); + for (size_t i = 0; i < x.size(); ++i) { + if (res[i] == -1) { + continue; + } + + if (std::find(exclude_axes.begin(), exclude_axes.end(), res[i]) != + exclude_axes.end()) { + res[i] = -1; + } + } + return res; +} + SpmdInfo LabelSmoothInferSpmd(const DistMetaTensor& label, const DistMetaTensor& prior_dist, float epsilon) { if (prior_dist.initialized()) { + const auto& label_dims_attr = label.dist_attr(); + const auto& prior_dims_attr = prior_dist.dist_attr(); + const auto& label_dims_mapping = label_dims_attr.dims_mapping(); + const auto& prior_dims_mapping = prior_dims_attr.dims_mapping(); + size_t label_rank = label.dims().size(); + size_t prior_rank = prior_dist.dims().size(); + + if (label_rank > prior_rank) { + size_t extra_dims_count = label_rank - prior_rank; + std::vector extra_dims_mapping_values; + extra_dims_mapping_values.reserve(extra_dims_count); + for (size_t i = 0; i < extra_dims_count; ++i) { + extra_dims_mapping_values.push_back(label_dims_mapping[i]); + } + const auto& new_prior_dims_mapping = ReplicateWithExcludeAxes( + prior_dims_mapping, extra_dims_mapping_values); + TensorDistAttr new_prior_dist_attr = prior_dims_attr; + new_prior_dist_attr.set_dims_mapping(new_prior_dims_mapping); + DistMetaTensor modified_prior_dist(prior_dist.dims(), + new_prior_dist_attr); + VLOG(4) << "LabelSmoothInferSpmd call ElementwiseBinaryInferSpmd:"; + return ElementwiseBinaryInferSpmd(label, modified_prior_dist); + } VLOG(4) << "LabelSmoothInferSpmd call ElementwiseBinaryInferSpmd:"; return ElementwiseBinaryInferSpmd(label, prior_dist); } diff --git a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py index 62abea990d0437..ce995e0b1a6781 100644 --- a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py +++ b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py @@ -145,18 +145,21 @@ def test_label_smooth_forward(self): self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1]) self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) - # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 0]: - # Because reuse the ElementwiseBinaryInferSpmd, it will cause error now when meet 1 and 0 with same axes. + # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 0] --> [0, -1, 1], [-1, 1], [0, -1, 1] self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 0]) - with self.assertRaises(NotImplementedError) as cm: - result_dist_attrs = self.rule.infer_forward( - self.label_dist_tensor_spec, - self.prior_dist_dist_tensor_spec, - self.attrs['epsilon'], - ) - self.assertIn( - "is Sharded by two different mesh dimension", str(cm.exception) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], ) + 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), 2) + 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[1].dims_mapping, [-1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, -1] --> [0, -1, 1], [-1, 1], [0, -1, 1] self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, -1])