[Semi-Auto] add split spmd rule #55397
Conversation
|
你的PR提交成功，感谢你对开源项目的贡献! |
| @@ -145,6 +146,9 @@ REGISTER_SPMD_RULE(lookup_table_v2, EmbeddingSPMDRule); | |||
| REGISTER_SPMD_RULE(softmax, SoftmaxSPMDRule); | |||
| REGISTER_SPMD_RULE(log_softmax, SoftmaxSPMDRule); | |||
|
|
|||
| // split rule | |||
| REGISTER_SPMD_RULE(split, SplitSPMDRule); | |||
There was a problem hiding this comment.
name string for regstration should be op names which are defined in phi yaml.
for split them are : split_with_num & split.
and you should take care about the different in attribute among these two ops
|
|
||
| // step1: Build Einsum Notation | ||
| int64_t ndim = input_specs[0].shape().size(); | ||
| Attribute section_attr = GetAttr("num_or_sections", attrs); |
There was a problem hiding this comment.
follow definition in phi yaml：
- op : split_with_num
args : (Tensor x, int num, Scalar(int) axis) - op : split
args : (Tensor x, IntArray sections, Scalar(int) axis)
there is no num_or_sections
| input_axes_vec.emplace_back(input_axes); | ||
|
|
||
| // get einsum notation for output | ||
| std::string output_axes(input_axes); |
There was a problem hiding this comment.
we should support input tensor sharded on split_axis.
the main logic for split spmd forward infer should be:
- for axes other than the split_axis, treat as broadcast axes, copy dims from input to outputs.
- for split_axis:
the infered dim_mapping for all outputs are replicated
the infered dst_dim_mapping for all outputs is also replicated
There was a problem hiding this comment.
use a special char for split axis, like: "k"
| """ | ||
|
|
||
| def setUp(self): | ||
| self.rule = get_spmd_rule("split") |
There was a problem hiding this comment.
for both split and split_with_num
|
Sorry to inform you that fe38b69's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually. |
* add split spmd rule * add pytest in cmake file * small fix
* add split spmd rule * add pytest in cmake file * small fix
* add split spmd rule * add pytest in cmake file * small fix
PR types
New features
PR changes
Others
Description
Pcard-70448
Add split spmd rule for auto parallel. It infers the output dims mapping as following: the splitted axis cannot be sharded, set the dims mapping of splitted axis in input and outputs to -1. For other axes in output, set their dims mapping equal to input's one.