paddlepaddle--paddle/test/auto_parallel/test_static_sequence_parallel_pass.py

# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import unittest

import paddle

sys.path.append("../legacy_test")

import paddle.nn.functional as F
from paddle import nn, static, utils
from paddle.base import ParamAttr
from paddle.distributed.auto_parallel.static.dist_context import (
    DistributedContext,
)
from paddle.distributed.auto_parallel.static.parallelizer_v2 import Parallelizer
from paddle.distributed.auto_parallel.static.planner_v2 import Planner
from paddle.distributed.auto_parallel.strategy import Strategy
from paddle.distributed.fleet import auto

paddle.enable_static()
BATCH_SIZE = 4
SEQ_LEN = 512
HIDDEN_SIZE = 1024
MESH_0 = auto.ProcessMesh([0, 1, 2, 3], dim_names=["x"])


class MLPLayer(nn.Layer):
    def __init__(
        self,
        hidden_size=1024,
        intermediate_size=4 * 1024,
        dropout_ratio=0.1,
        initializer_range=0.02,
    ):
        super().__init__()
        d_model = hidden_size
        dim_feedforward = intermediate_size
        weight_attr = ParamAttr(
            initializer=paddle.nn.initializer.Normal(
                mean=0.0, std=initializer_range
            )
        )
        bias_attr = False

        self.norm0 = paddle.nn.LayerNorm(d_model, epsilon=1e-5)
        self.norm0.bias.stop_gradient = True
        self.norm1 = paddle.nn.LayerNorm(d_model, epsilon=1e-5)
        self.norm1.bias.stop_gradient = True
        self.linear0 = paddle.nn.Linear(
            d_model, dim_feedforward, weight_attr, bias_attr=bias_attr
        )
        auto.shard_tensor(self.linear0.weight, MESH_0, [None, "x"])
        self.linear1 = paddle.nn.Linear(
            dim_feedforward, d_model, weight_attr, bias_attr=bias_attr
        )
        auto.shard_tensor(self.linear1.weight, MESH_0, ["x", None])
        self.dropout = paddle.nn.Dropout(dropout_ratio, mode="upscale_in_train")

    def forward(self, input):
        # sp region
        auto.shard_tensor(input, MESH_0, ["x", None, None])
        out = self.norm0(input)
        auto.shard_tensor(input, MESH_0, ["x", None, None])
        out = F.gelu(out, approximate=True)

        # tp region
        auto.shard_tensor(out, MESH_0, [None, None, None])
        out = self.linear0(out)
        out = F.gelu(out, approximate=True)
        out = self.linear1(out)
        auto.shard_tensor(out, MESH_0, [None, None, None])

        # sp region
        out = self.dropout(out)
        auto.shard_tensor(out, MESH_0, ["x", None, None])
        out = F.gelu(out, approximate=True)
        out = self.norm1(out)

        return out


class HybridParallelNet(nn.Layer):
    def __init__(
        self,
        hidden_size=1024,
    ):
        super().__init__()
        self.mlp0 = MLPLayer(hidden_size, hidden_size * 4)
        self.mlp1 = MLPLayer(hidden_size, hidden_size * 4)

    def forward(self, input):
        out = self.mlp0(input)
        out = self.mlp1(out)

        return out


def get_hybrid_parallel_model(train_program, start_program):
    with (
        static.program_guard(train_program, start_program),
        utils.unique_name.guard(),
    ):
        batch_size = BATCH_SIZE
        hidden_size = HIDDEN_SIZE
        sequence_len = SEQ_LEN

        input = static.data(
            name="input",
            shape=[batch_size, sequence_len, hidden_size],
            dtype='float32',
        )
        network = HybridParallelNet(hidden_size=HIDDEN_SIZE)

        predict = network(input)
        error_cost = paddle.sum(predict)

    return error_cost, train_program, start_program


def get_dist_prog(rank=2):
    train_program = paddle.static.Program()
    startup_program = paddle.static.Program()

    loss, train_program, startup_program = get_hybrid_parallel_model(
        train_program, startup_program
    )
    opt = paddle.optimizer.AdamW(learning_rate=0.00001)
    strategy = Strategy()
    strategy.auto_mode = "semi"
    strategy.sp_optimization.enable = True
    dist_context = DistributedContext(
        train_program, startup_program, opt, loss, strategy=strategy
    )
    planner = Planner("train", dist_context)
    planner.plan()

    parallelizer = Parallelizer(
        "train",
        planner.completer,
        dist_context,
    )
    parallelizer.parallel(rank=rank)

    return (
        dist_context.dist_main_programs[rank],
        dist_context.dist_startup_programs[rank],
    )


class TestGradSync(unittest.TestCase):
    def test_decoder_dp_sp(self):
        dist_main_prog, dist_startup_prog = get_dist_prog()

        with open("test_static_sequence_parallel.txt", "w+") as f:
            f.write(str(dist_main_prog))

        ops = dist_main_prog.global_block().ops
        sp_ring_id = None
        allgather_count = 0
        reducescatter_count = 0
        allreduce_count = 0

        for op in ops:
            # check sequence parallel allgather
            if op.type == "all_gather":
                assert int(op.attr("nranks")) == 4, (
                    "sequence parallel allgather error with nranks [{}]".format(
                        op.attr("nranks")
                    )
                )
                if sp_ring_id is None:
                    sp_ring_id = int(op.attr("ring_id"))
                else:
                    assert sp_ring_id == int(op.attr("ring_id")), (
                        "sequence parallel allgather error with ring_id [{}]".format(
                            op.attr("ring_id")
                        )
                    )
                allgather_count += 1

            # check sequence parallel reducescatter
            elif op.type == "reduce_scatter":
                assert int(op.attr("nranks")) == 4, (
                    "sequence parallel reducescatter error with nranks [{}]".format(
                        op.attr("nranks")
                    )
                )
                assert sp_ring_id == int(op.attr("ring_id")), (
                    "sequence parallel reducescatter error with ring_id [{}]".format(
                        op.attr("ring_id")
                    )
                )
                reducescatter_count += 1

            # check sequence parallel grad sync
            elif op.type == "c_allreduce_sum":
                assert "layer_norm" in op.output_arg_names[0], (
                    f"sequence parallel reducescatter error grad sync var [{op.output_arg_names[0]}]"
                )
                assert sp_ring_id == int(op.attr("ring_id")), (
                    "sequence parallel reducescatter error with ring_id [{}]".format(
                        op.attr("ring_id")
                    )
                )
                allreduce_count += 1

        assert allgather_count == 4, (
            f"sequence parallel should have 4 allgather, but got [{allgather_count}]"
        )
        assert reducescatter_count == 4, (
            f"sequence parallel should have 4 allgather, but got [{reducescatter_count}]"
        )
        assert allreduce_count == 4, (
            f"sequence parallel should have 4 allgather, but got [{allreduce_count}]"
        )


if __name__ == "__main__":
    unittest.main()