test_auto_parallel_mapper.py

#   Copyright (c) 2021 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.

from __future__ import print_function

import tempfile
import unittest
import os
import json
import collections
import math
import numpy as np

import paddle
import paddle.nn as nn
import paddle.fluid as fluid
import paddle.nn.functional as F
import paddle.tensor as tensor
import paddle.utils as utils
import paddle.static as static
from paddle.fluid import core
from paddle.fluid import layers
from paddle.fluid.framework import _non_static_mode
from paddle.nn.layer.transformer import _convert_param_attr_to_list
from paddle.fluid.initializer import Normal, Constant, NumpyArrayInitializer
from paddle.distributed import fleet

import paddle.distributed.auto_parallel as auto
from paddle.distributed.auto_parallel.completion import Completer
from paddle.distributed.auto_parallel.parallelizer import AutoParallelizer
from paddle.distributed.auto_parallel.dist_context import DistributedContext
from paddle.distributed.auto_parallel.partitioner import Partitioner
from paddle.distributed.auto_parallel.reshard import Resharder
from paddle.distributed.auto_parallel.process_group import get_all_process_groups
from paddle.distributed.auto_parallel.process_group import new_process_group
from paddle.distributed.auto_parallel.cluster import Cluster
from paddle.distributed.auto_parallel.cluster import DeviceType
from paddle.distributed.auto_parallel.cluster import LinkType
from paddle.distributed.auto_parallel.utils import check_distributed_attr_for_program
from paddle.distributed.auto_parallel.utils import print_program_with_dist_attr
from paddle.distributed.auto_parallel.mapper import build_process_graph
from paddle.distributed.auto_parallel.mapper import build_cluster_graph
from paddle.distributed.auto_parallel.mapper import mapping
from paddle.distributed.auto_parallel.mapper import get_dtype_bytes
from paddle.distributed.auto_parallel.mapper import get_comm_volume

if os.getenv("CUDA_VISIBLE_DEVICES") is not None:
    os.environ["CUDA_VISIBLE_DEVICES"] = ""

paddle.enable_static()
_global_parallel_strategy = None
_global_process_mesh = None
_global_num_stages = None

