PR types

New features

PR changes

Others

Description

Pcard-73145

支持Shard到Replicate的状态转换，要求：

输入输出的process_mesh为一维；
输出输出不跨mesh；（process_mesh无变化）
输入的shard状态为均匀切分，Tensor的切分维度能被对应组的进程数整除。

以4卡为例，输入输出都是一维process_mesh，[0, 1, 2, 3]，输出为二维Replicate状态，out_dims_mapping = [-1, -1]，输入为二维Shard状态，in_tensor_shape = [4, 8]

• 用process_mesh的0维切分输入的0维，in_dims_mapping = [0, -1]，每个进程上最终有形状为[16, 8]的物理tensor。
• 用process_mesh的0维切分输入的1维，in_dims_mapping = [-1, 0]，每个进程上最终有形状为[4, 32]的物理tensor。（待实现）

TODO：

• 静态检查，限制tensor的切分维度被对应组的进程数 整除
• 支持切分输入的非0维，需要在all_gather后，实现split和concat