Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[auto parallel] complete elementwise ops spmd rules of LLaMa2 for eager semi auto parallel #58474

Merged
merged 5 commits into from
Oct 31, 2023
Merged

[auto parallel] complete elementwise ops spmd rules of LLaMa2 for eager semi auto parallel #58474

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9857325
【pir】Modify comment of #58401 (#58420)
xiaoguoguo626807 Oct 27, 2023
8bf445a
add elementwise spmd rules for eager semi auto parallel
cxxly Oct 30, 2023
b2196ae
support bool dtype for s_to_r_reshared
cxxly Oct 30, 2023
809f4a5
increase test_semi_auto_parallel_basic timeout for pass ci
cxxly Oct 30, 2023
1d8a01e
prompt rtol tolerance from 1e-5 to 1e-6
cxxly Oct 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions paddle/phi/api/yaml/backward.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -432,6 +432,7 @@
infer_meta :
func : UnchangedInferMeta
param : [x]
spmd_rule : ElementwiseUnaryGradInferSpmd
kernel :
func : cos_grad
backward : cos_double_grad
Expand Down Expand Up @@ -708,6 +709,7 @@
infer_meta :
func : UnchangedInferMeta
param : [out]
spmd_rule : ElementwiseUnaryGradInferSpmd
kernel :
func : exp_grad
inplace : (out_grad -> x_grad)
Expand Down Expand Up @@ -1907,6 +1909,7 @@
infer_meta :
func : UnchangedInferMeta
param : [out]
spmd_rule : ElementwiseUnaryGradInferSpmd
kernel :
func : rsqrt_grad
backward : rsqrt_double_grad
Expand Down Expand Up @@ -2062,6 +2065,7 @@
infer_meta :
func : UnchangedInferMeta
param : [x]
spmd_rule : ElementwiseUnaryGradInferSpmd
kernel :
func : silu_grad
backward : silu_double_grad
Expand All @@ -2088,6 +2092,7 @@
infer_meta :
func : UnchangedInferMeta
param : [x]
spmd_rule : ElementwiseUnaryGradInferSpmd
kernel :
func : sin_grad
backward : sin_double_grad
Expand Down
3 changes: 3 additions & 0 deletions paddle/phi/api/yaml/legacy_backward.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,7 @@
infer_meta :
func : GeneralBinaryGradInferMeta
param : [x, y]
spmd_rule : ElementwiseBinaryGradInferSpmd
kernel :
func : divide_grad
composite : divide_grad(x, y, out, out_grad, axis, x_grad, y_grad)
Expand Down Expand Up @@ -226,6 +227,7 @@
infer_meta :
func : GeneralBinaryGradInferMeta
param: [x, y]
spmd_rule : ElementwiseBinaryGradInferSpmd
composite : elementwise_pow_grad(x, y, out_grad, x_grad, y_grad)
kernel :
func : elementwise_pow_grad
Expand Down Expand Up @@ -440,6 +442,7 @@
infer_meta :
func : GeneralBinaryGradInferMeta
param : [x, y]
spmd_rule : ElementwiseBinaryGradInferSpmd
kernel :
func : multiply_grad
composite: multiply_grad(x, y, out_grad, axis, x_grad, y_grad)
Expand Down
5 changes: 5 additions & 0 deletions paddle/phi/api/yaml/legacy_ops.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,7 @@
output : Tensor(out)
infer_meta :
func : ElementwiseInferMeta
spmd_rule : ElementwiseBinaryInferSpmd
kernel :
func : divide
inplace: (x -> out)
Expand Down Expand Up @@ -348,6 +349,7 @@
output : Tensor(out)
infer_meta :
func : ElementwiseInferMeta
spmd_rule: ElementwiseBinaryInferSpmd
kernel :
func : elementwise_pow
backward : elementwise_pow_grad
Expand Down Expand Up @@ -415,6 +417,7 @@
output : Tensor(out)
infer_meta :
func : CompareInferMeta
spmd_rule: ElementwiseBinaryInferSpmd
kernel :
func : equal
inplace: (x -> out)
Expand Down Expand Up @@ -783,6 +786,7 @@
output : Tensor
infer_meta :
func : ElementwiseInferMeta
spmd_rule : ElementwiseBinaryInferSpmd
kernel :
func : multiply {dense, dense -> dense},
multiply_sr {selected_rows, dense -> selected_rows}
Expand All @@ -803,6 +807,7 @@
output : Tensor(out)
infer_meta :
func : CompareInferMeta
spmd_rule : ElementwiseBinaryInferSpmd
kernel :
func : not_equal
inplace: (x -> out)
Expand Down
7 changes: 7 additions & 0 deletions paddle/phi/api/yaml/ops.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -327,6 +327,7 @@
output : Tensor(out)
infer_meta :
func : ElementwiseInferMeta
spmd_rule : ElementwiseBinaryInferSpmd
kernel :
func : bitwise_and
backend : x
Expand All @@ -337,6 +338,7 @@
output : Tensor(out)
infer_meta :
func : UnchangedInferMeta
spmd_rule : ElementwiseUnaryInferSpmd
kernel :
func : bitwise_not
backend : x
Expand Down Expand Up @@ -550,6 +552,7 @@
output : Tensor(out)
infer_meta :
func : UnchangedInferMeta
spmd_rule : ElementwiseUnaryInferSpmd
kernel :
func : cos
inplace: (x -> out)
Expand Down Expand Up @@ -819,6 +822,7 @@
output : Tensor(out)
infer_meta :
func : UnchangedInferMeta
spmd_rule : ElementwiseUnaryInferSpmd
kernel :
func : exp
inplace : (x -> out)
Expand Down Expand Up @@ -2180,6 +2184,7 @@
output : Tensor(out)
infer_meta :
func : UnchangedInferMeta
spmd_rule : ElementwiseUnaryInferSpmd
kernel :
func : rsqrt
inplace : (x -> out)
Expand Down Expand Up @@ -2360,6 +2365,7 @@
output : Tensor
infer_meta :
func : UnchangedInferMeta
spmd_rule : ElementwiseUnaryInferSpmd
kernel :
func : silu
backward : silu_grad
Expand All @@ -2369,6 +2375,7 @@
output : Tensor(out)
infer_meta :
func : UnchangedInferMeta
spmd_rule : ElementwiseUnaryInferSpmd
kernel :
func : sin
inplace: (x -> out)
Expand Down
4 changes: 2 additions & 2 deletions paddle/phi/core/distributed/auto_parallel/reshard_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,14 +76,14 @@ CommContext* CreateOrGetCommContext(const DeviceContext& dev_ctx,
do { \
if (phi::CPUContext::classof(dev_ctx)) { \
VLOG(4) << "Call `" << #fn_name << "` in Resharding on GPU."; \
PD_VISIT_FLOATING_AND_INTEGRAL_TYPES( \
PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES( \
dtype, #fn_name, ([&] { \
fn_name<data_t>(static_cast<const CPUContext&>(*dev_ctx), \
__VA_ARGS__); \
})); \
} else if (phi::GPUContext::classof(dev_ctx)) { \
VLOG(4) << "Call `" << #fn_name << "` in Resharding on CPU."; \
PD_VISIT_FLOATING_AND_INTEGRAL_TYPES( \
PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES( \
dtype, #fn_name, ([&] { \
fn_name<data_t>(static_cast<const GPUContext&>(*dev_ctx), \
__VA_ARGS__); \
Expand Down
27 changes: 27 additions & 0 deletions paddle/phi/core/visit_type.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,33 @@ namespace phi {
} \
}()

///////// BOOL and Floating and Integral Dispatch Marco ///////////

#define PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES(TYPE, NAME, ...) \
[&] { \
const auto& __dtype__ = TYPE; \
switch (__dtype__) { \
PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::BOOL, bool, __VA_ARGS__) \
PD_PRIVATE_CASE_TYPE( \
NAME, ::paddle::DataType::FLOAT32, float, __VA_ARGS__) \
PD_PRIVATE_CASE_TYPE( \
NAME, ::paddle::DataType::FLOAT64, double, __VA_ARGS__) \
PD_PRIVATE_CASE_TYPE(NAME, ::paddle::DataType::INT32, int, __VA_ARGS__) \
PD_PRIVATE_CASE_TYPE( \
NAME, ::paddle::DataType::INT64, int64_t, __VA_ARGS__) \
PD_PRIVATE_CASE_TYPE( \
NAME, ::paddle::DataType::INT8, int8_t, __VA_ARGS__) \
PD_PRIVATE_CASE_TYPE( \
NAME, ::paddle::DataType::UINT8, uint8_t, __VA_ARGS__) \
PD_PRIVATE_CASE_TYPE( \
NAME, ::paddle::DataType::INT16, int16_t, __VA_ARGS__) \
default: \
PD_THROW("function " #NAME " is not implemented for data type `", \
__dtype__, \
"`"); \
} \
}()

///////// Floating and Complex Dispatch Marco ///////////

#define PD_VISIT_FLOATING_AND_COMPLEX_TYPES(TYPE, NAME, ...) \
Expand Down
19 changes: 19 additions & 0 deletions paddle/phi/infermeta/spmd_rules/elementwise.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,6 +314,13 @@ SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
return {{out_grad.dist_attr(), out_grad.dist_attr()}, {out_grad.dist_attr()}};
}

SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
const DistMetaTensor& out,
const DistMetaTensor& out_grad) {
return {{out_grad.dist_attr(), out_grad.dist_attr(), out_grad.dist_attr()},
{out_grad.dist_attr()}};
}

SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
const DistMetaTensor& y,
const DistMetaTensor& out_grad,
Expand Down Expand Up @@ -381,5 +388,17 @@ SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
{x_grad_dist_attr, y_grad_dist_attr}};
}

SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
const DistMetaTensor& y,
const DistMetaTensor& out,
const DistMetaTensor& out_grad,
int64_t axis) {
// The out's dist_attr is the same with out_grad's dist_attr, reuse
// ElementwiseBinaryGradInferSpmd(x, y, out_grad, axis) to infer dist_attrs of
// {{x, y, out_grad}, {x_grad, y_grad}}, then insert out's dist_attr into it.
SpmdInfo info = ElementwiseBinaryGradInferSpmd(x, y, out_grad, axis);
info.first.emplace(info.first.begin() + 2, out_grad.dist_attr());
return info;
}
} // namespace distributed
} // namespace phi
10 changes: 10 additions & 0 deletions paddle/phi/infermeta/spmd_rules/elementwise.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,10 @@ SpmdInfo ElementwiseUnaryInferSpmdReverse(const DistMetaTensor& x,
SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
const DistMetaTensor& out_grad);

SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

缺少单规则模块的单测，可以参考 test/auto_parallel/spmd_rules/test_elementwise_rule.py

Copy link
Contributor Author