cluster_json = """
{
  "machines": [
    {
      "hostname": "machine0",
      "addr": "0.0.0.1",
      "port": "768",
      "devices": [
        {
          "global_id": 0,
          "local_id": 0,
          "type": "GPU",
          "model": "A100-SXM4-40GB",
          "sp_gflops": 19500,
          "dp_gflops": 9700,
          "memory": 40
        },
        {
          "global_id": 1,
          "local_id": 1,
          "type": "GPU",
          "model": "A100-SXM4-40GB",
          "sp_gflops": 19500,
          "dp_gflops": 9700,
          "memory": 40
        },
        {
          "global_id": 2,
          "local_id": 2,
          "type": "GPU",
          "model": "A100-SXM4-40GB",
          "sp_gflops": 19500,
          "dp_gflops": 9700,
          "memory": 40
        },
        {
          "global_id": 3,
          "local_id": 3,
          "type": "GPU",
          "model": "A100-SXM4-40GB",
          "sp_gflops": 19500,
          "dp_gflops": 9700,
          "memory": 40
        },
        {
          "global_id": 4,
          "local_id": 0,
          "type": "NIC"
        }
      ],
      "links": [
        {
          "source_global_id": 0,
          "target_global_id": 1,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 0,
          "target_global_id": 2,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 0,
          "target_global_id": 3,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 0,
          "target_global_id": 4,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 1,
          "target_global_id": 0,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 1,
          "target_global_id": 2,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 1,
          "target_global_id": 3,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 1,
          "target_global_id": 4,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 2,
          "target_global_id": 0,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 2,
          "target_global_id": 1,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 2,
          "target_global_id": 3,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 2,
          "target_global_id": 4,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 3,
          "target_global_id": 0,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 3,
          "target_global_id": 1,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 3,
          "target_global_id": 2,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 3,
          "target_global_id": 4,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 4,
          "target_global_id": 9,
          "type": "NET",
          "bandwidth": 1
        }
      ]
    },
    {
      "hostname": "machine1",
      "addr": "0.0.0.2",
      "port": "768",
      "devices": [
        {
          "global_id": 5,
          "local_id": 0,
          "type": "GPU",
          "model": "Tesla V100-SXM2-32GB",
          "sp_gflops": 15700,
          "dp_gflops": 7800,
          "memory": 32
        },
        {
          "global_id": 6,
          "local_id": 1,
          "type": "GPU",
          "model": "Tesla V100-SXM2-32GB",
          "sp_gflops": 15700,
          "dp_gflops": 7800,
          "memory": 32
        },
        {
          "global_id": 7,
          "local_id": 2,
          "type": "GPU",
          "model": "Tesla V100-SXM2-32GB",
          "sp_gflops": 15700,
          "dp_gflops": 7800,
          "memory": 32
        },
        {
          "global_id": 8,
          "local_id": 3,
          "type": "GPU",
          "model": "Tesla V100-SXM2-32GB",
          "sp_gflops": 15700,
          "dp_gflops": 7800,
          "memory": 32
        },
        {
          "global_id": 9,
          "local_id": 0,
          "type": "NIC"
        }
      ],
      "links": [
        {
          "source_global_id": 5,
          "target_global_id": 6,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 5,
          "target_global_id": 7,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 5,
          "target_global_id": 8,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 5,
          "target_global_id": 9,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 6,
          "target_global_id": 5,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 6,
          "target_global_id": 7,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 6,
          "target_global_id": 8,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 6,
          "target_global_id": 9,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 7,
          "target_global_id": 5,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 7,
          "target_global_id": 6,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 7,
          "target_global_id": 8,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 7,
          "target_global_id": 9,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 8,
          "target_global_id": 5,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 8,
          "target_global_id": 6,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 8,
          "target_global_id": 7,
          "type": "NVL",
          "bandwidth": 42
        },
        {
          "source_global_id": 8,
          "target_global_id": 9,
          "type": "PHB",
          "bandwidth": 12
        },
        {
          "source_global_id": 9,
          "target_global_id": 4,
          "type": "NET",
          "bandwidth": 1
        }
      ]
    }
  ]
}
"""


class MLPLayer(nn.Layer):

    def __init__(self,
                 hidden_size=64,
                 intermediate_size=4 * 64,
                 initializer_range=0.02):
        super(MLPLayer, self).__init__()
        d_model = hidden_size
        dim_feedforward = intermediate_size
        np.random.seed(2021)
        arr0 = np.random.normal(0, 0.02, size=(d_model, dim_feedforward))
        arr1 = np.random.normal(0, 0.02, size=(dim_feedforward, d_model))
        arr2 = np.random.normal(0, 0.02, size=(d_model, dim_feedforward))
        arr3 = np.random.normal(0, 0.02, size=(dim_feedforward, d_model))
        weight_attr0 = paddle.ParamAttr(initializer=NumpyArrayInitializer(arr0))
        weight_attr1 = paddle.ParamAttr(initializer=NumpyArrayInitializer(arr1))
        weight_attr2 = paddle.ParamAttr(initializer=NumpyArrayInitializer(arr2))
        weight_attr3 = paddle.ParamAttr(initializer=NumpyArrayInitializer(arr3))
        bias_attr = None
        self.linear0 = nn.Linear(d_model,
                                 dim_feedforward,
                                 weight_attr0,
                                 bias_attr=bias_attr)
        self.linear1 = nn.Linear(dim_feedforward,
                                 d_model,
                                 weight_attr1,
                                 bias_attr=bias_attr)
        self.norm = nn.LayerNorm(d_model, epsilon=1e-5)
        self.linear2 = nn.Linear(d_model,
                                 dim_feedforward,
                                 weight_attr2,
                                 bias_attr=bias_attr)
        self.linear3 = nn.Linear(dim_feedforward,
                                 d_model,
                                 weight_attr3,
                                 bias_attr=bias_attr)

    def forward(self, input):
        if _global_parallel_strategy == "dp_mp_pp":
            auto.shard_tensor(self.linear0.weight,
                              dist_attr={
                                  "process_mesh": _global_process_mesh[0],
                                  "dims_mapping": [-1, 1]
                              })
            auto.shard_tensor(self.linear1.weight,
                              dist_attr={
                                  "process_mesh": _global_process_mesh[0],
                                  "dims_mapping": [1, -1]
                              })
            auto.shard_tensor(self.linear2.weight,
                              dist_attr={
                                  "process_mesh": _global_process_mesh[1],
                                  "dims_mapping": [-1, 1]
                              })
            auto.shard_tensor(self.linear3.weight,
                              dist_attr={
                                  "process_mesh": _global_process_mesh[1],
                                  "dims_mapping": [1, -1]
                              })

        out = self.norm(input)
        out = self.linear0(out)
        out = F.gelu(out, approximate=True)
        out = self.linear1(out)

        auto.shard_tensor(out,
                          dist_attr={
                              "process_mesh": _global_process_mesh[1],
                              "dims_mapping": [0, -1]
                          })
        out = self.linear2(out)
        out = F.gelu(out, approximate=True)
        out = self.linear3(out)
        return out


