s_to_r reshard support unbalanced split

[jeff41404]

PR Category

Auto Parallel

PR Types

New features

Description

pcard-86471
s_to_r reshard should support unbalanced split or will error in some model (Dit etc)

import paddle
import paddle.distributed as dist
process_mesh = dist.ProcessMesh([0, 1], dim_names=['dp'])

class ParallelLabelEmbedder(nn.Layer):
    def __init__(self, num_classes, hidden_size, dropout_prob):
        ......
        self.embedding_table = nn.Embedding(embedding_dim, hidden_size) # parameter shape=(1001, 1024)
        self.embedding_table.weight = dist.shard_tensor(self.embedding_table.weight, process_mesh, [dist.Replicate()])
        # Although we DO NOT shard the parameters in the forward and backward processes, 
        # But when using sharding parallel(eg. sharding stage 1), the optimizer will shard the parameters
        # at the 0th dimension, (1001/2) resulting unbalanced splitting. 
        # and using all_gather to update parameters, an error will occur

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[JZ-LIANG]

LGTM