[Semi-Auto] LayerNorm Parallel Rule #55130
Conversation
|
你的PR提交成功，感谢你对开源项目的贡献! |
| std::pair<std::vector<TensorDistAttr>, std::vector<TensorDistAttr>> | ||
| LayerNormSPMDRule::InferForward(const std::vector<DistTensorSpec>& input_specs, | ||
| const paddle::framework::AttributeMap& attrs) { | ||
| // step0: verify input args based on matmul logic |
| } | ||
| std::string out_axes = x_axes; | ||
|
|
||
| VLOG(4) << "LayerNormSPMDRule build Einsum notation (x,scale,bias->out): [" |
There was a problem hiding this comment.
should include all layer_norm outputs (x,scale,bias->out, mean, var)
| input_dist_attrs.emplace_back(ReplicatedOnMesh(input_specs[2].dist_attr())); | ||
|
|
||
| // Step2.4. handle input and out tensor partial | ||
| std::vector<int64_t> partial_on_dims; |
There was a problem hiding this comment.
LayerNorm activation output would not be partial
| << "]; out dims_mapping: [" << str_join(out_dims_mapping) | ||
| << "], partial_on_dims: [" << str_join(partial_on_dims) << "]"; | ||
|
|
||
| return {input_dist_attrs, {output_dist_attr_dst}}; |
There was a problem hiding this comment.
should infer the distattr of output （variance and mean）
| namespace distributed { | ||
| namespace auto_parallel { | ||
|
|
||
| TensorDistAttr GetInferedDistAttr( |
|
|
||
| SPMDRuleBase* layer_norm_rule = SPMDRuleMap::Instance().Get("layer_norm"); | ||
|
|
||
| // ijk[1, -1, -1],k[-1],k[-1] --> ijk[1, -1, -1] partial[1] |
There was a problem hiding this comment.
I will fix it in the next pr.
| VLOG(4) << "test1 done."; | ||
|
|
||
| // ijk[1, 0, -1],k[0],k[0] --> ijk[1, 0, -1] | ||
| x_dist_tensor_spec.set_dims_mapping({1, 0, -1}); |
There was a problem hiding this comment.
GOOD test case !
there are two kind of error in the case:
- multiple batch axes (i,j) are sharded in input activation, which is not supported by now.
- same mesh dimension ( 0 ) is sharding two different tensor axes (j & k).
it inspires me that we are missing one important precondition checking in InferForward (when support sharding on bias in future):
when [ijkl,y(kl),y(kl)->ijkl,x(ij),x(ij) (x,scale,bias->out,mean,variance, begin_norm_axis=2, x=ij, y=kl)]
thought "y" and "kl" are represented by different char, but they are same tensor axis and should be sharded by same mesh dim.
| attrs);); | ||
| VLOG(4) << "test2 done."; | ||
|
|
||
| // ijk[0, -1, -1],z[-1],z[1] --> ijk[0, 1, -1, -1], z=jk |
There was a problem hiding this comment.
ijk[0, -1, -1], k[-1], k[1] --> ijk[0, -1, -1], x[0], x[0], x=ij
There was a problem hiding this comment.
I will fix it in the next pr.
* add layernorm spmd rule * add ut * follow comments
PR types
New features
PR changes
OPs
Description
Pcard-70448
add spmd rule for layer_norm