def mlp_forward(train_program, start_program):
    with static.program_guard(train_program,start_program), \
        utils.unique_name.guard():
        batch_size = 4
        hidden_size = 64
        input = static.data(name="input",
                            shape=[batch_size, hidden_size],
                            dtype='float32')
        label = static.data(name="label",
                            shape=[batch_size, 1],
                            dtype='float32')

        if _global_parallel_strategy == "dp_mp_pp":
            auto.shard_tensor(input,
                              dist_attr={
                                  "process_mesh": _global_process_mesh[0],
                                  "dims_mapping": [0, -1]
                              })
        mlp = MLPLayer(hidden_size=hidden_size,
                       intermediate_size=4 * hidden_size,
                       initializer_range=0.02)
        predict = mlp(input)
        error_cost = paddle.nn.functional.square_error_cost(predict, label)
        loss = paddle.mean(error_cost)
    return loss, train_program, start_program


def get_dist_prog(train_program, startup_program, dist_context, rank_id):
    loss, train_program, startup_program = mlp_forward(train_program,
                                                       startup_program)

    fleet._user_defined_strategy = fleet.DistributedStrategy()
    fleet.user_defined_optimizer = paddle.fluid.optimizer.AdamOptimizer()
    parallelizer = AutoParallelizer(fleet)
    parallelizer._dist_context = dist_context

    # auto completion
    completer = Completer(dist_context)
    complete_train_program = completer.complete_forward_annotation(
        train_program)
    dist_context.block_state.parse_forward_blocks(complete_train_program)
    params_grads = parallelizer._generate_backward(complete_train_program,
                                                   startup_program,
                                                   loss,
                                                   parameter_list=None,
                                                   no_grad_set=None,
                                                   callbacks=None)

    partitioner = Partitioner(dist_context, rank_id)
    dist_train_program, dist_startup_prog, dist_params_grads = partitioner.partition(
        complete_train_program, startup_program, params_grads)

    partitioned_optimize_ops = parallelizer._apply_optimize(
        dist_train_program, dist_startup_prog, dist_params_grads)

    resharder = Resharder(dist_train_program, dist_startup_prog, rank_id,
                          dist_context, dist_params_grads)
    resharder.reshard()
    return dist_train_program, dist_startup_prog


def is_in_machine(device_local_id, machine):
    for device in machine.devices.values():
        if device_local_id == device.local_id:
            return True
    return False


def get_device_local_ids(machine):
    local_ids = []
    for device in machine.devices.values():
        local_ids.append[device.local_id]
    return local_ids


