PR Category

Auto Parallel

PR Types

New features

Description

Pcard-67164

Add corresponding dist_op of the MoE api (#63904) in PIR dist dialect, used for MoE model as following:

• local_tensors_from_dtensor: get the tensor list of a dist_tensor on sub-mesh, e.g. for a dist tensor with mesh=[0,1], placements=[Shard(0)], get its sub-mesh list [DistTensor(mesh=[0],placements=[Replicate()]), DistTensor(mesh=[1], placements=[Replicate()])]
• dtensor_from_local_tensors: the opposite operation of local_tensors_from_dtensor, get the global-mesh dist tensor from sub-mesh dist tensors, e.g. for a sub-mesh list [DistTensor(mesh=[0],placements=[Replicate()]), DistTensor(mesh=[1], placements=[Replicate()])], get the global mesh dist tensor with mesh=[0,1], placements=[Shard(0)].