static support mp_layers

[wangxicoding]

PR types

New features

PR changes

APIs

Describe

c_softmax_with_cross_entropy API support static graph.
mp_layers接口支持静态图，mp_layers实现动静统一。

import paddle
import paddle.distributed.fleet as fleet

paddle.enable_static()

class ColumnLinearNet(paddle.nn.Layer):
    def __init__(self, input_size, output_size):
        super(ColumnLinearNet, self).__init__()
        self.parallel_linear = fleet.meta_parallel.ColumnParallelLinear(
            in_features=input_size,
            out_features=output_size,
            weight_attr=None,
            has_bias=True,
            gather_output=True,
            name="test_column_linear")

    def forward(self, x):
        output = self.parallel_linear(x)
        return output

strategy = fleet.DistributedStrategy()
strategy.sharding = True
strategy.sharding_configs = {
    "mp_degree": 2,
    "sharding_degree": 2,
}
fleet.init(is_collective=True, strategy=strategy)

input_size, output_size = 28, 64
model_a = ColumnLinearNet(input_size, output_size)

x = paddle.static.data(name='x', shape=[None, input_size])
y = model_a(x)

python -m paddle.distributed.launch --gpus 0,1,2,3 test.py

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[kuizhiqing]

lgtm for new group part

[JZ-LIANG]

LGTM

[JZ-LIANG]

or we could explain the parallel2id mapping here again?

[JZ-LIANG]

    pre-assign ring ids
        mp: 0
        sharding: 1
        pure-dp: 2
        global: 3
        pp: >= 20
    if one parallelism is not enable: -1
    and only support parallelism hierarchy: mp --> sharding --> pp --> dp      

done

[ForFishes]

LGTM