class TestAutoParallelMapper(unittest.TestCase):

    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()

    def tearDown(self):
        self.temp_dir.cleanup()

    def test_mapper_dp_mp_pp(self):
        cluster_json_path = os.path.join(self.temp_dir.name,
                                         "auto_parallel_cluster.json")
        cluster_json_object = json.loads(cluster_json)
        with open(cluster_json_path, "w") as cluster_json_file:
            json.dump(cluster_json_object, cluster_json_file)
        cluster = Cluster()
        cluster.build_from_file(cluster_json_path)

        global _global_parallel_strategy
        _global_parallel_strategy = "dp_mp_pp"
        global _global_num_stages
        _global_num_stages = 2
        global _global_process_mesh
        _global_process_mesh = [[[0, 1], [2, 3]], [[4, 5], [6, 7]]]
        processes = [0, 1, 2, 3, 4, 5, 6, 7]

        dist_programs = {}
        for rank_id in processes:
            train_program = static.Program()
            startup_program = static.Program()
            dist_context = DistributedContext()
            dist_train_program, dist_startup_prog = get_dist_prog(
                train_program, startup_program, dist_context, rank_id)
            # if rank_id == 0:
            #   print_program_with_dist_attr(dist_train_program, dist_context)
            dist_programs[rank_id] = [dist_train_program, None]

        rank_mapping = mapping(dist_programs, cluster)

        all_mapped_ranks = set()
        for machine_id, machine_mapping in rank_mapping.items():
            machine = cluster.machines[machine_id]
            machine_mapped_ranks = set()
            machine_mapped_device_local_ids = set()
            for rank, device_ids in machine_mapping["ranks"].items():
                # Only allow one process to one device mapping
                self.assertEqual(len(device_ids), 1)
                self.assertTrue(is_in_machine(device_ids[0], machine))
                machine_mapped_ranks.add(rank)
                machine_mapped_device_local_ids.add(device_ids[0])
            self.assertEqual(len(machine_mapped_ranks),
                             len(machine_mapped_device_local_ids))
            all_mapped_ranks.update(machine_mapped_ranks)
        self.assertEqual(set(processes), all_mapped_ranks)

    def test_mapper_misc(self):
        self.assertEqual(get_dtype_bytes(paddle.float64), 8)
        self.assertEqual(get_dtype_bytes(paddle.float32), 4)
        self.assertEqual(get_dtype_bytes(paddle.float16), 2)
        self.assertEqual(get_dtype_bytes(paddle.bfloat16), 2)
        self.assertEqual(get_dtype_bytes(paddle.int64), 8)
        self.assertEqual(get_dtype_bytes(paddle.int32), 4)
        self.assertEqual(get_dtype_bytes(paddle.int16), 2)
        self.assertEqual(get_dtype_bytes(paddle.int8), 1)
        self.assertEqual(get_dtype_bytes(paddle.uint8), 1)
        self.assertRaises(ValueError, get_dtype_bytes, "unknown type")
        train_program = static.Program()
        startup_program = static.Program()
        ring_id = 0
        root_id = 0
        nranks = 2
        with fluid.program_guard(train_program, startup_program):
            input = layers.data(name="input", shape=[10, 10], dtype='float32')
            output = train_program.current_block().create_var(
                name="outofbroadcast",
                dtype='float32',
                type=core.VarDesc.VarType.LOD_TENSOR,
                persistable=False,
                stop_gradient=False)
            broadcast_op = train_program.global_block().append_op(
                type="c_broadcast",
                inputs={'X': input},
                attrs={
                    'ring_id': ring_id,
                    'root': root_id
                },
                outputs={'Out': output})
            self.assertEqual(get_comm_volume(broadcast_op, 0, 1), 400)
            self.assertEqual(get_comm_volume(broadcast_op, 1, 0), None)
            allgather_op = train_program.global_block().append_op(
                type="c_allgather",
                inputs={'X': input},
                attrs={
                    'ring_id': ring_id,
                    'nranks': nranks
                },
                outputs={'Out': output})
            self.assertEqual(get_comm_volume(allgather_op, 0, 1), 400)
            self.assertEqual(get_comm_volume(allgather_op, 0, 0), None)
            reduce_op = train_program.global_block().append_op(
                type="c_reduce_sum",
                inputs={'X': input},
                attrs={
                    'ring_id': ring_id,
                    'root_id': root_id
                },
                outputs={'Out': output})
            self.assertEqual(get_comm_volume(reduce_op, 0, 1), None)
            self.assertEqual(get_comm_volume(reduce_op, 1, 0), 400)
            cast_op = train_program.global_block().append_op(
                type="cast",
                inputs={"X": input},
                outputs={"Out": output},
                attrs={
                    "in_dtype": fluid.core.VarDesc.VarType.FP32,
                    "out_dtype": fluid.core.VarDesc.VarType.FP32
                })
            self.assertRaises(ValueError, get_comm_volume, cast_op, 0, 1)


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