PR Category

Auto Parallel

PR Types

Bug fixes

Description

In a distributed scenario, the output shape of the nonzero method is incorrect (the first dimension is -1), which is abnormal. This issue arises because the shape of nonzero is generated after the computation, but the framework does not place it into ops_infer_shape_in_runtime. The following code can reproduce the corresponding problem.

import paddle
import paddle.distributed as dist
x = paddle.to_tensor([1, 0, 1, 1])
mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
x = dist.shard_tensor(x,mesh,[dist.Replicate()])
out = paddle.nonzero(x)
print(out.shape)

Added a unit test for MoE layers with shared experts, and that multiple experts can be saved on a device