[spmd_rules] Add SPMD rules for ArgSort operator #72388
Conversation
|
你的PR提交成功，感谢你对开源项目的贡献! |
Yeenyeong
force-pushed
the
dev-argsort-spmd-rule
branch
3 times, most recently
from
561556f to
e14325b
Compare
| x_tensor_dist_attr = dist_attr[0] | ||
| y_tensor_dist_attr = dist_attr[1] | ||
|
|
||
| self.assertEqual(len(dist_attr), 2) |
There was a problem hiding this comment.
this line should be before L77, first determine the length, then retrieve the element
| x_tensor_dist_attr = dist_attr[0] | ||
| y_tensor_dist_attr = dist_attr[1] | ||
|
|
||
| self.assertEqual(len(dist_attr), 2) |
| x_tensor_dist_attr = dist_attr[0] | ||
| y_tensor_dist_attr = dist_attr[1] | ||
|
|
||
| self.assertEqual(len(dist_attr), 2) |
| SpmdInfo ArgSortInferSpmdBase(const DistMetaTensor& x, | ||
| int axis, | ||
| bool descending, | ||
| bool stable); | ||
|
|
||
| SpmdInfo ArgSortGradInferSpmdBase(const DistMetaTensor& indices, | ||
| const DistMetaTensor& x, | ||
| const DistMetaTensor& out_grad, | ||
| int axis, | ||
| bool descending, | ||
| bool stable); |
There was a problem hiding this comment.
What is the reason for ArgSortInferSpmdBase and ArgSortGradInferSpmdBase? for other operators reuse these codes?
Yeenyeong
force-pushed
the
dev-argsort-spmd-rule
branch
2 times, most recently
from
f62d6e7 to
51a9572
Compare
|
@jeff41404 Code has been modified according issues you highlighted. |
|
to generate the correct code, |
Yeenyeong
force-pushed
the
dev-argsort-spmd-rule
branch
from
51a9572 to
62735f0
Compare
| @@ -0,0 +1,37 @@ | |||
| /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. | |||
| @@ -0,0 +1,184 @@ | |||
| # Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. | |||
Yeenyeong
force-pushed
the
dev-argsort-spmd-rule
branch
from
62735f0 to
b4026b7
Compare
Yeenyeong
force-pushed
the
dev-argsort-spmd-rule
branch
from
4be10d4 to
0a50085
Compare
|
/re-run approval |
PR Category
Auto Parallel
PR Types
New features
Description
pcard-90510
The version upon committing incorporates SPMD rules for the ArgSort operator, covering both the forward and backward passes.
The SPMD rules for both the forward and backward passes simply replicate the sorting axis across one worker mesh dimension. For example, dimension mapping [0, 1, -1], given sorting axis of 1, will be transformed into dimension mapping [0, -1, -1].
The rules that involve splitting the sorting axis are deeply coupled with the reduce communication mechanism. Specifically, merging sorted sequences during reduction is not currently supported and remains a topic for future work.