wehub-resource-sync
428 lines
17 KiB
C++
428 lines
17 KiB
C++
// 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/batch_norm.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/spmd_rule_macro_define.h"
|
|
#include "paddle/phi/infermeta/spmd_rules/utils.h"
|
|
|
|
namespace phi::distributed {
|
|
|
|
SpmdInfo BatchNormInferSpmd(const DistMetaTensor& x,
|
|
const DistMetaTensor& mean,
|
|
const DistMetaTensor& variance,
|
|
const DistMetaTensor& scale,
|
|
const DistMetaTensor& bias,
|
|
const bool is_test,
|
|
const float momentum,
|
|
const float epsilon,
|
|
const std::string& data_format,
|
|
const bool use_global_stats,
|
|
const bool trainable_statistics) {
|
|
// Step0: verify input args based on batch_norm logic
|
|
auto x_shape = vectorize(x.dims());
|
|
auto mean_shape = vectorize(mean.dims());
|
|
auto variance_shape = vectorize(variance.dims());
|
|
auto scale_shape = vectorize(scale.dims());
|
|
auto bias_shape = vectorize(bias.dims());
|
|
int x_ndim = static_cast<int>(x_shape.size());
|
|
int mean_ndim = static_cast<int>(mean_shape.size());
|
|
int variance_ndim = static_cast<int>(variance_shape.size());
|
|
int scale_ndim = static_cast<int>(scale_shape.size());
|
|
int bias_ndim = static_cast<int>(bias_shape.size());
|
|
TensorDistAttr x_dist_attr_src = x.dist_attr();
|
|
TensorDistAttr mean_dist_attr_src = mean.dist_attr();
|
|
TensorDistAttr variance_dist_attr_src = variance.dist_attr();
|
|
TensorDistAttr scale_dist_attr_src = scale.dist_attr();
|
|
TensorDistAttr bias_dist_attr_src = bias.dist_attr();
|
|
std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
|
|
std::vector<int64_t> mean_dims_mapping = mean.dist_attr().dims_mapping();
|
|
std::vector<int64_t> variance_dims_mapping =
|
|
variance.dist_attr().dims_mapping();
|
|
std::vector<int64_t> scale_dims_mapping = scale.dist_attr().dims_mapping();
|
|
std::vector<int64_t> bias_dims_mapping = bias.dist_attr().dims_mapping();
|
|
|
|
PADDLE_ENFORCE_GE(
|
|
x_ndim,
|
|
2,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of x in batch_norm should be greater than 1, but got [%d].",
|
|
x_ndim));
|
|
PADDLE_ENFORCE_LE(
|
|
x_ndim,
|
|
5,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of x in batch_norm should be less than 6, but got [%d].",
|
|
x_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
mean_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of mean in batch_norm should be 1, but got [%d].",
|
|
mean_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
variance_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of variance in batch_norm should be 1, but got [%d].",
|
|
variance_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
scale_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of scale in batch_norm should be 1, but got [%d].",
|
|
scale_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
bias_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of bias in batch_norm should be 1, but got [%d].",
|
|
bias_ndim));
|
|
|
|
// Step1: Build Einsum Notation
|
|
|
|
std::string alphabet = "ijklmnopqrstuvwxyz";
|
|
// get input notation
|
|
// The mean and variance was flatten at C axis
|
|
std::string x_axes(x_ndim, '1');
|
|
for (int i = 0; i < x_ndim; ++i) {
|
|
x_axes[i] = alphabet[i];
|
|
}
|
|
int c_index = data_format[1] == 'C' ? 1 : x_ndim - 1;
|
|
std::string mean_axes(1, x_axes[c_index]);
|
|
std::string variance_axes(1, x_axes[c_index]);
|
|
std::string scale_axes(1, x_axes[c_index]);
|
|
std::string bias_axes(1, x_axes[c_index]);
|
|
|
|
// get output notation
|
|
std::string out_axes = x_axes;
|
|
|
|
// Step2: Sharding Propagation
|
|
// Step2.1: merge input sharding
|
|
// Only C axis can be shard.
|
|
auto c_dim =
|
|
x_dims_mapping[c_index]; // type: "NC"、"NCL"、"NLC"、"NCHW"、"NHWC"" and
|
|
// "NCDHW"
|
|
|
|
for (int i = 0; i < x_ndim; ++i) {
|
|
x_dims_mapping[i] = -1;
|
|
}
|
|
x_dims_mapping[c_index] = c_dim;
|
|
std::unordered_map<std::string, int64_t> axis_to_dim_map =
|
|
ShardingMergeForTensors({{x_axes, x_dims_mapping}});
|
|
|
|
// Step2.2: infer output dims mapping
|
|
TensorDistAttr out_dist_attr = CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
TensorDistAttr mean_dist_attr = CopyTensorDistAttrForOutput(mean.dist_attr());
|
|
TensorDistAttr variance_dist_attr =
|
|
CopyTensorDistAttrForOutput(variance.dist_attr());
|
|
TensorDistAttr saved_mean_dist_attr =
|
|
CopyTensorDistAttrForOutput(mean.dist_attr());
|
|
TensorDistAttr saved_variance_dist_attr =
|
|
CopyTensorDistAttrForOutput(variance.dist_attr());
|
|
TensorDistAttr reserve_space_dist_attr =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
out_dist_attr.set_dims_mapping(
|
|
GetDimsMappingForAxes(out_axes, axis_to_dim_map));
|
|
mean_dist_attr.set_dims_mapping(
|
|
GetDimsMappingForAxes(mean_axes, axis_to_dim_map));
|
|
variance_dist_attr.set_dims_mapping(
|
|
GetDimsMappingForAxes(variance_axes, axis_to_dim_map));
|
|
saved_mean_dist_attr.set_dims_mapping(
|
|
GetDimsMappingForAxes(mean_axes, axis_to_dim_map));
|
|
saved_variance_dist_attr.set_dims_mapping(
|
|
GetDimsMappingForAxes(variance_axes, axis_to_dim_map));
|
|
reserve_space_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
|
|
|
|
// Step2.3: update input dims mapping
|
|
// mean, variance, mean_out, variance_out and
|
|
TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
TensorDistAttr scale_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(scale.dist_attr());
|
|
TensorDistAttr bias_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(bias.dist_attr());
|
|
TensorDistAttr mean_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(mean.dist_attr());
|
|
TensorDistAttr variance_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(variance.dist_attr());
|
|
scale_dist_attr_dst.set_dims_mapping(std::vector<int64_t>{-1});
|
|
bias_dist_attr_dst.set_dims_mapping(std::vector<int64_t>{-1});
|
|
variance_dist_attr_dst.set_dims_mapping(
|
|
GetDimsMappingForAxes(variance_axes, axis_to_dim_map));
|
|
mean_dist_attr_dst.set_dims_mapping(
|
|
GetDimsMappingForAxes(mean_axes, axis_to_dim_map));
|
|
|
|
x_dist_attr_dst.set_dims_mapping(x_dims_mapping);
|
|
|
|
LOG_SPMD_INPUT(x);
|
|
LOG_SPMD_INPUT(mean);
|
|
LOG_SPMD_INPUT(variance);
|
|
LOG_SPMD_INPUT(scale);
|
|
LOG_SPMD_INPUT(bias);
|
|
LOG_SPMD_OUTPUT(out_dist_attr);
|
|
LOG_SPMD_OUTPUT(mean_dist_attr);
|
|
LOG_SPMD_OUTPUT(variance_dist_attr);
|
|
LOG_SPMD_OUTPUT(saved_mean_dist_attr);
|
|
LOG_SPMD_OUTPUT(saved_variance_dist_attr);
|
|
LOG_SPMD_OUTPUT(reserve_space_dist_attr);
|
|
return {{x_dist_attr_dst,
|
|
mean_dist_attr_dst,
|
|
variance_dist_attr_dst,
|
|
scale_dist_attr_dst,
|
|
bias_dist_attr_dst},
|
|
{out_dist_attr,
|
|
mean_dist_attr,
|
|
variance_dist_attr,
|
|
saved_mean_dist_attr,
|
|
saved_variance_dist_attr,
|
|
reserve_space_dist_attr}};
|
|
}
|
|
SpmdInfo BatchNormInferSpmdStatic(const DistMetaTensor& x,
|
|
const DistMetaTensor& mean,
|
|
const DistMetaTensor& variance,
|
|
const DistMetaTensor& scale,
|
|
const DistMetaTensor& bias) {
|
|
return BatchNormInferSpmd(x, mean, variance, scale, bias);
|
|
}
|
|
SpmdInfo BatchNormGradInferSpmd(const DistMetaTensor& x,
|
|
const DistMetaTensor& scale,
|
|
const DistMetaTensor& bias,
|
|
const DistMetaTensor& mean_out,
|
|
const DistMetaTensor& variance_out,
|
|
const DistMetaTensor& saved_mean,
|
|
const DistMetaTensor& saved_variance,
|
|
const DistMetaTensor& reserve_space,
|
|
const DistMetaTensor& out_grad,
|
|
const float momentum,
|
|
const float epsilon,
|
|
const std::string& data_format,
|
|
const bool is_test,
|
|
const bool use_global_stats,
|
|
const bool trainable_statistics) {
|
|
auto x_shape = vectorize(x.dims());
|
|
auto scale_shape = vectorize(scale.dims());
|
|
auto bias_shape = vectorize(bias.dims());
|
|
auto mean_out_shape = vectorize(mean_out.dims());
|
|
auto variance_out_shape = vectorize(variance_out.dims());
|
|
auto saved_mean_shape = vectorize(saved_mean.dims());
|
|
auto saved_variance_shape = vectorize(saved_variance.dims());
|
|
auto reserve_space_shape = vectorize(reserve_space.dims());
|
|
auto out_grad_shape = vectorize(out_grad.dims());
|
|
int x_ndim = static_cast<int>(x_shape.size());
|
|
int scale_ndim = static_cast<int>(scale_shape.size());
|
|
int bias_ndim = static_cast<int>(bias_shape.size());
|
|
int mean_out_ndim = static_cast<int>(mean_out_shape.size());
|
|
int variance_out_ndim = static_cast<int>(variance_out_shape.size());
|
|
int saved_mean_ndim = static_cast<int>(saved_mean_shape.size());
|
|
int saved_variance_ndim = static_cast<int>(saved_variance_shape.size());
|
|
int reserve_space_ndim = static_cast<int>(reserve_space_shape.size());
|
|
int out_grad_ndim = static_cast<int>(out_grad_shape.size());
|
|
TensorDistAttr x_dist_attr_src = x.dist_attr();
|
|
std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
|
|
TensorDistAttr scale_dist_attr_src = scale.dist_attr();
|
|
TensorDistAttr bias_dist_attr_src = bias.dist_attr();
|
|
TensorDistAttr mean_out_dist_attr_src = mean_out.dist_attr();
|
|
TensorDistAttr variance_out_dist_attr_src = variance_out.dist_attr();
|
|
TensorDistAttr saved_mean_dist_attr_src = saved_mean.dist_attr();
|
|
TensorDistAttr saved_variance_dist_attr_src = saved_variance.dist_attr();
|
|
TensorDistAttr reserve_space_dist_attr_src = reserve_space.dist_attr();
|
|
TensorDistAttr out_grad_dist_attr_src = out_grad.dist_attr();
|
|
PADDLE_ENFORCE_GE(
|
|
x_ndim,
|
|
2,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of x in batch_norm should be greater than 1, but got [%d].",
|
|
x_ndim));
|
|
PADDLE_ENFORCE_LE(
|
|
x_ndim,
|
|
5,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of x in batch_norm should be less than 6, but got [%d].",
|
|
x_ndim));
|
|
PADDLE_ENFORCE_EQ(out_grad_ndim,
|
|
x_ndim,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of out_grad in batch_norm should be equal "
|
|
"with x, but got out_grad:[%d] and x:[%d].",
|
|
out_grad_ndim,
|
|
x_ndim));
|
|
PADDLE_ENFORCE_EQ(
|
|
mean_out_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of mean_out in batch_norm should be 1, but got [%d].",
|
|
mean_out_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
variance_out_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of variance_out in batch_norm should be 1, but got [%d].",
|
|
variance_out_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
scale_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of scale in batch_norm should be 1, but got [%d].",
|
|
scale_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
bias_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of bias in batch_norm should be 1, but got [%d].",
|
|
bias_ndim));
|
|
PADDLE_ENFORCE_EQ(
|
|
saved_mean_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of saved_mean in batch_norm should be 1, but got [%d].",
|
|
saved_mean_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
saved_variance_ndim,
|
|
1,
|
|
common::errors::InvalidArgument(
|
|
"The ndim of saved_variance in batch_norm should be 1, but got [%d].",
|
|
saved_variance_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
reserve_space_ndim,
|
|
1,
|
|
common::errors::InvalidArgument("The ndim of reserve_space_ndim in "
|
|
"batch_norm should be 1, but got [%d].",
|
|
reserve_space_ndim));
|
|
|
|
std::string alphabet = "ijklmnopqrstuvwxyz";
|
|
// get input notation
|
|
// The mean and variance was flatten at C axis
|
|
std::string x_axes(x_ndim, '1');
|
|
std::string out_grad_axes(out_grad_ndim, '1');
|
|
|
|
for (int i = 0; i < x_ndim; ++i) {
|
|
x_axes[i] = alphabet[i];
|
|
out_grad_axes[i] = alphabet[i];
|
|
}
|
|
int c_index = data_format[1] == 'C' ? 1 : x_ndim - 1;
|
|
std::string mean_out_axes(1, x_axes[c_index]);
|
|
std::string variance_out_axes(1, x_axes[c_index]);
|
|
std::string scale_axes(1, x_axes[c_index]);
|
|
std::string bias_axes(1, x_axes[c_index]);
|
|
std::string saved_mean_axes(1, x_axes[c_index]);
|
|
std::string saved_variance_axes(1, x_axes[c_index]);
|
|
std::string reserve_space_axes(1, x_axes[c_index]);
|
|
|
|
auto c_dim =
|
|
x_dims_mapping[c_index]; // Only C axis can be sharded. ndim Type:
|
|
// type: "NC"、"NCL"、"NLC"、"NCHW"、"NHWC"" and
|
|
// "NCDHW"
|
|
|
|
for (int i = 0; i < x_ndim; ++i) {
|
|
x_dims_mapping[i] = -1;
|
|
}
|
|
x_dims_mapping[c_index] = c_dim;
|
|
|
|
std::unordered_map<std::string, int64_t> axis_to_dim_map =
|
|
ShardingMergeForTensors({{x_axes, x_dims_mapping}});
|
|
// infer output spmdinfo
|
|
TensorDistAttr x_grad_dist_attr =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
x_grad_dist_attr.set_dims_mapping(x_dims_mapping);
|
|
TensorDistAttr scale_grad_dist_attr =
|
|
CopyTensorDistAttrForOutput(scale.dist_attr());
|
|
scale_grad_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
|
|
TensorDistAttr bias_grad_dist_attr =
|
|
CopyTensorDistAttrForOutput(bias.dist_attr());
|
|
bias_grad_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
|
|
// infer input spmdinfo
|
|
TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
x_dist_attr_dst.set_dims_mapping(x_dims_mapping);
|
|
TensorDistAttr mean_out_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
mean_out_dist_attr_dst.set_dims_mapping(
|
|
GetDimsMappingForAxes(mean_out_axes, axis_to_dim_map));
|
|
TensorDistAttr variance_out_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
variance_out_dist_attr_dst.set_dims_mapping(
|
|
GetDimsMappingForAxes(variance_out_axes, axis_to_dim_map));
|
|
TensorDistAttr scale_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
scale_dist_attr_dst.set_dims_mapping(std::vector<int64_t>{-1});
|
|
TensorDistAttr bias_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
bias_dist_attr_dst.set_dims_mapping(std::vector<int64_t>{-1});
|
|
TensorDistAttr saved_mean_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
saved_mean_dist_attr_dst.set_dims_mapping(
|
|
GetDimsMappingForAxes(saved_mean_axes, axis_to_dim_map));
|
|
TensorDistAttr saved_variance_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
saved_variance_dist_attr_dst.set_dims_mapping(
|
|
GetDimsMappingForAxes(saved_variance_axes, axis_to_dim_map));
|
|
TensorDistAttr reserve_space_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
reserve_space_dist_attr_dst.set_dims_mapping(std::vector<int64_t>{-1});
|
|
TensorDistAttr out_grad_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
out_grad_dist_attr_dst.set_dims_mapping(
|
|
GetDimsMappingForAxes(out_grad_axes, axis_to_dim_map));
|
|
|
|
// partial grad dim
|
|
std::vector<int64_t> partial_on_dims;
|
|
for (int i = 0; i < x_ndim; ++i) {
|
|
auto mapping = x_dims_mapping[i];
|
|
if (mapping != -1) {
|
|
partial_on_dims.push_back(mapping);
|
|
}
|
|
}
|
|
scale_grad_dist_attr.set_partial_status(partial_on_dims);
|
|
bias_grad_dist_attr.set_partial_status(partial_on_dims);
|
|
|
|
LOG_SPMD_INPUT(x);
|
|
LOG_SPMD_INPUT(scale);
|
|
LOG_SPMD_INPUT(bias);
|
|
LOG_SPMD_INPUT(mean_out);
|
|
LOG_SPMD_INPUT(variance_out);
|
|
LOG_SPMD_INPUT(saved_mean);
|
|
LOG_SPMD_INPUT(saved_variance);
|
|
LOG_SPMD_INPUT(reserve_space);
|
|
LOG_SPMD_INPUT(out_grad);
|
|
LOG_SPMD_OUTPUT(x_grad_dist_attr);
|
|
LOG_SPMD_OUTPUT(scale_grad_dist_attr);
|
|
LOG_SPMD_OUTPUT(bias_grad_dist_attr);
|
|
|
|
return {{x_dist_attr_dst,
|
|
scale_dist_attr_dst,
|
|
bias_dist_attr_dst,
|
|
mean_out_dist_attr_dst,
|
|
variance_out_dist_attr_dst,
|
|
saved_mean_dist_attr_dst,
|
|
saved_variance_dist_attr_dst,
|
|
reserve_space_dist_attr_dst,
|
|
out_grad_dist_attr_dst},
|
|
{x_grad_dist_attr, scale_grad_dist_attr, bias_grad_dist_attr}};
|
|
}
|
|
|
|
} // namespace phi::distributed
|