chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,96 @@
|
||||
if(WIN32)
|
||||
cc_test(
|
||||
device_mesh_test
|
||||
SRCS device_mesh_test.cc
|
||||
DEPS type_info)
|
||||
cc_test(
|
||||
process_mesh_test
|
||||
SRCS process_mesh_test.cc
|
||||
DEPS type_info)
|
||||
else()
|
||||
cc_test(device_mesh_test SRCS device_mesh_test.cc)
|
||||
cc_test(process_mesh_test SRCS process_mesh_test.cc)
|
||||
endif()
|
||||
|
||||
cc_test(
|
||||
dist_attr_test
|
||||
SRCS dist_attr_test.cc
|
||||
DEPS proto_desc)
|
||||
|
||||
if(WITH_DISTRIBUTE)
|
||||
cc_library(
|
||||
spmd_rule_test_util
|
||||
SRCS spmd_rule_test_util.cc
|
||||
DEPS gtest)
|
||||
cc_test(
|
||||
dist_tensor_test
|
||||
SRCS dist_tensor_test.cc
|
||||
DEPS phi common)
|
||||
|
||||
paddle_test(spmd_rule_test SRCS spmd_rule_test.cc DEPS spmd_rule_test_util
|
||||
phi)
|
||||
|
||||
paddle_test(softmax_grad_spmd_rule_test SRCS softmax_grad_spmd_rule_test.cc
|
||||
DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(tile_spmd_rule_test SRCS tile_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
paddle_test(tile_co_shard_spmd_rule_test SRCS tile_co_shard_spmd_rule_test.cc
|
||||
DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(
|
||||
fused_linear_param_grad_add_spmd_rule_test SRCS
|
||||
fused_linear_param_grad_add_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(
|
||||
cross_entropy_softmax_spmd_rule_test SRCS
|
||||
cross_entropy_softmax_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(expand_spmd_rule_test SRCS expand_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(expand_as_spmd_rule_test SRCS expand_as_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(matmul_co_shard_spmd_rule_test SRCS
|
||||
matmul_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(custom_op_spmd_rule_test SRCS custom_op_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(fused_rms_norm_spmd_rule_test SRCS
|
||||
fused_rms_norm_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
paddle_test(moe_gate_dispatch_spmd_rule_test SRCS
|
||||
moe_gate_dispatch_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(moe_combine_spmd_rule_test SRCS moe_combine_spmd_rule_test.cc
|
||||
DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(softmax_co_shard_spmd_rule_test SRCS
|
||||
softmax_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(
|
||||
index_select_co_shard_spmd_rule_test SRCS
|
||||
index_select_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(reshape_co_shard_spmd_rule_test SRCS
|
||||
reshape_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(argsort_co_shard_spmd_rule_test SRCS
|
||||
argsort_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
paddle_test(transpose_co_shard_spmd_rule_test SRCS
|
||||
transpose_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
endif()
|
||||
|
||||
if(WIN32)
|
||||
cc_test(
|
||||
dist_mapper_test
|
||||
SRCS dist_mapper_test.cc
|
||||
DEPS type_info)
|
||||
else()
|
||||
cc_test(
|
||||
dist_mapper_test
|
||||
SRCS dist_mapper_test.cc
|
||||
DEPS phi)
|
||||
endif()
|
||||
@@ -0,0 +1,205 @@
|
||||
/* 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct ArgSortTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_indices_dims_mapping;
|
||||
|
||||
// unused attribute
|
||||
bool descending = true;
|
||||
bool stable = true;
|
||||
};
|
||||
|
||||
struct ArgSortGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> indices_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_indices_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
// unused attribute
|
||||
bool descending = true;
|
||||
bool stable = true;
|
||||
};
|
||||
|
||||
TEST(ArgSortInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<ArgSortTestCase> test_cases = {
|
||||
// shape = [16, 32, 48], axis = -1
|
||||
// [[0,1],[2],[]] -> [[],[2],[]], [[],[2],[]]
|
||||
{{16, 32, 48},
|
||||
{{0, 1}, {2}, {}},
|
||||
-1,
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}}},
|
||||
|
||||
// shape = [16, 32, 48], axis = 2
|
||||
// [[0],[],[1,2]] -> [[0],[],[]], [[0],[],[]]
|
||||
{{16, 32, 48},
|
||||
{{0}, {}, {1, 2}},
|
||||
2,
|
||||
{{0}, {}, {}},
|
||||
{{0}, {}, {}},
|
||||
{{0}, {}, {}}},
|
||||
|
||||
// shape = [10, 32, 48, 24], axis = 1
|
||||
// [[0,1],[2],[],[]] -> [[0,1],[],[],[]], [[0,1],[],[],[]]
|
||||
{{10, 32, 48, 24},
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{0, 1}, {}, {}, {}}}};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::ArgSortInferSpmd(
|
||||
x, tc.axis, tc.descending, tc.stable);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(2));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[1],
|
||||
tc.expected_indices_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(ArgSortGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<ArgSortGradTestCase> test_cases = {
|
||||
// shape = [16, 32, 48], axis = -1
|
||||
// [[0,1],[2],[]], [[0,1],[2],[]], [[0,1],[2],[]] -> [[0,1],[2],[]],
|
||||
// [[0,1],[2],[]], [[0,1],[2],[]], [[0,1],[2],[]]
|
||||
{{16, 32, 48},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
-1,
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}}},
|
||||
// axis = 2
|
||||
// [[0,1],[],[2]], [[0,1],[],[2]], [[0,1],[],[2]] -> [[0,1],[],[]],
|
||||
// [[0,1],[],[]], [[0,1],[],[]], [[0,1],[],[]]
|
||||
{{16, 32, 48},
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
2,
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
// [10, 32, 48, 24], axis = 1
|
||||
// [[0],[1,2],[]], [[0],[1,2],[]], [[0],[1,2],[]] -> [[0],[],[]],
|
||||
// [[0],[],[]], [[0],[],[]], [[0],[],[]]
|
||||
{{10, 32, 48, 24},
|
||||
{{0}, {1, 2}, {}, {}},
|
||||
{{0}, {1, 2}, {}, {}},
|
||||
{{0}, {1, 2}, {}, {}},
|
||||
1,
|
||||
{{0}, {}, {}, {}},
|
||||
{{0}, {}, {}, {}},
|
||||
{{0}, {}, {}, {}},
|
||||
{{0}, {}, {}, {}}}};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr indices_dist_attr = TensorDistAttr();
|
||||
indices_dist_attr.set_process_mesh(process_mesh);
|
||||
indices_dist_attr.set_dims_mapping(tc.indices_dims_mapping);
|
||||
indices_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor indices = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), indices_dist_attr);
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), x_dist_attr);
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.input_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::ArgSortGradInferSpmd(
|
||||
indices, x, out_grad, tc.axis, tc.descending, tc.stable);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_indices_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[2],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,127 @@
|
||||
/* Copyright (c) 2022 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(CrossEntropyInferSpmd, Ctor) {
|
||||
std::vector<int64_t> x_shape = {32, 48};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
TensorDistAttr label_dist_attr = TensorDistAttr();
|
||||
label_dist_attr.set_process_mesh(process_mesh);
|
||||
label_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
label_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
// forward
|
||||
{
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor label(phi::make_ddim(x_shape),
|
||||
label_dist_attr);
|
||||
int axis = 1;
|
||||
|
||||
auto spmdinfo =
|
||||
CrossEntropyWithSoftmaxInferSpmd(x, label, false, true, true, 1, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 2UL);
|
||||
check_dim_mapping(spmdinfo.first[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[1], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[1], {0, -1});
|
||||
check_partial_dims(spmdinfo.second[0], {});
|
||||
|
||||
VLOG(4) << "Test CrossEntropyWithSoftmaxInferSpmd sharding on other axes."
|
||||
<< std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
// test sharding along softmax axis.
|
||||
{
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1}));
|
||||
label_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor label(phi::make_ddim(x_shape),
|
||||
label_dist_attr);
|
||||
int axis = 1;
|
||||
|
||||
auto spmdinfo =
|
||||
CrossEntropyWithSoftmaxInferSpmd(x, label, false, true, true, 1, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 2UL);
|
||||
check_dim_mapping(spmdinfo.first[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[1], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[1], {0, -1});
|
||||
check_partial_dims(spmdinfo.second[0], {});
|
||||
|
||||
VLOG(4) << "Test CrossEntropyWithSoftmaxInferSpmd sharding on other axes."
|
||||
<< std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
// backward
|
||||
{
|
||||
std::vector<int64_t> loss_shape = {32, 1};
|
||||
// Sharding along softmax axis.
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>{0, 1});
|
||||
label_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1}));
|
||||
auto label = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
label_dist_attr);
|
||||
auto softmax =
|
||||
phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
|
||||
auto loss_dist_attr = x_dist_attr;
|
||||
loss_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
auto loss_grad = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(loss_shape), x_dist_attr);
|
||||
|
||||
int axis = 1;
|
||||
auto spmdinfo = CrossEntropyWithSoftmaxGradInferSpmd(
|
||||
label, softmax, loss_grad, true, true, true, 1, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 3UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
check_dim_mapping(spmdinfo.first[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[1], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[2], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[0], {0, -1});
|
||||
check_partial_dims(spmdinfo.second[0], {});
|
||||
|
||||
VLOG(4)
|
||||
<< "Test CrossEntropyWithSoftmaxGradInferSpmd sharding on softmax axis."
|
||||
<< std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,89 @@
|
||||
/* 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/phi/api/ext/op_meta_info.h"
|
||||
#include "paddle/phi/api/ext/spmd_infer.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
TEST(CustomOp, Ctor) {
|
||||
// test with concat rule
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<std::vector<int64_t>> shapes = {
|
||||
{16, 16, 16}, {4, 16, 16}, {2, 16, 16}};
|
||||
std::vector<std::vector<int64_t>> dim_mappings = {
|
||||
{-1, 0, 1}, {-1, 1, 0}, {-1, -1, 0}};
|
||||
std::vector<std::vector<int64_t>> partial_status = {{}, {}, {1}};
|
||||
|
||||
auto build_inputs = [&] {
|
||||
std::vector<phi::distributed::DistMetaTensor> inputs;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
auto t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(dim_mappings[i]);
|
||||
t_dist_attr.set_dynamic_dims({false, false, false});
|
||||
auto input = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(shapes[i]), t_dist_attr);
|
||||
inputs.push_back(input);
|
||||
}
|
||||
return inputs;
|
||||
};
|
||||
|
||||
// test 1, inputs are aligned according to cost, and partial status is cleared
|
||||
auto inputs = build_inputs();
|
||||
|
||||
auto forward_spmd_func =
|
||||
PD_INFER_SPMD_RULE(phi::distributed::ConcatInferSpmd);
|
||||
int axis = 0;
|
||||
std::vector<CustomSpmdInferTensorArg> infer_inputs = {inputs};
|
||||
std::vector<CustomSpmdInferAttrArg> attrs = {axis};
|
||||
|
||||
auto inferred_dist_attrs = forward_spmd_func(infer_inputs, attrs);
|
||||
// list of tensor => single tensor
|
||||
EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
|
||||
inferred_dist_attrs.first[0]));
|
||||
EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
|
||||
inferred_dist_attrs.second[0]));
|
||||
auto& inputs_infer1 =
|
||||
PADDLE_GET_CONST(std::vector<phi::distributed::TensorDistAttr>,
|
||||
inferred_dist_attrs.first[0]);
|
||||
|
||||
for (auto e : inputs_infer1) {
|
||||
check_dim_mapping(e, {-1, 1, 0});
|
||||
check_partial_dims(e, {});
|
||||
}
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {-1, 1, 0});
|
||||
check_partial_dims(inferred_dist_attrs.second[0], {});
|
||||
}
|
||||
|
||||
TEST(CustomOp, Register) {
|
||||
OpMetaInfoBuilder builder("test_custom_op_spmd", 0);
|
||||
auto iter = OpMetaInfoMap::Instance().GetMap().find("test_custom_op_spmd");
|
||||
EXPECT_TRUE(iter != OpMetaInfoMap::Instance().GetMap().end());
|
||||
EXPECT_TRUE(OpMetaInfoHelper::GetInferSpmdFn(iter->second[0]) == nullptr);
|
||||
builder.SetInferSpmdFn(PD_INFER_SPMD_RULE(phi::distributed::ConcatInferSpmd));
|
||||
EXPECT_TRUE(OpMetaInfoHelper::GetInferSpmdFn(iter->second[0]) != nullptr);
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,96 @@
|
||||
/* Copyright (c) 2022 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 <iostream>
|
||||
#include <sstream>
|
||||
|
||||
#include "paddle/phi/core/distributed/auto_parallel/device_mesh.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(DeviceMesh, Ctor) {
|
||||
std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<int64_t> device_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
std::string device_type = "GPU";
|
||||
int64_t size = shape[0] * shape[1];
|
||||
|
||||
DeviceMesh device_mesh("mesh", shape, device_ids, dim_names);
|
||||
for (int64_t i = 0; i < shape[0]; ++i) {
|
||||
for (int64_t j = 0; j < shape[1]; ++j) {
|
||||
int64_t global_id = i * shape[1] + j;
|
||||
int64_t local_id = j;
|
||||
int64_t machine_id = i;
|
||||
device_mesh.add_device(
|
||||
Device(global_id, local_id, machine_id, device_type));
|
||||
}
|
||||
}
|
||||
for (int64_t i = 0; i < size; ++i) {
|
||||
for (int64_t j = 0; j < size; ++j) {
|
||||
device_mesh.add_link(Link(i, j, "NVL"));
|
||||
}
|
||||
}
|
||||
|
||||
EXPECT_EQ(device_mesh.name(), "mesh");
|
||||
EXPECT_EQ(device_mesh.shape(), shape);
|
||||
EXPECT_EQ(device_mesh.device_ids(), device_ids);
|
||||
EXPECT_EQ(device_mesh.dim_names()[0], "x");
|
||||
EXPECT_EQ(device_mesh.dim_names()[1], "y");
|
||||
EXPECT_EQ(device_mesh.device_type(), device_type);
|
||||
EXPECT_EQ(device_mesh.size(), size);
|
||||
EXPECT_EQ(device_mesh.ndim(), static_cast<int64_t>(shape.size()));
|
||||
EXPECT_EQ(device_mesh.dim_size(0), shape[0]);
|
||||
EXPECT_EQ(device_mesh.dim_size(-1), shape[1]);
|
||||
EXPECT_EQ(device_mesh.dim_size("x"), shape[0]);
|
||||
EXPECT_EQ(device_mesh.dim_size("y"), shape[1]);
|
||||
EXPECT_EQ(device_mesh.empty(), false);
|
||||
EXPECT_EQ(device_mesh.contains(0), true);
|
||||
EXPECT_EQ(device_mesh.contains(6), false);
|
||||
EXPECT_EQ(device_mesh.device(3).global_id(), 3);
|
||||
EXPECT_EQ(device_mesh.device(3).local_id(), 0);
|
||||
EXPECT_EQ(device_mesh.device(3).machine_id(), 1);
|
||||
EXPECT_EQ(device_mesh.device(3).type(), "GPU");
|
||||
EXPECT_EQ(device_mesh.link(3, 4).source_id(), 3);
|
||||
EXPECT_EQ(device_mesh.link(3, 4).target_id(), 4);
|
||||
EXPECT_EQ(device_mesh.link(3, 4).type(), "NVL");
|
||||
for (int64_t i = 0; i < shape[0]; ++i) {
|
||||
for (int64_t j = 0; j < shape[1]; ++j) {
|
||||
int64_t global_id = i * shape[1] + j;
|
||||
int64_t local_id = j;
|
||||
int64_t machine_id = i;
|
||||
auto device = device_mesh.devices().at(global_id);
|
||||
EXPECT_EQ(device, Device(global_id, local_id, machine_id, device_type));
|
||||
}
|
||||
}
|
||||
for (int64_t i = 0; i < size; ++i) {
|
||||
for (int64_t j = 0; j < size; ++j) {
|
||||
EXPECT_EQ(device_mesh.links().at(i).at(j), Link(i, j, "NVL"));
|
||||
}
|
||||
}
|
||||
std::stringstream sstream;
|
||||
sstream << device_mesh;
|
||||
EXPECT_EQ(sstream.str(), device_mesh.to_string());
|
||||
auto proto = phi::distributed::to_proto(device_mesh);
|
||||
DeviceMesh new_device_mesh = DeviceMesh::from_proto(proto);
|
||||
EXPECT_EQ(device_mesh, new_device_mesh);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,182 @@
|
||||
/* Copyright (c) 2022 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 <iostream>
|
||||
#include <sstream>
|
||||
#include "glog/logging.h"
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#include "paddle/fluid/distributed/auto_parallel/dist_attr.h"
|
||||
#include "paddle/fluid/framework/block_desc.h"
|
||||
#include "paddle/fluid/framework/op_desc.h"
|
||||
#include "paddle/fluid/framework/program_desc.h"
|
||||
#include "paddle/fluid/framework/var_desc.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using paddle::framework::ProgramDesc;
|
||||
using paddle::framework::VarDesc;
|
||||
|
||||
using paddle::distributed::auto_parallel::get_tensor_shape;
|
||||
using paddle::distributed::auto_parallel::OperatorDistAttr;
|
||||
|
||||
TEST(DistAttr, ctor) {
|
||||
ProgramDesc program;
|
||||
auto* global_block = program.MutableBlock(0);
|
||||
auto* x = global_block->Var("X");
|
||||
x->SetType(paddle::framework::proto::VarType::DENSE_TENSOR);
|
||||
x->SetLoDLevel(0);
|
||||
x->SetDataType(paddle::framework::proto::VarType::FP32);
|
||||
x->SetShape({1000, 784});
|
||||
|
||||
auto* y = global_block->Var("Y");
|
||||
y->SetType(paddle::framework::proto::VarType::DENSE_TENSOR);
|
||||
y->SetLoDLevel(0);
|
||||
y->SetDataType(paddle::framework::proto::VarType::FP32);
|
||||
y->SetShape({784, 100});
|
||||
|
||||
auto* op = global_block->AppendOp();
|
||||
op->SetType("mul");
|
||||
op->SetInput("X", {x->Name()});
|
||||
op->SetInput("Y", {y->Name()});
|
||||
|
||||
auto* out = global_block->Var("Out");
|
||||
out->SetType(paddle::framework::proto::VarType::DENSE_TENSOR);
|
||||
out->SetShape({1000, 100});
|
||||
op->SetOutput("Out", {out->Name()});
|
||||
|
||||
auto get_dist_attr = [](const VarDesc* var_desc) {
|
||||
auto shape = get_tensor_shape(var_desc);
|
||||
return TensorDistAttr(shape);
|
||||
};
|
||||
|
||||
std::vector<int64_t> shape = {2, 4};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> shape2 = {2, 2};
|
||||
std::vector<int64_t> process_ids2 = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names2 = {"a", "b"};
|
||||
ProcessMesh process_mesh2(shape2, process_ids2, dim_names2);
|
||||
|
||||
auto x_dist_attr = get_dist_attr(x);
|
||||
auto y_dist_attr = get_dist_attr(y);
|
||||
auto out_dist_attr = get_dist_attr(out);
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
x_dist_attr.set_batch_dim(0);
|
||||
x_dist_attr.set_chunk_id(0);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({true, false}));
|
||||
x_dist_attr.mark_annotated("process_mesh");
|
||||
x_dist_attr.mark_annotated("dims_mapping");
|
||||
EXPECT_EQ(x_dist_attr.process_mesh(), process_mesh);
|
||||
EXPECT_EQ(x_dist_attr.dims_mapping(), std::vector<int64_t>({0, -1}));
|
||||
EXPECT_EQ(x_dist_attr.batch_dim(), 0);
|
||||
EXPECT_EQ(x_dist_attr.chunk_id(), 0);
|
||||
EXPECT_EQ(x_dist_attr.dynamic_dims(), std::vector<bool>({true, false}));
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("process_mesh"), true);
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("dims_mapping"), true);
|
||||
EXPECT_EQ(x_dist_attr.verify(get_tensor_shape(x)), true);
|
||||
x_dist_attr.clear_annotated();
|
||||
EXPECT_EQ(x_dist_attr.annotated().empty(), true);
|
||||
|
||||
std::stringstream x_sstream;
|
||||
x_sstream << x_dist_attr;
|
||||
EXPECT_EQ(x_sstream.str(), x_dist_attr.to_string());
|
||||
auto x_proto = phi::distributed::to_proto(x_dist_attr);
|
||||
TensorDistAttr new_x_dist_attr = get_dist_attr(x);
|
||||
new_x_dist_attr.from_proto(x_proto);
|
||||
EXPECT_EQ(x_dist_attr, new_x_dist_attr);
|
||||
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 0}));
|
||||
y_dist_attr.set_batch_dim(-1);
|
||||
y_dist_attr.set_chunk_id(0);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>({false, true}));
|
||||
x_dist_attr.mark_annotated("batch_dim");
|
||||
x_dist_attr.mark_annotated("dynamic_dims");
|
||||
EXPECT_EQ(y_dist_attr.process_mesh(), process_mesh);
|
||||
EXPECT_EQ(y_dist_attr.dims_mapping(), std::vector<int64_t>({-1, 0}));
|
||||
EXPECT_EQ(y_dist_attr.batch_dim(), -1);
|
||||
EXPECT_EQ(y_dist_attr.chunk_id(), 0);
|
||||
EXPECT_EQ(y_dist_attr.dynamic_dims(), std::vector<bool>({false, true}));
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("batch_dim"), true);
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("dynamic_dims"), true);
|
||||
EXPECT_EQ(x_dist_attr.verify(get_tensor_shape(y)), true);
|
||||
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1}));
|
||||
out_dist_attr.set_batch_dim(1);
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
EXPECT_EQ(out_dist_attr.process_mesh(), process_mesh);
|
||||
EXPECT_EQ(out_dist_attr.dims_mapping(), std::vector<int64_t>({0, 1}));
|
||||
EXPECT_EQ(out_dist_attr.batch_dim(), 1);
|
||||
EXPECT_EQ(out_dist_attr.dynamic_dims(), std::vector<bool>({false, false}));
|
||||
EXPECT_EQ(out_dist_attr.verify(get_tensor_shape(out)), true);
|
||||
|
||||
OperatorDistAttr mul_dist_attr(*op);
|
||||
EXPECT_EQ(mul_dist_attr.impl_type(),
|
||||
paddle::distributed::auto_parallel::kDefault);
|
||||
EXPECT_EQ(mul_dist_attr.impl_idx(), 0);
|
||||
EXPECT_EQ(mul_dist_attr.chunk_id(), 0);
|
||||
EXPECT_EQ(mul_dist_attr.is_recompute(), false);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("process_mesh"), false);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_type"), false);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_idx"), false);
|
||||
mul_dist_attr.set_input_dist_attr(x->Name(), x_dist_attr);
|
||||
mul_dist_attr.set_input_dist_attr(y->Name(), y_dist_attr);
|
||||
mul_dist_attr.set_output_dist_attr(out->Name(), out_dist_attr);
|
||||
mul_dist_attr.set_process_mesh(process_mesh2);
|
||||
mul_dist_attr.set_impl_type("dist_mul");
|
||||
mul_dist_attr.set_impl_idx(0);
|
||||
mul_dist_attr.set_chunk_id(1);
|
||||
mul_dist_attr.set_is_recompute(true);
|
||||
mul_dist_attr.mark_annotated("process_mesh");
|
||||
mul_dist_attr.mark_annotated("impl_type");
|
||||
mul_dist_attr.mark_annotated("impl_idx");
|
||||
EXPECT_NE(mul_dist_attr.input_dist_attr(x->Name()), x_dist_attr);
|
||||
EXPECT_NE(mul_dist_attr.input_dist_attr(y->Name()), y_dist_attr);
|
||||
EXPECT_NE(mul_dist_attr.output_dist_attr(out->Name()), out_dist_attr);
|
||||
EXPECT_EQ(mul_dist_attr.process_mesh(), process_mesh2);
|
||||
EXPECT_EQ(mul_dist_attr.input_dist_attr(x->Name()).process_mesh(),
|
||||
process_mesh2);
|
||||
EXPECT_EQ(mul_dist_attr.input_dist_attr(y->Name()).process_mesh(),
|
||||
process_mesh2);
|
||||
EXPECT_EQ(mul_dist_attr.impl_type(), "dist_mul");
|
||||
EXPECT_EQ(mul_dist_attr.impl_idx(), 0);
|
||||
EXPECT_EQ(mul_dist_attr.chunk_id(), 1);
|
||||
EXPECT_EQ(mul_dist_attr.is_recompute(), true);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("process_mesh"), true);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_type"), true);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_idx"), true);
|
||||
EXPECT_EQ(mul_dist_attr.verify(op), true);
|
||||
mul_dist_attr.clear_annotated();
|
||||
EXPECT_EQ(mul_dist_attr.annotated().empty(), true);
|
||||
|
||||
std::stringstream mul_sstream;
|
||||
mul_sstream << mul_dist_attr;
|
||||
EXPECT_EQ(mul_sstream.str(), mul_dist_attr.to_string());
|
||||
auto mul_proto = mul_dist_attr.to_proto();
|
||||
OperatorDistAttr new_mul_dist_attr(*op);
|
||||
new_mul_dist_attr.from_proto(mul_proto);
|
||||
EXPECT_EQ(mul_dist_attr, new_mul_dist_attr);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,73 @@
|
||||
/* Copyright (c) 2022 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/core/distributed/auto_parallel/dist_mapper.h"
|
||||
#include <map>
|
||||
#include <sstream>
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(DistributedMapper, Ctor) {
|
||||
std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<int64_t> device_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
std::string device_type = "GPU";
|
||||
int64_t size = shape[0] * shape[1];
|
||||
|
||||
DeviceMesh device_mesh("device_mesh", shape, device_ids, dim_names);
|
||||
for (int64_t i = 0; i < shape[0]; ++i) {
|
||||
for (int64_t j = 0; j < shape[1]; ++j) {
|
||||
int64_t global_id = i * shape[1] + j;
|
||||
int64_t local_id = j;
|
||||
int64_t machine_id = i;
|
||||
device_mesh.add_device(
|
||||
Device(global_id, local_id, machine_id, device_type));
|
||||
}
|
||||
}
|
||||
for (int64_t i = 0; i < size; ++i) {
|
||||
for (int64_t j = 0; j < size; ++j) {
|
||||
device_mesh.add_link(Link(i, j, "NVL"));
|
||||
}
|
||||
}
|
||||
|
||||
DistributedMapper dist_mapper;
|
||||
dist_mapper.add_device_mesh(device_mesh);
|
||||
std::map<int64_t, std::pair<std::string, std::vector<int64_t>>>
|
||||
process_id_to_device_ids;
|
||||
process_id_to_device_ids[0] = {"device_mesh", {5}};
|
||||
process_id_to_device_ids[1] = {"device_mesh", {4}};
|
||||
process_id_to_device_ids[2] = {"device_mesh", {3}};
|
||||
process_id_to_device_ids[3] = {"device_mesh", {2}};
|
||||
process_id_to_device_ids[4] = {"device_mesh", {1}};
|
||||
process_id_to_device_ids[5] = {"device_mesh", {0}};
|
||||
dist_mapper.set_process_id_to_device_ids(process_id_to_device_ids);
|
||||
|
||||
EXPECT_EQ(dist_mapper.device_meshes().at("device_mesh"), device_mesh);
|
||||
EXPECT_EQ(dist_mapper.device_mesh("device_mesh"), device_mesh);
|
||||
EXPECT_EQ(dist_mapper.process_id_to_device_ids(), process_id_to_device_ids);
|
||||
std::stringstream sstream;
|
||||
sstream << dist_mapper;
|
||||
EXPECT_EQ(sstream.str(), dist_mapper.to_string());
|
||||
auto proto = phi::distributed::to_proto(dist_mapper);
|
||||
DistributedMapper new_dist_mapper = DistributedMapper::from_proto(proto);
|
||||
EXPECT_EQ(dist_mapper, new_dist_mapper);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,72 @@
|
||||
/* 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/phi/core/distributed/auto_parallel/dist_tensor.h"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
|
||||
#include "test/cpp/phi/core/allocator.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace tests {
|
||||
|
||||
TEST(dist_tensor, constructor) {
|
||||
auto fancy_allocator =
|
||||
std::unique_ptr<Allocator>(new phi::tests::FancyAllocator);
|
||||
auto* alloc = fancy_allocator.get();
|
||||
|
||||
DataType dtype{DataType::FLOAT16};
|
||||
DDim dims({3, 4});
|
||||
DenseTensorMeta meta(dtype, dims);
|
||||
|
||||
auto dist_attr = TensorDistAttr(common::vectorize(dims));
|
||||
|
||||
std::vector<int64_t> mesh_shape = {1};
|
||||
std::vector<int64_t> process_ids = {0};
|
||||
std::vector<std::string> dim_names = {"x"};
|
||||
ProcessMesh mesh(mesh_shape, process_ids, dim_names);
|
||||
dist_attr.set_process_mesh(mesh);
|
||||
|
||||
// copy construct
|
||||
std::shared_ptr<DenseTensor> x1 = std::make_shared<DenseTensor>(alloc, meta);
|
||||
DistTensor dist_x1(x1, dist_attr);
|
||||
EXPECT_TRUE(dist_x1.defined());
|
||||
EXPECT_TRUE(dist_x1.initialized());
|
||||
EXPECT_TRUE(dist_x1.valid());
|
||||
EXPECT_EQ(dist_x1.numel(), 12L);
|
||||
EXPECT_EQ(dist_x1.local_dims()[0], 3L);
|
||||
EXPECT_EQ(dist_x1.local_dims()[1], 4L);
|
||||
|
||||
// empty construct
|
||||
DistTensor dist_x2(dims, dist_attr);
|
||||
EXPECT_TRUE(!dist_x2.defined());
|
||||
EXPECT_TRUE(!dist_x2.initialized());
|
||||
// allocate error test
|
||||
bool caught_exception = false;
|
||||
try {
|
||||
dist_x2.AllocateFrom(alloc, phi::DataType::FLOAT32, 12L, false);
|
||||
} catch (common::enforce::EnforceNotMet& error) {
|
||||
caught_exception = true;
|
||||
EXPECT_NE(std::string(error.what()).find("Unavailable"), 0UL);
|
||||
}
|
||||
EXPECT_TRUE(caught_exception);
|
||||
}
|
||||
|
||||
} // namespace tests
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,116 @@
|
||||
/* Copyright (c) 2022 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(ExpandAsInferSpmd, Ctor) {
|
||||
// Sharding along axes besides softmax axis.
|
||||
std::vector<int64_t> x_shape = {1, 48};
|
||||
std::vector<int64_t> y_shape = {2, 32, 48};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
TensorDistAttr y_dist_attr = TensorDistAttr();
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor y(phi::make_ddim(y_shape), y_dist_attr);
|
||||
|
||||
// test info forward
|
||||
auto spmdinfo = ExpandAsInferSpmd(x, y, y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test ExpandAsInferSpmd" << std::endl << std::endl << std::endl;
|
||||
|
||||
// test info forward without y
|
||||
// x [1, 48], target [2, 32, 48]: [-1, -1] -> [-1, -1, -1]
|
||||
spmdinfo = ExpandAsInferSpmd(x, phi::distributed::DistMetaTensor(), y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
const phi::distributed::ArgDistAttr& attr = spmdinfo.first[1];
|
||||
if (paddle::holds_alternative<phi::distributed::TensorDistAttr>(attr)) {
|
||||
EXPECT_EQ(paddle::get<phi::distributed::TensorDistAttr>(attr),
|
||||
phi::distributed::TensorDistAttr());
|
||||
} else {
|
||||
FAIL() << "forward_info.first[1] is not TensorDistAttr";
|
||||
}
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({-1, -1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test ExpandAsInferSpmd" << std::endl << std::endl << std::endl;
|
||||
|
||||
// test info reverse
|
||||
spmdinfo = ExpandAsInferSpmdReverse(x, y, y, y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test ExpandAsInferSpmdReverse" << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// test info grad
|
||||
spmdinfo = ExpandAsGradInferSpmd(x, y, y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
check_partial_dims(spmdinfo.second[0], {0, 1});
|
||||
VLOG(4) << "Test ExpandAsGradInferSpmd" << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,105 @@
|
||||
/* 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 "glog/logging.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
ProcessMesh CreateProcessMesh() {
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
return ProcessMesh(mesh_shape, process_ids, dim_names);
|
||||
}
|
||||
|
||||
phi::distributed::DistMetaTensor CreateDistMetaTensor(
|
||||
const std::vector<int64_t>& shape,
|
||||
const std::vector<int64_t>& dims_mapping,
|
||||
const ProcessMesh& process_mesh) {
|
||||
TensorDistAttr dist_attr;
|
||||
dist_attr.set_process_mesh(process_mesh);
|
||||
dist_attr.set_dims_mapping(dims_mapping);
|
||||
return phi::distributed::DistMetaTensor(phi::make_ddim(shape), dist_attr);
|
||||
}
|
||||
|
||||
TEST(ExpandInferSpmd, Ctor) {
|
||||
ProcessMesh process_mesh = CreateProcessMesh();
|
||||
|
||||
// Test case forward 1: Expand with shape {8, 2, 6, 1024, -1}
|
||||
auto x = CreateDistMetaTensor(
|
||||
{8, 2, 1, 1024, 128}, {0, -1, -1, 1, -1}, process_mesh);
|
||||
phi::IntArray shape = {8, 2, 6, 1024, -1};
|
||||
auto spmdinfo = ExpandInferSpmd(x, shape);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({0, -1, -1, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, -1, -1, 1, -1}));
|
||||
|
||||
// Test case forward 2: Expand with shape {2, -1}
|
||||
auto x1 = CreateDistMetaTensor({8}, {1}, process_mesh);
|
||||
phi::IntArray shape1 = {2, -1};
|
||||
auto spmdinfo1 = ExpandInferSpmd(x1, shape1);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo1.first[0]), std::vector<int64_t>({1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo1.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
|
||||
// Test case forward 3: Expand with shape {0, -1}
|
||||
auto x2 = CreateDistMetaTensor({8}, {1}, process_mesh);
|
||||
phi::IntArray shape2 = {0, -1};
|
||||
auto spmdinfo2 = ExpandInferSpmd(x2, shape2);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo2.first[0]), std::vector<int64_t>({1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo2.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
|
||||
// Test case backward 1: ExpandGrad with shape {0, -1}
|
||||
auto x3 = CreateDistMetaTensor({8}, {1}, process_mesh);
|
||||
auto out3 = CreateDistMetaTensor({2, 8}, {-1, 1}, process_mesh);
|
||||
phi::IntArray shape3 = {0, -1};
|
||||
auto spmdinfo3 = ExpandGradInferSpmd(x3, out3, shape3);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo3.first[0]), std::vector<int64_t>({1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo3.first[1]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo3.second[0]), std::vector<int64_t>({1}));
|
||||
|
||||
// Test case backward 2: ExpandGrad with shape {2, 2, -1}
|
||||
auto x4 = CreateDistMetaTensor({1, 8}, {-1, 1}, process_mesh);
|
||||
auto out4 = CreateDistMetaTensor({2, 2, 8}, {-1, -1, 1}, process_mesh);
|
||||
phi::IntArray shape4 = {2, 2, -1};
|
||||
auto spmdinfo4 = ExpandGradInferSpmd(x4, out4, shape4);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo4.first[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo4.first[1]),
|
||||
std::vector<int64_t>({-1, -1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo4.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
|
||||
// Test case backward 3: ExpandGrad with shape {2, 2, -1}
|
||||
auto x5 = CreateDistMetaTensor({1, 8}, {-1, 1}, process_mesh);
|
||||
auto out5 = CreateDistMetaTensor({2, 2, 8}, {-1, 0, 1}, process_mesh);
|
||||
phi::IntArray shape5 = {2, 2, -1};
|
||||
auto spmdinfo5 = ExpandGradInferSpmd(x5, out5, shape5);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo5.first[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo5.first[1]),
|
||||
std::vector<int64_t>({-1, 0, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo5.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_partial_dims(spmdinfo5.second[0]), std::set<int64_t>({0}));
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,103 @@
|
||||
/* Copyright (c) 2022 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(FusedLinearParamGradAddSPMDRule, Ctor) {
|
||||
// build input data class
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 4};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
// b s h
|
||||
std::vector<int64_t> x_shape = {2, 512, 1024};
|
||||
std::vector<int64_t> out_shape = {2, 512, 2048};
|
||||
std::vector<int64_t> weight_shape = {1024, 2048};
|
||||
std::vector<int64_t> bias_shape = {2048};
|
||||
|
||||
// test mp col split
|
||||
{
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor out(phi::make_ddim(out_shape),
|
||||
out_dist_attr);
|
||||
phi::distributed::DistMetaTensor dweight;
|
||||
phi::distributed::DistMetaTensor dbias;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
auto spmd_info =
|
||||
FusedLinearParamGradAddInferSpmd(x, out, dweight, dbias, 0, true);
|
||||
check_dim_mapping(spmd_info.second[0], {-1, 1});
|
||||
check_partial_dims(spmd_info.second[0], {0});
|
||||
check_dim_mapping(spmd_info.second[1], {1});
|
||||
check_partial_dims(spmd_info.second[1], {0});
|
||||
dweight = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(weight_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[0]));
|
||||
dbias = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(bias_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[1]));
|
||||
}
|
||||
}
|
||||
|
||||
// test mp row split
|
||||
{
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1}));
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor out(phi::make_ddim(out_shape),
|
||||
out_dist_attr);
|
||||
phi::distributed::DistMetaTensor dweight;
|
||||
phi::distributed::DistMetaTensor dbias;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
auto spmd_info =
|
||||
FusedLinearParamGradAddInferSpmd(x, out, dweight, dbias, 0, true);
|
||||
check_dim_mapping(spmd_info.second[0], {1, -1});
|
||||
check_partial_dims(spmd_info.second[0], {0});
|
||||
check_dim_mapping(spmd_info.second[1], {-1});
|
||||
check_partial_dims(spmd_info.second[1], {0});
|
||||
dweight = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(weight_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[0]));
|
||||
dbias = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(bias_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[1]));
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,110 @@
|
||||
/* Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
TEST(FusedRmsNormSPMDRule, test_fused_rms_norm) {
|
||||
// build input data class
|
||||
std::vector<int64_t> x_shape = {64, 32, 1024};
|
||||
std::vector<int64_t> scale_shape = {1024};
|
||||
std::vector<int64_t> variance_shape = {64, 32};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
TensorDistAttr scale_dist_attr = TensorDistAttr();
|
||||
scale_dist_attr.set_process_mesh(process_mesh);
|
||||
scale_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
|
||||
scale_dist_attr.set_dynamic_dims(std::vector<bool>({false}));
|
||||
|
||||
x_dist_attr.set_dims_mapping({1, -1, -1});
|
||||
scale_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
|
||||
|
||||
phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor scale(common::make_ddim(scale_shape),
|
||||
scale_dist_attr);
|
||||
auto inferred_dist_attrs = phi::distributed::RmsNormInferSpmd(x, scale, 0.5);
|
||||
|
||||
size_t input_size = 2;
|
||||
size_t output_size = 2;
|
||||
EXPECT_EQ(inferred_dist_attrs.first.size(), input_size);
|
||||
EXPECT_EQ(inferred_dist_attrs.second.size(), output_size);
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {1, -1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {1, -1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {1, -1});
|
||||
|
||||
VLOG(4) << "test1 done.";
|
||||
|
||||
x_dist_attr.set_dims_mapping({1, 0, -1});
|
||||
scale_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
|
||||
x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
|
||||
scale = phi::distributed::DistMetaTensor(common::make_ddim(scale_shape),
|
||||
scale_dist_attr);
|
||||
|
||||
inferred_dist_attrs = phi::distributed::RmsNormInferSpmd(x, scale, 0.5);
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {1, 0, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {1, 0, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {1, 0});
|
||||
VLOG(4) << "test2 done.";
|
||||
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
phi::distributed::DistMetaTensor out(common::make_ddim(x_shape),
|
||||
out_dist_attr);
|
||||
|
||||
TensorDistAttr invvar_dist_attr = TensorDistAttr();
|
||||
invvar_dist_attr.set_process_mesh(process_mesh);
|
||||
invvar_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
invvar_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
phi::distributed::DistMetaTensor invvar(common::make_ddim(variance_shape),
|
||||
invvar_dist_attr);
|
||||
|
||||
inferred_dist_attrs =
|
||||
phi::distributed::RmsNormInferSpmdReverse(x, scale, out, invvar, 0.5);
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {0, 1});
|
||||
VLOG(4) << "test3 done.";
|
||||
|
||||
x_dist_attr.set_dims_mapping({0, 1, -1});
|
||||
x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
|
||||
inferred_dist_attrs =
|
||||
phi::distributed::RmsNormGradInferSpmd(x, scale, invvar, out, 0.5);
|
||||
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[2], {0, 1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[3], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {-1});
|
||||
check_partial_dims(inferred_dist_attrs.second[1], {0, 1});
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,286 @@
|
||||
/* 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 <set>
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct IndexSelectTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
std::vector<int64_t> index_shape;
|
||||
std::vector<std::vector<int64_t>> index_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_index_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
};
|
||||
|
||||
struct IndexSelectGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
std::vector<int64_t> index_shape;
|
||||
std::vector<std::vector<int64_t>> index_dims_mapping;
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_index_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
std::set<int64_t> partial_dims;
|
||||
};
|
||||
|
||||
TEST(IndexSelectInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<IndexSelectTestCase> test_cases = {
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0,1],[2],[]], [[]] -> [[0,1],[],[]], [[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0,1],[2],[]], [[2]] -> [[0,1],[],[]], [[2]], [[0,1],[2],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{2}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{2}},
|
||||
{{0, 1}, {2}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0,1],[2],[]], [[0]] -> [[0,1],[],[]], [[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{0}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[2],[],[]], [[0,1]] -> [[2],[],[]], [[0,1]], [[2],[0,1],[]]
|
||||
{{8, 16, 32},
|
||||
{{2}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
1,
|
||||
{{2}, {}, {}},
|
||||
{{0, 1}},
|
||||
{{2}, {0, 1}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0],[],[]], [[0,1]] -> [[0],[],[]], [[1]], [[0],[1],[]]
|
||||
{{8, 16, 32},
|
||||
{{0}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
1,
|
||||
{{0}, {}, {}},
|
||||
{{1}},
|
||||
{{0}, {1}, {}}},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr index_dist_attr = TensorDistAttr();
|
||||
index_dist_attr.set_process_mesh(process_mesh);
|
||||
index_dist_attr.set_dims_mapping(tc.index_dims_mapping);
|
||||
index_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.index_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor index = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.index_shape), index_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::IndexSelectInferSpmd(x, index, tc.axis);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.first[1],
|
||||
tc.expected_index_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(IndexSelectGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<IndexSelectGradTestCase> test_cases = {
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0,1],[2],[]], [[]], [[0,1], [], []] -> [[0,1],[],[]], [[]],
|
||||
// [[0,1],[],[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{}},
|
||||
{8, 8, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0,1],[2],[]], [[2]], [[0,1],[2],[]] -> [[0,1],[],[]], [[2]],
|
||||
// [[0,1],[2],[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{2}},
|
||||
{8, 8, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{2}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{2}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0,1],[2],[]], [[0]], [[0,1],[],[]] -> [[0,1],[],[]], [[]],
|
||||
// [[0,1],[],[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{0}},
|
||||
{8, 8, 32},
|
||||
{{0, 1}, {}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[2],[],[]], [[0,1]], [[2],[0,1],[]] -> [[2],[],[]], [[0,1]],
|
||||
// [[2],[0,1],[]], [[2],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{2}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
{8, 8, 32},
|
||||
{{2}, {0, 1}, {}},
|
||||
1,
|
||||
{{2}, {}, {}},
|
||||
{{0, 1}},
|
||||
{{2}, {0, 1}, {}},
|
||||
{{2}, {}, {}},
|
||||
{0, 1}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0],[],[]], [[0,1]], [[0],[1],[]] -> [[0],[],[]], [[1]], [[0],[1],[]],
|
||||
// [[0],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
{8, 8, 32},
|
||||
{{0}, {1}, {}},
|
||||
1,
|
||||
{{0}, {}, {}},
|
||||
{{1}},
|
||||
{{0}, {1}, {}},
|
||||
{{0}, {}, {}},
|
||||
{1}},
|
||||
};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr index_dist_attr = TensorDistAttr();
|
||||
index_dist_attr.set_process_mesh(process_mesh);
|
||||
index_dist_attr.set_dims_mapping(tc.index_dims_mapping);
|
||||
index_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.index_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor index = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.index_shape), index_dist_attr);
|
||||
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::IndexSelectGradInferSpmd(x, index, out_grad, tc.axis);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_index_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[2],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
if (!tc.partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(backward_spmd_info.second[0]), true);
|
||||
check_partial_dims(backward_spmd_info.second[0], tc.partial_dims);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,505 @@
|
||||
/* 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 <set>
|
||||
#include "paddle/phi/infermeta/spmd_rules/bmm.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct MatmulTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<int64_t> y_shape;
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping;
|
||||
|
||||
// attribute
|
||||
bool trans_x;
|
||||
bool trans_y;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_y_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
|
||||
std::set<int64_t> partial_dims;
|
||||
};
|
||||
|
||||
struct MatmulGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<int64_t> y_shape;
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping;
|
||||
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// attribute
|
||||
bool trans_x;
|
||||
bool trans_y;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_y_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_y_grad_dims_mapping;
|
||||
|
||||
std::set<int64_t> x_grad_partial_dims;
|
||||
std::set<int64_t> y_grad_partial_dims;
|
||||
};
|
||||
|
||||
TEST(MatmulInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<MatmulTestCase> test_cases = {
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1], []] ,[[],[2]] -> [[0,1], []] ,[[],[2]],[[0,1],[2]]
|
||||
{{64, 32},
|
||||
{{0, 1}, {}},
|
||||
{32, 48},
|
||||
{{}, {2}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {}},
|
||||
{{}, {2}},
|
||||
{{0, 1}, {2}},
|
||||
{}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1], [2]] ,[[],[]] -> [[0,1], [2]] ,[[2],[]],[[0,1],[]], partial: 2
|
||||
{{64, 32},
|
||||
{{0, 1}, {2}},
|
||||
{32, 48},
|
||||
{{}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {2}},
|
||||
{{2}, {}},
|
||||
{{0, 1}, {}},
|
||||
{2}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[], []] ,[[0,1],[2]] -> [[],[0,1]] ,[[0,1],[2],[[],[2]], partial:
|
||||
// {0,1}
|
||||
{{64, 32},
|
||||
{{}, {}},
|
||||
{32, 48},
|
||||
{{0, 1}, {2}},
|
||||
false,
|
||||
false,
|
||||
{{}, {0, 1}},
|
||||
{{0, 1}, {2}},
|
||||
{{}, {2}},
|
||||
{0, 1}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0], [1]] ,[[2],[0]] -> [[0], [1,2]] ,[[1,2],[]],[[0],[]], partial:
|
||||
// {1,2}
|
||||
{{64, 32},
|
||||
{{0}, {1}},
|
||||
{32, 48},
|
||||
{{2}, {0}},
|
||||
false,
|
||||
false,
|
||||
{{0}, {1, 2}},
|
||||
{{1, 2}, {}},
|
||||
{{0}, {}},
|
||||
{1, 2}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1], [2]] ,[[0],[]] -> [[0,1], [2]] ,[[2],[]],[[0,1],[]], partial: 2
|
||||
{{64, 32},
|
||||
{{0, 1}, {2}},
|
||||
{32, 48},
|
||||
{{0}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {2}},
|
||||
{{2}, {}},
|
||||
{{0, 1}, {}},
|
||||
{2}},
|
||||
|
||||
// [512, 48, 64, 32], [1, 32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1],[2],[],[]] ,[[],[],[]] -> [[0,1],[2],[],[]]
|
||||
// ,[[],[],[]],[[0,1],[2],[],[]],
|
||||
// partial: {}
|
||||
{{512, 48, 64, 32},
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
{{}, {}, {}},
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
{}},
|
||||
|
||||
// [512, 48, 32, 64], [1, 32, 48], trans_x=true, trans_y=false
|
||||
// [[0],[],[1,2],[]] ,[[],[],[2]] -> [[0],[],[1],[]]
|
||||
// ,[[],[1],[2]],[[0],[],[],[2]],
|
||||
// partial: {1}
|
||||
{{512, 48, 32, 64},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
true,
|
||||
false,
|
||||
{{0}, {}, {1}, {}},
|
||||
{{}, {1}, {2}},
|
||||
{{0}, {}, {}, {2}},
|
||||
{1}},
|
||||
|
||||
// [512, 48, 64, 32], [1, 48, 32], trans_x=false, trans_y=true
|
||||
// [[0],[],[1,2],[]] ,[[],[0],[]] -> [[0],[],[1,2],[]]
|
||||
// ,[[],[],[]],[[0],[],[1,2],[]],
|
||||
// partial: {}
|
||||
{{512, 48, 64, 32},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{1, 48, 32},
|
||||
{{}, {0}, {}},
|
||||
false,
|
||||
true,
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{}},
|
||||
|
||||
// [512, 48, 32, 64], [1, 48, 32], trans_x=true, trans_y=true
|
||||
// [[],[],[0,1],[2]] ,[[],[0,1],[2]] -> [[],[],[],[2]]
|
||||
// ,[[],[0,1],[]],[[],[],[2],[0,1]],
|
||||
// partial: {}
|
||||
{{512, 48, 32, 64},
|
||||
{{}, {}, {0, 1}, {2}},
|
||||
{1, 48, 32},
|
||||
{{}, {0, 1}, {2}},
|
||||
true,
|
||||
true,
|
||||
{{}, {}, {}, {2}},
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {}, {2}, {0, 1}},
|
||||
{}},
|
||||
};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr y_dist_attr = TensorDistAttr();
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(tc.y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(tc.y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.y_shape), y_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::MatmulInferSpmd(x, y, tc.trans_x, tc.trans_y);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.first[1],
|
||||
tc.expected_y_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
if (!tc.partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(forward_spmd_info.second[0]), true);
|
||||
check_partial_dims(forward_spmd_info.second[0], tc.partial_dims);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(MatmulGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<MatmulGradTestCase> test_cases = {
|
||||
// [64, 32], [32, 48], [64,48], trans_x=false, trans_y=false
|
||||
// [[0,1], []] ,[[],[2]], [[0,1],[2]] -> [[0,1], []]
|
||||
// ,[[],[2]],[[0,1],[2]], [[0,1],[]], [[],[2]], x_partial: {2}, y_partial:
|
||||
// {0,1}
|
||||
{{64, 32},
|
||||
{{0, 1}, {}},
|
||||
{32, 48},
|
||||
{{}, {2}},
|
||||
{64, 48},
|
||||
{{0, 1}, {2}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {}},
|
||||
{{}, {2}},
|
||||
{{0, 1}, {2}},
|
||||
{{0, 1}, {}},
|
||||
{{}, {2}},
|
||||
{2},
|
||||
{0, 1}},
|
||||
// [1024,512,64,32], [1,32,48], [1024,512,64,48], trans_x=false,
|
||||
// trans_y=false
|
||||
// [[0],[],[1,2],[]] ,[[],[],[2]], [[0],[],[1,2],[]] -> [[0],[],[1,2],[]]
|
||||
// ,[[],[],[]], [[0],[],[1,2],[]], [[0],[],[1,2],[]], [[],[],[]],
|
||||
// x_grad_partial: {}, y_grad_partial: {0,1,2}
|
||||
{{1024, 512, 64, 32},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
{1024, 512, 64, 48},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{},
|
||||
{0, 1, 2}},
|
||||
// [1024,512,64,32], [1,32,48], [1024,512,64,48], trans_x=false,
|
||||
// trans_y=false
|
||||
// [[],[0],[1,2],[]] ,[[],[],[2]], [[],[0],[1,2],[]] -> [[],[0],[1,2],[]]
|
||||
// ,[[],[],[]], [[],[0],[1,2],[]], [[],[0],[1,2],[]], [[],[],[]],
|
||||
// x_grad_partial: {}, y_grad_partial: {0,1,2}
|
||||
{{1024, 512, 64, 32},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
false,
|
||||
false,
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{},
|
||||
{0, 1, 2}},
|
||||
// [1024,512,32,64], [1,32,48], [1024,512,64,48], trans_x=true,
|
||||
// trans_y=false
|
||||
// [[],[0],[1,2],[]] ,[[],[],[2]], [[],[0],[],[2]] -> [[],[0],[1],[]]
|
||||
// ,[[],[1],[2]], [[],[0],[],[2]], [[],[0],[1],[]], [[],[1],[2]],
|
||||
// x_grad_partial: {2}, y_grad_partial: {0}
|
||||
{{1024, 512, 32, 64},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {0}, {}, {2}},
|
||||
true,
|
||||
false,
|
||||
{{}, {0}, {1}, {}},
|
||||
{{}, {1}, {2}},
|
||||
{{}, {0}, {}, {2}},
|
||||
{{}, {0}, {1}, {}},
|
||||
{{}, {1}, {2}},
|
||||
{2},
|
||||
{0}},
|
||||
// [1024,512,32,64], [1,48,32], [1024,512,64,48], trans_x=true,
|
||||
// trans_y=true
|
||||
// [[],[],[1,2],[]] ,[[],[],[0]], [[],[],[],[]] -> [[],[],[0,1,2],[]]
|
||||
// ,[[],[],[0,1,2]], [[],[],[],[]], [[],[],[0,1,2],[]], [[],[],[0,1,2]],
|
||||
// x_grad_partial: {}, y_grad_partial: {}
|
||||
{{1024, 512, 32, 64},
|
||||
{{}, {}, {1, 2}, {}},
|
||||
{1, 48, 32},
|
||||
{{}, {}, {0}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {}, {}, {}},
|
||||
true,
|
||||
true,
|
||||
{{}, {}, {1, 2, 0}, {}},
|
||||
{{}, {}, {1, 2, 0}},
|
||||
{{}, {}, {}, {}},
|
||||
{{}, {}, {1, 2, 0}, {}},
|
||||
{{}, {}, {1, 2, 0}},
|
||||
{},
|
||||
{}},
|
||||
// [1024,512,64,32], [1,48,32], [1024,512,64,48], trans_x=false,
|
||||
// trans_y=true
|
||||
// [[],[],[0],[1,2]] ,[[],[],[0]], [[],[],[0],[]] -> [[],[],[0],[1,2]]
|
||||
// ,[[],[],[1,2]], [[],[],[0],[]], [[],[],[0],[1,2]],
|
||||
// [[],[],[1,2]],
|
||||
// x_grad_partial: {}, y_grad_partial: {0}
|
||||
{{1024, 512, 64, 32},
|
||||
{{}, {}, {0}, {1, 2}},
|
||||
{1, 48, 32},
|
||||
{{}, {}, {0}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {}, {0}, {}},
|
||||
false,
|
||||
true,
|
||||
{{}, {}, {0}, {1, 2}},
|
||||
{{}, {}, {1, 2}},
|
||||
{{}, {}, {0}, {}},
|
||||
{{}, {}, {0}, {1, 2}},
|
||||
{{}, {}, {1, 2}},
|
||||
{},
|
||||
{0}}};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr y_dist_attr = TensorDistAttr();
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(tc.y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(tc.y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.y_shape), y_dist_attr);
|
||||
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::MatmulGradInferSpmd(
|
||||
x, y, out_grad, tc.trans_x, tc.trans_y);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(2));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_y_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[2],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
if (!tc.x_grad_partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(backward_spmd_info.second[0]), true);
|
||||
check_partial_dims(backward_spmd_info.second[0], tc.x_grad_partial_dims);
|
||||
}
|
||||
check_multi_dims_mapping(backward_spmd_info.second[1],
|
||||
tc.expected_y_grad_dims_mapping);
|
||||
if (!tc.y_grad_partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(backward_spmd_info.second[1]), true);
|
||||
check_partial_dims(backward_spmd_info.second[1], tc.y_grad_partial_dims);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BmmInferSpmd, CoShard) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> x_shape = {4, 16, 8};
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping = {{0, 1}, {2}, {}};
|
||||
TensorDistAttr x_dist_attr;
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
|
||||
|
||||
std::vector<int64_t> y_shape = {4, 8, 32};
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping = {{0, 1}, {}, {}};
|
||||
TensorDistAttr y_dist_attr;
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);
|
||||
|
||||
auto bmm_spmd_info = phi::distributed::BmmInferSpmd(x, y);
|
||||
|
||||
ASSERT_EQ(bmm_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
ASSERT_EQ(bmm_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
|
||||
check_multi_dims_mapping(bmm_spmd_info.first[0], x_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_spmd_info.first[0]));
|
||||
check_multi_dims_mapping(bmm_spmd_info.first[1], y_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_spmd_info.first[1]));
|
||||
|
||||
const std::vector<std::vector<int64_t>> expected_out_dims_mapping = {
|
||||
{0, 1}, {2}, {}};
|
||||
check_multi_dims_mapping(bmm_spmd_info.second[0], expected_out_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_spmd_info.second[0]));
|
||||
}
|
||||
|
||||
TEST(BmmGradInferSpmd, CoShard) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> x_shape = {4, 16, 8};
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping = {{0, 1}, {2}, {}};
|
||||
TensorDistAttr x_dist_attr;
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
|
||||
|
||||
std::vector<int64_t> y_shape = {4, 8, 32};
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping = {{0, 1}, {}, {}};
|
||||
TensorDistAttr y_dist_attr;
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);
|
||||
|
||||
std::vector<int64_t> out_grad_shape = {4, 16, 32};
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping = {{0, 1}, {2}, {}};
|
||||
TensorDistAttr out_grad_dist_attr;
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad(common::make_ddim(out_grad_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
auto bmm_grad_spmd_info = phi::distributed::BmmGradInferSpmd(x, y, out_grad);
|
||||
|
||||
ASSERT_EQ(bmm_grad_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
ASSERT_EQ(bmm_grad_spmd_info.second.size(), static_cast<size_t>(2));
|
||||
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.first[0], x_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.first[0]));
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.first[1], y_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.first[1]));
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.first[2], out_grad_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.first[2]));
|
||||
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.second[0], x_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.second[0]));
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.second[1], y_dims_mapping);
|
||||
EXPECT_TRUE(is_partial(bmm_grad_spmd_info.second[1]));
|
||||
check_partial_dims(bmm_grad_spmd_info.second[1], {2});
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,144 @@
|
||||
/* Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using phi::distributed::ArgDistAttr;
|
||||
using phi::distributed::DistMetaTensor;
|
||||
|
||||
void test_moe_combine_spmd(
|
||||
const std::vector<std::vector<int64_t>>& input_shapes,
|
||||
const std::vector<std::vector<int64_t>>& input_dims_mappings,
|
||||
const std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>& expected_dims_mappings,
|
||||
bool test_bwd_spmd = false) {
|
||||
size_t num_inputs = 0;
|
||||
if (test_bwd_spmd) {
|
||||
num_inputs = 4;
|
||||
} else {
|
||||
num_inputs = 3;
|
||||
}
|
||||
|
||||
EXPECT_EQ(input_shapes.size(), num_inputs)
|
||||
<< "The number of input_shapes must be" << num_inputs << ", but got "
|
||||
<< input_shapes.size();
|
||||
EXPECT_EQ(input_dims_mappings.size(), num_inputs)
|
||||
<< "The number of input_dims_mapping must be" << num_inputs
|
||||
<< ", but got " << input_dims_mappings.size();
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"dp", "mp", "pp"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<DistMetaTensor> dist_meta_tensors;
|
||||
for (size_t i = 0; i < num_inputs; ++i) {
|
||||
TensorDistAttr dist_attr = TensorDistAttr();
|
||||
dist_attr.set_process_mesh(process_mesh);
|
||||
|
||||
const std::vector<int64_t>& shape = input_shapes[i];
|
||||
const std::vector<int64_t>& dim_mapping = input_dims_mappings[i];
|
||||
EXPECT_EQ(shape.size(), dim_mapping.size())
|
||||
<< "The size of shape and dim_mapping for input " << i
|
||||
<< " must be the same, but got " << shape.size()
|
||||
<< " != " << dim_mapping.size();
|
||||
|
||||
dist_attr.set_dims_mapping(dim_mapping);
|
||||
dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
|
||||
|
||||
dist_meta_tensors.push_back(
|
||||
DistMetaTensor(common::make_ddim(shape), dist_attr));
|
||||
}
|
||||
|
||||
phi::distributed::SpmdInfo spmd_info;
|
||||
if (test_bwd_spmd) {
|
||||
spmd_info = phi::distributed::MoECombineGradInferSpmd(dist_meta_tensors[0],
|
||||
dist_meta_tensors[1],
|
||||
dist_meta_tensors[2],
|
||||
dist_meta_tensors[3]);
|
||||
} else {
|
||||
spmd_info = phi::distributed::MoECombineInferSpmd(
|
||||
dist_meta_tensors[0], dist_meta_tensors[1], dist_meta_tensors[2]);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < 2; ++i) {
|
||||
std::vector<ArgDistAttr> dist_attrs;
|
||||
std::vector<std::vector<int64_t>> dims_mappings;
|
||||
if (i == 0) {
|
||||
dist_attrs = spmd_info.first;
|
||||
dims_mappings = expected_dims_mappings.first;
|
||||
} else {
|
||||
dist_attrs = spmd_info.second;
|
||||
dims_mappings = expected_dims_mappings.second;
|
||||
}
|
||||
EXPECT_EQ(dist_attrs.size(), dims_mappings.size())
|
||||
<< "The size of dist_attr and expected_dims must be the same, but got "
|
||||
<< dist_attrs.size() << " != " << dims_mappings.size();
|
||||
|
||||
for (size_t j = 0; j < dist_attrs.size(); ++j) {
|
||||
const ArgDistAttr& inferred_attr = dist_attrs[j];
|
||||
const std::vector<int64_t>& expected_dims_mapping = dims_mappings[j];
|
||||
check_dim_mapping(inferred_attr, expected_dims_mapping);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(MoECombineSPMDRule, test_moe_combine_spmd) {
|
||||
// forward: x, combine_weights, scatter_index -> y
|
||||
// backward: x, combine_weights, scatter_index, grad_y -> grad_x,
|
||||
// grad_combine_weights
|
||||
|
||||
int s = 1024, h = 512, k = 2;
|
||||
const std::vector<std::vector<int64_t>>& forward_input_shapes = {
|
||||
{s * k, h}, {s, k}, {s, k}};
|
||||
const std::vector<std::vector<int64_t>>& backward_input_shapes = {
|
||||
{s * k, h}, {s, k}, {s, k}, {s, h}};
|
||||
|
||||
// replicated case, forward
|
||||
std::vector<std::vector<int64_t>> input_dims_mappings = {
|
||||
{-1, -1}, {-1, -1}, {-1, -1}};
|
||||
std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>
|
||||
expected_dims_mappings = {{{-1, -1}, {-1, -1}, {-1, -1}}, {{-1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
forward_input_shapes, input_dims_mappings, expected_dims_mappings);
|
||||
|
||||
// replicated case, backward
|
||||
input_dims_mappings = {{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}};
|
||||
expected_dims_mappings = {{{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}},
|
||||
{{-1, -1}, {-1, -1}, {-1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
backward_input_shapes, input_dims_mappings, expected_dims_mappings, true);
|
||||
|
||||
// mp case, forward
|
||||
input_dims_mappings = {{1, -1}, {1, -1}, {-1, -1}};
|
||||
expected_dims_mappings = {{{1, -1}, {1, -1}, {1, -1}}, {{1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
forward_input_shapes, input_dims_mappings, expected_dims_mappings);
|
||||
|
||||
// mp case, backward
|
||||
input_dims_mappings = {{1, -1}, {1, -1}, {-1, -1}, {1, -1}};
|
||||
expected_dims_mappings = {{{1, -1}, {1, -1}, {1, -1}, {1, -1}},
|
||||
{{1, -1}, {1, -1}, {1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
backward_input_shapes, input_dims_mappings, expected_dims_mappings, true);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,188 @@
|
||||
/* Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using phi::distributed::ArgDistAttr;
|
||||
using phi::distributed::DistMetaTensor;
|
||||
|
||||
void test_moe_gate_dispatch_spmd(
|
||||
const std::vector<std::vector<int64_t>>& input_shapes,
|
||||
const std::vector<std::vector<int64_t>>& input_dims_mappings,
|
||||
const std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>& expected_dims_mappings,
|
||||
int64_t k,
|
||||
int64_t capacity,
|
||||
bool use_pad,
|
||||
bool test_bwd_spmd = false,
|
||||
bool optional = true) {
|
||||
size_t num_inputs = 0;
|
||||
if (test_bwd_spmd) {
|
||||
num_inputs = 5;
|
||||
} else {
|
||||
num_inputs = 3;
|
||||
}
|
||||
|
||||
EXPECT_EQ(input_shapes.size(), num_inputs)
|
||||
<< "The number of input_shapes must be" << num_inputs << ", but got "
|
||||
<< input_shapes.size();
|
||||
EXPECT_EQ(input_dims_mappings.size(), num_inputs)
|
||||
<< "The number of input_dims_mapping must be" << num_inputs
|
||||
<< ", but got " << input_dims_mappings.size();
|
||||
|
||||
std::vector<int64_t> mesh_shape = {4};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"dpmp"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<DistMetaTensor> dist_meta_tensors;
|
||||
for (size_t i = 0; i < num_inputs; ++i) {
|
||||
TensorDistAttr dist_attr = TensorDistAttr();
|
||||
dist_attr.set_process_mesh(process_mesh);
|
||||
|
||||
const std::vector<int64_t>& shape = input_shapes[i];
|
||||
const std::vector<int64_t>& dim_mapping = input_dims_mappings[i];
|
||||
EXPECT_EQ(shape.size(), dim_mapping.size())
|
||||
<< "The size of shape and dim_mapping for input " << i
|
||||
<< " must be the same, but got " << shape.size()
|
||||
<< " != " << dim_mapping.size();
|
||||
|
||||
dist_attr.set_dims_mapping(dim_mapping);
|
||||
dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
|
||||
|
||||
dist_meta_tensors.push_back(
|
||||
DistMetaTensor(common::make_ddim(shape), dist_attr));
|
||||
}
|
||||
|
||||
phi::distributed::SpmdInfo spmd_info;
|
||||
if (test_bwd_spmd) {
|
||||
spmd_info =
|
||||
phi::distributed::MoEGateDispatchGradInferSpmd(dist_meta_tensors[0],
|
||||
dist_meta_tensors[1],
|
||||
dist_meta_tensors[2],
|
||||
dist_meta_tensors[3],
|
||||
dist_meta_tensors[4],
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
} else {
|
||||
phi::distributed::DistMetaTensor uninitialized_tensor;
|
||||
spmd_info = phi::distributed::MoEGateDispatchInferSpmd(
|
||||
dist_meta_tensors[0],
|
||||
dist_meta_tensors[1],
|
||||
optional ? dist_meta_tensors[2] : uninitialized_tensor,
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < 2; ++i) {
|
||||
std::vector<ArgDistAttr> dist_attrs;
|
||||
std::vector<std::vector<int64_t>> dims_mappings;
|
||||
if (i == 0) {
|
||||
dist_attrs = spmd_info.first;
|
||||
dims_mappings = expected_dims_mappings.first;
|
||||
} else {
|
||||
dist_attrs = spmd_info.second;
|
||||
dims_mappings = expected_dims_mappings.second;
|
||||
}
|
||||
EXPECT_EQ(dist_attrs.size(), dims_mappings.size())
|
||||
<< "The size of dist_attr and expected_dims must be the same, but got "
|
||||
<< dist_attrs.size() << " != " << dims_mappings.size();
|
||||
|
||||
for (size_t j = 0; j < dist_attrs.size(); ++j) {
|
||||
const ArgDistAttr& inferred_attr = dist_attrs[j];
|
||||
const std::vector<int64_t>& expected_dims_mapping = dims_mappings[j];
|
||||
check_dim_mapping(inferred_attr, expected_dims_mapping);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(MoECombineSPMDRule, test_moe_gate_dispatch_spmd) {
|
||||
int64_t s = 1024, h = 512, k = 2, e = 8, capacity = 1024;
|
||||
bool use_pad = true;
|
||||
const std::vector<std::vector<int64_t>>& forward_input_shapes = {
|
||||
{s, h}, {s, e}, {e}};
|
||||
const std::vector<std::vector<int64_t>>& backward_input_shapes = {
|
||||
{s, k}, {k, s}, {s, k}, {e, capacity, h}, {s, k}};
|
||||
|
||||
// replicated case, forward
|
||||
std::vector<std::vector<int64_t>> input_dims_mappings = {
|
||||
{-1, -1}, {-1, -1}, {-1}};
|
||||
std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>
|
||||
expected_dims_mappings = {
|
||||
{{-1, -1}, {-1, -1}, {-1}},
|
||||
{{-1, -1, -1}, {-1, -1}, {-1, -1}, {-1}, {-1, -1}}};
|
||||
test_moe_gate_dispatch_spmd(forward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
|
||||
// replicated case, backward
|
||||
input_dims_mappings = {{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1, -1}, {-1, -1}};
|
||||
expected_dims_mappings = {
|
||||
{{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1, -1}, {-1, -1}},
|
||||
{{-1, -1}, {-1, -1}}};
|
||||
test_moe_gate_dispatch_spmd(backward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad,
|
||||
true);
|
||||
|
||||
// ep case, forward
|
||||
input_dims_mappings = {{0, -1}, {-1, -1}, {-1}};
|
||||
expected_dims_mappings = {{{0, -1}, {0, -1}, {-1}},
|
||||
{{-1, 0, -1}, {0, -1}, {-1, 0}, {-1}, {0, -1}}};
|
||||
test_moe_gate_dispatch_spmd(forward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
|
||||
// ep case, backward
|
||||
input_dims_mappings = {{0, -1}, {-1, 0}, {0, -1}, {-1, 0, -1}, {0, -1}};
|
||||
expected_dims_mappings = {{{0, -1}, {-1, 0}, {0, -1}, {-1, 0, -1}, {0, -1}},
|
||||
{{0, -1}, {0, -1}}};
|
||||
test_moe_gate_dispatch_spmd(backward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad,
|
||||
true);
|
||||
|
||||
// ep, corr_bias is none case, forward
|
||||
input_dims_mappings = {{0, -1}, {-1, -1}, {-1}};
|
||||
expected_dims_mappings = {{{0, -1}, {0, -1}, {}},
|
||||
{{-1, 0, -1}, {0, -1}, {-1, 0}, {-1}, {0, -1}}};
|
||||
test_moe_gate_dispatch_spmd(forward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad,
|
||||
false,
|
||||
false);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,54 @@
|
||||
/* Copyright (c) 2022 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/core/distributed/auto_parallel/process_mesh.h"
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(ProcessMesh, Ctor) {
|
||||
std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
int64_t size = shape[0] * shape[1];
|
||||
ProcessMesh process_mesh(shape, process_ids, dim_names);
|
||||
EXPECT_EQ(process_mesh.shape(), shape);
|
||||
EXPECT_EQ(process_mesh.process_ids(), process_ids);
|
||||
EXPECT_EQ(process_mesh.dim_names()[0], "x");
|
||||
EXPECT_EQ(process_mesh.dim_names()[1], "y");
|
||||
EXPECT_EQ(process_mesh.size(), size);
|
||||
EXPECT_EQ(process_mesh.ndim(), static_cast<int64_t>(shape.size()));
|
||||
EXPECT_EQ(process_mesh.dim_size(0), shape[0]);
|
||||
EXPECT_EQ(process_mesh.dim_size(-1), shape[1]);
|
||||
EXPECT_EQ(process_mesh.dim_size("x"), shape[0]);
|
||||
EXPECT_EQ(process_mesh.dim_size("y"), shape[1]);
|
||||
EXPECT_EQ(process_mesh.empty(), false);
|
||||
EXPECT_EQ(process_mesh.contains(0), true);
|
||||
EXPECT_EQ(process_mesh.contains(6), false);
|
||||
std::stringstream sstream;
|
||||
sstream << process_mesh;
|
||||
EXPECT_EQ(sstream.str(), process_mesh.to_string());
|
||||
auto proto = phi::distributed::to_proto(process_mesh);
|
||||
ProcessMesh new_process_mesh = ProcessMesh::from_proto(proto);
|
||||
EXPECT_EQ(process_mesh, new_process_mesh);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,153 @@
|
||||
/* 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct ReshapeTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> input_dims_mapping;
|
||||
|
||||
// shape attribute
|
||||
std::vector<int64_t> target_shape;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_input_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
};
|
||||
|
||||
TEST(Reshape, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
// test flatten
|
||||
// [4, 6, 8] -> [192]:
|
||||
// [[0], [1], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1]]
|
||||
|
||||
// [4, 6, 8] -> [192]: [[ ], [0], [1]] -> [[ ], [ ], [ ]], [[ ]]
|
||||
|
||||
// [4, 6, 8] -> [192]:
|
||||
// [[0, 1], [ ], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1]]
|
||||
|
||||
// [2, 12, 8] -> [192]:
|
||||
// [[0], [1], [ ]] -> [[0], [ ], [ ]], [[0]]
|
||||
|
||||
// test split
|
||||
|
||||
// [128] -> [4, 6, 8]:
|
||||
// [[0, 1]] -> [[0, 1]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [128] -> [6, 4, 8]:
|
||||
// [[0, 1]] -> [[ ], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [2, 12, 8]
|
||||
// [[0], [1], [ ]] -> [[0], [ ], [ ]], [[0], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [2, 12, 8]
|
||||
// [[0, 1], [ ], [ ]] -> [[ ], [ ], [ ]], [[ ], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [12, 2, 8]:
|
||||
// [[0], [1], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [12, 2, 8]:
|
||||
// [[0, 1], [ ], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[0], [1], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[ ], [0], [1]] -> [[ ], [ ], [ ]], [[ ], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[0], [ ], [1]] -> [[0], [ ], [ ]], [[0], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[0, 1], [ ], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
std::vector<ReshapeTestCase> test_cases = {
|
||||
// input_shape, input_dims_mapping, target_shape,
|
||||
// expected_input_dims_mapping, expected_output_dims_mapping
|
||||
|
||||
// test flatten
|
||||
{{4, 6, 8}, {{0}, {1}, {}}, {192}, {{0, 1}, {}, {}}, {{0, 1}}},
|
||||
{{4, 6, 8}, {{}, {0}, {1}}, {192}, {{}, {}, {}}, {{}}},
|
||||
{{4, 6, 8}, {{0, 1}, {}, {}}, {192}, {{0, 1}, {}, {}}, {{0, 1}}},
|
||||
{{2, 12, 8}, {{0}, {1}, {}}, {192}, {{0}, {}, {}}, {{0}}},
|
||||
|
||||
// test split
|
||||
{{192}, {{0, 1}}, {4, 6, 8}, {{0, 1}}, {{0, 1}, {}, {}}},
|
||||
{{192}, {{0, 1}}, {6, 4, 8}, {{}}, {{}, {}, {}}},
|
||||
|
||||
// test combination
|
||||
{{4, 6, 8}, {{0}, {1}, {}}, {2, 12, 8}, {{0}, {}, {}}, {{0}, {}, {}}},
|
||||
{{4, 6, 8}, {{0, 1}, {}, {}}, {2, 12, 8}, {{}, {}, {}}, {{}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0}, {1}, {}},
|
||||
{12, 2, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{12, 2, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0}, {1}, {}},
|
||||
{8, 6, 4},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8}, {{}, {0}, {1}}, {8, 6, 4}, {{}, {}, {}}, {{}, {}, {}}},
|
||||
{{4, 6, 8}, {{0}, {}, {1}}, {8, 6, 4}, {{0}, {}, {}}, {{0}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{8, 6, 4},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8}, {{}, {}, {0, 1}}, {24, 2, 4}, {{}, {}, {}}, {{}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{}, {}, {0, 1}},
|
||||
{24, 4, 2},
|
||||
{{}, {}, {0, 1}},
|
||||
{{}, {0, 1}, {}}},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.input_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::ReshapeInferSpmd(x, tc.target_shape);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_input_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
}
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,243 @@
|
||||
/* 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct SoftmaxTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> input_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_input_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
};
|
||||
|
||||
struct SoftmaxGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> out_shape;
|
||||
std::vector<std::vector<int64_t>> out_dims_mapping;
|
||||
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
};
|
||||
|
||||
TEST(SoftmaxInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<SoftmaxTestCase> test_cases = {
|
||||
// shape = [32, 48, 128], axis = 0
|
||||
// [[0,1],[2],[]] -> [[],[2],[]], [[],[2],[]]
|
||||
{{32, 48, 128}, {{0, 1}, {2}, {}}, 0, {{}, {2}, {}}, {{}, {2}, {}}},
|
||||
{{32, 48, 128}, {{0, 1}, {2}, {}}, -3, {{}, {2}, {}}, {{}, {2}, {}}},
|
||||
|
||||
// shape = [32, 48, 128], axis = 1
|
||||
// [[0,1],[2],[]] -> [[0, 1],[],[]], [[0, 1],[],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}}};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.input_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::SoftmaxInferSpmd(x, tc.axis);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_input_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(SoftmaxGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<SoftmaxGradTestCase> test_cases = {
|
||||
// out_shape = [32, 48, 128], out_grad_shape = [32, 48, 128], axis = 0
|
||||
// [[0,1],[2],[]], [[0,1],[2],[]] -> [[],[2],[]], [[],[2],[]], [[],[2],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
0,
|
||||
{{}, {2}, {}},
|
||||
{{}, {2}, {}},
|
||||
{{}, {2}, {}}},
|
||||
// axis = 0
|
||||
// [[0,1],[2],[]], [[0],[1,2],[]] -> [[],[1,2],[]], [[],[1, 2],[]],
|
||||
// [[],[1,2],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
{32, 48, 128},
|
||||
{{0}, {1, 2}, {}},
|
||||
0,
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}}},
|
||||
// axis = 1
|
||||
// [[0,1],[2],[]], [[2],[0,1],[]] -> [[0,1,2],[],[]], [[0, 1, 2],[],[]],
|
||||
// [[0, 1, 2],[],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
{32, 48, 128},
|
||||
{{2}, {0, 1}, {}},
|
||||
1,
|
||||
{{0, 1, 2}, {}, {}},
|
||||
{{0, 1, 2}, {}, {}},
|
||||
{{0, 1, 2}, {}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1],[]], [[],[0,1],[]] -> [[],[0,1],[]], [[],[0,1],[]],
|
||||
// [[],[0,1],[]]
|
||||
{{32, 48, 128},
|
||||
{{0}, {1}, {}},
|
||||
{32, 48, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1],[]], [[0,1],[],[]] -> [[0,1],[],[]], [[0, 1],[],[]],
|
||||
// [[0,1],[],[]]
|
||||
{{32, 48, 128},
|
||||
{{0}, {1}, {}},
|
||||
{32, 48, 128},
|
||||
{{0, 1}, {}, {}},
|
||||
2,
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1,2],[]], [[],[0,1],[]] -> [[0],[1,2],[]], [[0],[1,2],[]],
|
||||
// [[0],[1,2],[]]
|
||||
{{32, 48, 128},
|
||||
{{0}, {1, 2}, {}},
|
||||
{32, 48, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1,2],[]], [[],[0,1],[]] -> [[0],[1,2],[]], [[0],[1,2],[]],
|
||||
// [[0],[1,2],[]]
|
||||
{{2, 4, 128},
|
||||
{{0}, {1, 2}, {}},
|
||||
{2, 4, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}}},
|
||||
// axis = 2
|
||||
// [[],[1,2],[]], [[],[0,1],[]] -> [[],[1,2],[]], [[],[1,2],[]],
|
||||
// [[],[1,2],[]]
|
||||
{{2, 4, 128},
|
||||
{{}, {1, 2}, {}},
|
||||
{2, 4, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}}},
|
||||
// axis = 1
|
||||
// [[0,1],[],[]], [[],[],[2]] -> [[0,1],[],[2]], [[0,1],[],[2]],
|
||||
// [[0,1],[],[2]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {}, {}},
|
||||
{32, 48, 128},
|
||||
{{}, {}, {2}},
|
||||
1,
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}}},
|
||||
// Note: just for pass coverage ci: axis = 2
|
||||
// [[0],[0,1],[]], [[],[],[]] -> [[],[0,1],[]], [[],[0,1],[]],
|
||||
// [[],[0,1],[]]
|
||||
{{2, 4, 128},
|
||||
{{0}, {0, 1}, {}},
|
||||
{2, 4, 128},
|
||||
{{}, {}, {}},
|
||||
2,
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}}}};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(tc.out_dims_mapping);
|
||||
out_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.out_shape), out_dist_attr);
|
||||
TensorDistAttr out_grad_attr = TensorDistAttr();
|
||||
out_grad_attr.set_process_mesh(process_mesh);
|
||||
out_grad_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::SoftmaxGradInferSpmd(out, out_grad, tc.axis);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
// [[0,1],[2]] [[2],[]]
|
||||
@@ -0,0 +1,143 @@
|
||||
/* Copyright (c) 2022 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(SoftmaxGradInferSpmd, Ctor) {
|
||||
// Sharding along axes besides softmax axis.
|
||||
std::vector<int64_t> x_shape = {36, 48};
|
||||
std::vector<int64_t> out_grad_shape = {36, 48};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1}));
|
||||
out_grad_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor out_grad(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
int axis = 1;
|
||||
|
||||
auto spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on other axes." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// Sharding along softmax axis.
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1}));
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1}));
|
||||
x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
axis = 1;
|
||||
|
||||
spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on softmax axis." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// Sharding on multi axes.
|
||||
x_shape = {10, 36, 48, 24};
|
||||
out_grad_shape = {10, 36, 48, 24};
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
|
||||
x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
axis = 1;
|
||||
|
||||
spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({0, -1, -1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, -1, -1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, -1, -1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on multi axes." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// Sharding on multi axes.
|
||||
x_shape = {10, 36, 48, 24};
|
||||
out_grad_shape = {10, 36, 48, 24};
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, 1, -1}));
|
||||
x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
axis = 1;
|
||||
|
||||
spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({0, -1, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, -1, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, -1, 1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on multi axes." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,118 @@
|
||||
/* Copyright (c) 2022 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
const std::vector<int64_t>& get_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.dims_mapping();
|
||||
}
|
||||
|
||||
bool is_partial(const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.is_partial();
|
||||
}
|
||||
|
||||
const std::set<int64_t> get_partial_dims(
|
||||
const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.partial_dims();
|
||||
}
|
||||
|
||||
const std::vector<std::vector<int64_t>>& get_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.multi_dims_mapping();
|
||||
}
|
||||
|
||||
void check_dim_mapping(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<int64_t>& dim_mapping,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(get_dims_mapping(dist_attr), dim_mapping) << line;
|
||||
}
|
||||
|
||||
void check_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<std::vector<int64_t>>& dim_mapping,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(get_multi_dims_mapping(dist_attr), dim_mapping) << line;
|
||||
}
|
||||
|
||||
void check_empty_dist_attr(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr),
|
||||
phi::distributed::TensorDistAttr());
|
||||
}
|
||||
|
||||
void check_partial_dims(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::set<int64_t>& dims,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(get_partial_dims(dist_attr), dims) << line;
|
||||
}
|
||||
|
||||
void clean_partial_status(phi::distributed::ArgDistAttr* dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
|
||||
auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
|
||||
tensor_attr.clean_partial_status();
|
||||
}
|
||||
|
||||
void clean_partial_dims(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
|
||||
auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
|
||||
tensor_attr.clean_partial_dims(dims);
|
||||
}
|
||||
|
||||
void set_partial_status(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
|
||||
auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
|
||||
tensor_attr.set_partial_status(dims);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,74 @@
|
||||
/* Copyright (c) 2022 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 <iostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
#include "glog/logging.h"
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
|
||||
#include "paddle/phi/core/distributed/type_defs.h"
|
||||
#include "paddle/phi/infermeta/spmd_rules/replicated.h"
|
||||
#include "paddle/phi/infermeta/spmd_rules/rules.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using phi::distributed::ProcessMesh;
|
||||
using phi::distributed::TensorDistAttr;
|
||||
|
||||
const std::vector<int64_t>& get_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
bool is_partial(const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
const std::set<int64_t> get_partial_dims(
|
||||
const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
const std::vector<std::vector<int64_t>>& get_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
void check_dim_mapping(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<int64_t>& dim_mapping,
|
||||
const std::string& line = "");
|
||||
|
||||
void check_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<std::vector<int64_t>>& dim_mapping,
|
||||
const std::string& line = "");
|
||||
|
||||
void check_empty_dist_attr(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::string& line = "");
|
||||
|
||||
void check_partial_dims(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::set<int64_t>& dims,
|
||||
const std::string& line = "");
|
||||
|
||||
void clean_partial_status(phi::distributed::ArgDistAttr* dist_attr);
|
||||
|
||||
void clean_partial_dims(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims);
|
||||
|
||||
void set_partial_status(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims);
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,215 @@
|
||||
/* 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/tile.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct TileTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
// repeat_times attribute
|
||||
phi::IntArray repeat_times;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
};
|
||||
|
||||
struct TileGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// repeat_times attribute
|
||||
phi::IntArray repeat_times;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
|
||||
std::set<int64_t> partial_dims;
|
||||
};
|
||||
|
||||
TEST(TileInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<TileTestCase> test_cases = {
|
||||
// shape = [8, 16, 24], repeat_times = {2, 2, 1, 1}
|
||||
// [[0],[],[1,2]] -> [[],[],[1,2]], [[],[],[],[1,2]]
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0}, {}, {1, 2}},
|
||||
phi::IntArray({2, 2, 1, 1}),
|
||||
{{}, {}, {1, 2}},
|
||||
{{}, {}, {}, {1, 2}},
|
||||
},
|
||||
|
||||
// shape = [8, 16, 24], repeat_times = {1, 2}
|
||||
// [[0,1],[],[2]] -> [[0,1],[],[]], [[0,1],[],[]]
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
phi::IntArray({1, 2}),
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
},
|
||||
|
||||
// shape = [8, 16, 24], repeat_times = {}
|
||||
// [[0,1],[],[2]] -> [[0,1],[],[2]], [[0,1],[],[2]]
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
phi::IntArray({}),
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::TileInferSpmdDynamic(x, tc.repeat_times);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TileGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<TileGradTestCase> test_cases = {
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [2, 16, 16, 24], repeat_times =
|
||||
// {2, 2, 1, 1}
|
||||
// [[0],[],[1,2]], [[],[],[],[1,2]] -> [[],[],[1,2]], [[],[],[],[1,2]],
|
||||
// [[],[],[1,2]], partial on {}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0}, {}, {1, 2}},
|
||||
{2, 16, 16, 24},
|
||||
{{}, {}, {}, {1, 2}},
|
||||
phi::IntArray({2, 2, 1, 1}),
|
||||
{{}, {}, {1, 2}},
|
||||
{{}, {}, {}, {1, 2}},
|
||||
{{}, {}, {1, 2}},
|
||||
{},
|
||||
},
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [8, 16, 48], repeat_times = {1,
|
||||
// 2}
|
||||
// [[0,1],[],[2]], [[0,1],[],[2]] -> [[0,1],[],[]], [[0,1],[],[]]],
|
||||
// [[0,1],[],[]], partial on {}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
{8, 16, 48},
|
||||
{{0, 1}, {}, {2}},
|
||||
phi::IntArray({1, 2}),
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{},
|
||||
},
|
||||
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [8, 16, 24], repeat_times = {}
|
||||
// [[0,1],[],[2]], [[0],[1],[2]] -> [[0],[1],[2]], [[0],[1],[2]],
|
||||
// [[0],[1],[2]], partial on {}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
{8, 16, 24},
|
||||
{{0}, {1}, {2}},
|
||||
phi::IntArray({}),
|
||||
{{0}, {1}, {2}},
|
||||
{{0}, {1}, {2}},
|
||||
{{0}, {1}, {2}},
|
||||
{},
|
||||
},
|
||||
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [8, 16, 16, 24], repeat_times =
|
||||
// {8, 2, 1, 1}
|
||||
// [[0],[],[]], [[1,2],[],[],[]] -> [[],[],[]], [[1,2],[],[],[]],
|
||||
// [[],[],[]], partial on {1,2}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0}, {}, {}},
|
||||
{8, 16, 16, 24},
|
||||
{{1, 2}, {}, {}, {}},
|
||||
phi::IntArray({8, 2, 1, 1}),
|
||||
{{}, {}, {}},
|
||||
{{1, 2}, {}, {}, {}},
|
||||
{{}, {}, {}},
|
||||
{1, 2},
|
||||
},
|
||||
};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
TensorDistAttr out_grad_attr = TensorDistAttr();
|
||||
out_grad_attr.set_process_mesh(process_mesh);
|
||||
out_grad_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::TileGradInferSpmdDynamic(
|
||||
x, out_grad, tc.repeat_times);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,63 @@
|
||||
/* 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(Tile, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> shape = {6, 8, 10};
|
||||
std::vector<int64_t> dims_mapping = {0, -1, 1};
|
||||
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims({false, false, false});
|
||||
phi::distributed::DistMetaTensor x =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
|
||||
std::vector<int64_t> repeat_times = {2, 2, 1, 1};
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::TileInferSpmd(x, repeat_times);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_dim_mapping(forward_spmd_info.first[0], {-1, -1, 1});
|
||||
check_dim_mapping(forward_spmd_info.second[0], {-1, -1, -1, 1});
|
||||
check_partial_dims(forward_spmd_info.second[0], {});
|
||||
|
||||
// test backward
|
||||
auto out_grad_dist_attr =
|
||||
PADDLE_GET_CONST(TensorDistAttr, forward_spmd_info.second[0]);
|
||||
out_grad_dist_attr.set_dims_mapping({0, -1, -1, 1});
|
||||
phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim({2, 12, 8, 10}), out_grad_dist_attr);
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
TileGradInferSpmd(x, out_grad, repeat_times);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_dim_mapping(backward_spmd_info.first[0], {-1, -1, 1});
|
||||
check_dim_mapping(backward_spmd_info.first[1], {0, -1, -1, 1});
|
||||
check_dim_mapping(backward_spmd_info.second[0], {-1, -1, 1});
|
||||
check_partial_dims(backward_spmd_info.second[0], {0});
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,89 @@
|
||||
/* 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct TransposeTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> input_dims_mapping;
|
||||
|
||||
// shape attribute
|
||||
std::vector<int> perm;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_input_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
};
|
||||
|
||||
TEST(Transpose, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
//
|
||||
std::vector<TransposeTestCase> test_cases = {
|
||||
// input_shape, input_dims_mapping, perm,
|
||||
// expected_input_dims_mapping, expected_output_dims_mapping
|
||||
|
||||
{{64, 48, 36, 24},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{1, 0, 2, 3},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{}, {0, 1}, {}, {}}},
|
||||
{{64, 48, 36, 24},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{0, 1, 2, 3},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{0, 1}, {}, {}, {}}},
|
||||
{{64, 48, 36, 24},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{0, 2, 3, 1},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}, {}}},
|
||||
{{64, 48, 36, 24},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{-1, 0, -2, 1},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{{}, {}, {0, 1}, {}}},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.input_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::TransposeInferSpmd(x, tc.perm);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_input_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
}
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
Reference in New Issue
Block a user