wehub-resource-sync
243 lines
9.4 KiB
C++
243 lines
9.4 KiB
C++
/* 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/infermeta/spmd_rules/split.h"
|
|
|
|
#include "glog/logging.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/utils.h"
|
|
#include "paddle/phi/infermeta/spmd_rules/utils.h"
|
|
|
|
namespace phi::distributed {
|
|
|
|
SpmdInfo SplitWithNumInferSpmd(const DistMetaTensor& x, int num, int axis) {
|
|
// Step0: Verify input args based on split logic
|
|
auto x_shape = vectorize(x.dims());
|
|
int x_ndim = static_cast<int>(x_shape.size());
|
|
const auto& x_dist_attr_src = x.dist_attr();
|
|
std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
|
|
PADDLE_ENFORCE_EQ(
|
|
x_ndim,
|
|
x_dims_mapping.size(),
|
|
common::errors::InvalidArgument("The Tensor X's rank [%d] and X's "
|
|
"dims_mapping size [%d] are not matched.",
|
|
x_ndim,
|
|
x_dims_mapping.size()));
|
|
|
|
// Step1: Build Einsum Notation
|
|
std::string alphabet = "abcdefghijlmnopqrstuvwxyz";
|
|
if (axis < 0) {
|
|
axis += x_ndim;
|
|
}
|
|
|
|
// get einsum notation for input, use a special
|
|
// notation 'k' to mark the split axis in input
|
|
std::string x_axes = alphabet.substr(0, x_ndim);
|
|
x_axes[axis] = 'k';
|
|
|
|
// get einsum notation for output
|
|
std::string out_axes(x_axes);
|
|
// the split axis cannot be sharded, set its notation
|
|
// with the special '1' to set its dim mapping to -1.
|
|
out_axes[axis] = '1';
|
|
|
|
// Step2: Sharding Propagation
|
|
// Step2.1: merge input shardings
|
|
std::unordered_map<std::string, int64_t> axis_to_dim_map =
|
|
ShardingMergeForTensors({{x_axes, x_dims_mapping}});
|
|
|
|
// Step2.2: infer output dims mapping from merged input dims mapping
|
|
std::vector<int64_t> out_dims_mapping =
|
|
GetDimsMappingForAxes(out_axes, axis_to_dim_map);
|
|
|
|
// get the dist attributes for all outputs, the
|
|
// dist attributes are same for all outputs.
|
|
std::vector<TensorDistAttr> out_dist_attrs;
|
|
for (int i = 0; i < num; i++) {
|
|
out_dist_attrs.emplace_back(CopyTensorDistAttrForOutput(x_dist_attr_src));
|
|
out_dist_attrs[i].set_dims_mapping(out_dims_mapping);
|
|
}
|
|
|
|
// Step2.3 get new dist attribute for input. the split
|
|
// cannot be sharded, if it is sharded, set it to replicated.
|
|
TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
x_dims_mapping[axis] = -1;
|
|
x_dist_attr_dst.set_dims_mapping(x_dims_mapping);
|
|
|
|
// Step3 Handle input tensor partial (TODO)
|
|
VLOG(4) << "SplitWithNumInferSpmd:";
|
|
VLOG(4) << "Einsum Notation: " << x_axes << "-->" << out_axes;
|
|
VLOG(4) << "Input shape: [" << str_join(x_shape) << "] "
|
|
<< "src_dims_mapping: [" << str_join(x_dist_attr_src.dims_mapping())
|
|
<< "] "
|
|
<< "dst_dims_mapping: [" << str_join(x_dims_mapping) << "]";
|
|
for (int64_t i = 0; i < num; i++) {
|
|
VLOG(4) << "Output" << std::to_string(i) << " dims_mapping: ["
|
|
<< str_join(out_dims_mapping) << "]";
|
|
}
|
|
VLOG(4) << std::endl;
|
|
// TODO(liuzhenhai): remedy this
|
|
// should return list in list []
|
|
// return {{x_dist_attr_dst}, {out_dist_attrs}};
|
|
return {{x_dist_attr_dst}, ToArgDistAttr(out_dist_attrs)};
|
|
}
|
|
|
|
SpmdInfo SplitWithNumInferSpmdReverse(
|
|
const DistMetaTensor& x,
|
|
const std::vector<const DistMetaTensor*>& outs,
|
|
int num,
|
|
int axis) {
|
|
// Step0: Verify input args based on split logic
|
|
int nouts = static_cast<int>(outs.size());
|
|
int out_ndim = static_cast<int>(vectorize(outs[0]->dims()).size());
|
|
auto x_shape = vectorize(x.dims());
|
|
int x_ndim = static_cast<int>(x_shape.size());
|
|
const auto& x_dist_attr = x.dist_attr();
|
|
std::vector<int64_t> x_dims_mapping = x_dist_attr.dims_mapping();
|
|
PADDLE_ENFORCE_EQ(nouts,
|
|
num,
|
|
common::errors::InvalidArgument(
|
|
"The size of Output Tensors [%d] is not equal "
|
|
"to the specified split number [%d]",
|
|
nouts,
|
|
num));
|
|
PADDLE_ENFORCE_EQ(
|
|
x_ndim,
|
|
out_ndim,
|
|
common::errors::InvalidArgument("The Tensor X's rank [%d] is not equal "
|
|
"to the Tensor Out's rank [%d]",
|
|
x_ndim,
|
|
out_ndim));
|
|
for (int i = 0; i < num; i++) {
|
|
auto shape = vectorize(outs[i]->dims());
|
|
int ndim = static_cast<int>(shape.size());
|
|
auto dist_attr = outs[i]->dist_attr();
|
|
int dims_mapping_size = static_cast<int>(dist_attr.dims_mapping().size());
|
|
PADDLE_ENFORCE_EQ(ndim,
|
|
dims_mapping_size,
|
|
common::errors::InvalidArgument(
|
|
"The Tensor Out[%d]'s rank [%d] and Its "
|
|
"dims_mapping size [%d] are not matched.",
|
|
i,
|
|
ndim,
|
|
dims_mapping_size));
|
|
}
|
|
|
|
// Step1: Build Einsum Notation
|
|
if (axis < 0) {
|
|
axis += x_ndim;
|
|
}
|
|
std::string alphabet = "abcdefghijlmnopqrstuvwxyz";
|
|
|
|
// get einsum notation for input, use a special
|
|
// notation 'k' to mark the split axis in input
|
|
std::string x_axes = alphabet.substr(0, x_ndim);
|
|
x_axes[axis] = 'k';
|
|
|
|
// get einsum notation for output
|
|
std::string out_axes(x_axes);
|
|
out_axes[axis] = 'k';
|
|
|
|
// Step2: Sharding Propagation
|
|
// Step2.1: merge output shardings
|
|
std::vector<std::pair<std::string, std::vector<int64_t>>> axes_sharding_info;
|
|
for (int i = 0; i < nouts; i++) {
|
|
std::vector<int64_t> out_dims_mapping = outs[i]->dist_attr().dims_mapping();
|
|
axes_sharding_info.emplace_back(out_axes, out_dims_mapping);
|
|
}
|
|
std::unordered_map<std::string, int64_t> axis_to_dim_map =
|
|
ShardingMergeForTensors(axes_sharding_info);
|
|
|
|
// Step2.2: infer input dims mapping from output dims mapping
|
|
// the split axis in input is set to -1.
|
|
x_dims_mapping = GetDimsMappingForAxes(x_axes, axis_to_dim_map, true);
|
|
x_dims_mapping[axis] = -1;
|
|
|
|
auto x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr);
|
|
x_dist_attr_dst.set_dims_mapping(x_dims_mapping);
|
|
|
|
// step2.3 get new dist attribute for output. the split
|
|
// cannot be sharded, if it is sharded, set it to replicated.
|
|
std::vector<TensorDistAttr> out_dist_attrs;
|
|
for (int i = 0; i < nouts; i++) {
|
|
out_dist_attrs.emplace_back(
|
|
CopyTensorDistAttrForOutput(outs[i]->dist_attr()));
|
|
std::vector<int64_t> out_dims_mapping =
|
|
GetDimsMappingForAxes(out_axes, axis_to_dim_map, true);
|
|
out_dims_mapping[axis] = -1;
|
|
out_dist_attrs[i].set_dims_mapping(out_dims_mapping);
|
|
}
|
|
|
|
// step3 Handle input tensor partial (TODO)
|
|
|
|
VLOG(4) << "SplitWithNumInferSpmdReverse:";
|
|
VLOG(4) << "Einsum Notation: " << x_axes << "-->" << out_axes;
|
|
for (int i = 0; i < nouts; i++) {
|
|
VLOG(4) << "Output" << std::to_string(i) << " shape: ["
|
|
<< str_join(vectorize(outs[i]->dims())) << "] "
|
|
<< "src_dims_mapping: ["
|
|
<< str_join(outs[i]->dist_attr().dims_mapping()) << "] "
|
|
<< "dst_dims_mapping: ["
|
|
<< str_join(out_dist_attrs[i].dims_mapping()) << "]";
|
|
}
|
|
VLOG(4) << "Input shape: [" << str_join(x_shape) << "] "
|
|
<< "dims_mapping: [" << str_join(x_dims_mapping) << "]\n\n";
|
|
// TODO(liuzhenhai): remedy this
|
|
// return {{x_dist_attr}, {out_dist_attrs}};
|
|
return {{x_dist_attr_dst}, ToArgDistAttr(out_dist_attrs)};
|
|
}
|
|
|
|
SpmdInfo SplitInferSpmd(const DistMetaTensor& x,
|
|
const std::vector<int>& sections,
|
|
int axis) {
|
|
int num = static_cast<int>(sections.size());
|
|
return SplitWithNumInferSpmd(x, num, axis);
|
|
}
|
|
|
|
SpmdInfo SplitInferSpmdDynamic(const DistMetaTensor& x,
|
|
const std::vector<int64_t>& sections,
|
|
const Scalar& axis) {
|
|
int num = static_cast<int>(sections.size());
|
|
return SplitWithNumInferSpmdDynamic(x, num, axis);
|
|
}
|
|
|
|
SpmdInfo SplitInferSpmdReverse(const DistMetaTensor& x,
|
|
const std::vector<const DistMetaTensor*>& outs,
|
|
const std::vector<int>& sections,
|
|
int axis) {
|
|
int num = static_cast<int>(sections.size());
|
|
return SplitWithNumInferSpmdReverse(x, outs, num, axis);
|
|
}
|
|
|
|
SpmdInfo SplitWithNumInferSpmdDynamic(const DistMetaTensor& x,
|
|
int num,
|
|
const Scalar& axis) {
|
|
auto tmp = SplitWithNumInferSpmd(x, num, axis.to<int32_t>());
|
|
// bridge the diff concerning vector output between static and dynamic auto
|
|
// parallel ToDo(liuzhenhai): unify the difference between static and dynamic
|
|
SpmdInfo ret;
|
|
ret.first = tmp.first;
|
|
std::vector<TensorDistAttr> out_dist_attrs;
|
|
out_dist_attrs.reserve(tmp.second.size());
|
|
for (const auto& out : tmp.second) {
|
|
out_dist_attrs.push_back(PADDLE_GET_CONST(TensorDistAttr, out));
|
|
}
|
|
ret.second = {out_dist_attrs};
|
|
return ret;
|
|
}
|
|
|
|
} // namespace phi::distributed
|