PR Category

Auto Parallel

PR Types

Others

Description

• OperationDistAttr支持输入输出为Tensor列表的情况。

  • 目前的OperationDistAttr是由两个分别表示输入输出的 vector组成。当算子的输入或输出为list的时候，就会出现问题。以SplitWithNumOp算子为例，它只有一个类型为VectorType的输出，代表该输出的类型是Tensor数组；那么此时，算子的OperationDistAttr是无法表示输出的分布式属性的。
  • 本pr重构升级OperationDistAttr， 将它的两个分别表示输入输出的 vector，调整为 vectorpir::Attribute, 里面的每个元素是TensorDistAttr或者ArrayAttribute。由上层调用者去处理。
  • 假设算子A包含了两个输出，且都是Tensor数组，这两个输出数组的长度分别为4和5，那么此时，他的OperationDistAttr的 result_dist_attr的长度应该为2， 下标0是一个ArrayAttribute, 代表第一个输出数组中每个Tensor的分布式属性； 下标1也是一个ArrayAttribute， 代表第二个输出数组中的每个Tensor的分布式属性。

• 修改call_vjp的pybind接口，强约束张量的前反向的类型完全一致。

  • 目前的vjp接口中，反向算子在build时，输出反向张量的分布式属性来自于inferspmd推断得到的结果。该结果和前向张量的分布式属性不一定一致，这会导致在后面算子的组网失败。
  • 本pr中，inferspmd推断得到的分布式属性仅存储在op_dsit_attr中，输出反向张量的分布式属性和前向严格一致。

• 在partition pass中，在输出的分布式属性和op_dist_attr中的分布式属性不一致时，插入reshard op.

• 修改partition pass, 消除冗余的reshard。（比如两个相邻的reshard功能互逆时，可以同时移除）

Other

Pcard-67164