@cxxly cxxly Oct 31, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

新提一个PR补充

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

重载是为了适配 api 么？是不是可以加些注释说明一下

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

适配不同API签名，比如一元逐元素运算，部分算子反向只依赖输入，部分依赖输入输出，sin_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad), silu_grad(const Tensor& x, const Tensor& out, const Tensor& out_grad, Tensor* x_grad)

const DistMetaTensor& out,
const DistMetaTensor& out_grad);

SpmdInfo ElementwiseBinaryInferSpmd(const DistMetaTensor& x,
const DistMetaTensor& y);

Expand All @@ -42,5 +46,11 @@ SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
const DistMetaTensor& out_grad,
int64_t axis = -1);

SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同上

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

新提交一个PR补充

const DistMetaTensor& y,
const DistMetaTensor& out,
const DistMetaTensor& out_grad,
int64_t axis = -1);

} // namespace distributed
} // namespace phi
2 changes: 2 additions & 0 deletions paddle/phi/kernels/cpu/elementwise_divide_grad_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ PD_REGISTER_KERNEL(divide_grad,
int16_t,
int,
int64_t,
bool,
phi::dtype::complex<float>,
phi::dtype::complex<double>) {}

Expand All @@ -61,5 +62,6 @@ PD_REGISTER_KERNEL(divide_double_grad,
double,
int,
int64_t,
bool,
phi::dtype::complex<float>,
phi::dtype::complex<double>) {}
1 change: 1 addition & 0 deletions paddle/phi/kernels/cpu/elementwise_divide_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,5 +64,6 @@ PD_REGISTER_KERNEL(divide,
int16_t,
int,
int64_t,
bool,
complex64,
complex128) {}
2 changes: 2 additions & 0 deletions paddle/phi/kernels/gpu/elementwise_grad_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -329,6 +329,7 @@ PD_REGISTER_KERNEL(divide_grad,
int16_t,
int,
int64_t,
bool,
phi::dtype::complex<float>,
phi::dtype::complex<double>) {}

Expand All @@ -342,6 +343,7 @@ PD_REGISTER_KERNEL(divide_double_grad,
double,
int,
int64_t,
bool,
phi::dtype::complex<float>,
phi::dtype::complex<double>) {}

Expand Down
1 change: 1 addition & 0 deletions paddle/phi/kernels/kps/elementwise_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,7 @@ PD_REGISTER_KERNEL(divide,
int16_t,
int,
int64_t,
bool,
float16,
bfloat16,
complex64,
Expand Down
1 change: 1 addition & 0 deletions paddle/phi/kernels/legacy/cpu/elementwise_divide_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,5 +62,6 @@ PD_REGISTER_KERNEL(divide_raw,
double,
int,
int64_t,
bool,
complex64,
complex128) {}
1 change: 1 addition & 0 deletions paddle/phi/kernels/legacy/kps/elementwise_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@ PD_REGISTER_KERNEL(divide_raw,
int16_t,
int,
int64_t,
bool,
float16,
bfloat16,
complex64,
Expand Down
2 changes: 1 addition & 1 deletion test/auto_parallel/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
py_test_modules(test_semi_auto_parallel_basic MODULES
test_semi_auto_parallel_basic)
set_tests_properties(test_semi_auto_parallel_basic
PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 200)
py_test_modules(test_semi_auto_parallel_single_strategy MODULES
test_semi_auto_parallel_single_strategy)
set_tests_properties(test_semi_auto_parallel_single_strategy
Expand Down
Loading