mindspore2022/tests/ut/python/parallel/test_parallel_transformer.py

# Copyright 2021 Huawei Technologies Co., Ltd
#
# 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 os
import shutil
import glob

import numpy as np
import pytest
import mindspore.common.dtype as mstype
import mindspore.nn as nn
from mindspore import Tensor
from mindspore.context import set_auto_parallel_context, ParallelMode
from mindspore import context
from mindspore.ops import composite as C
from mindspore.ops import functional as F
import mindspore.ops as P
from mindspore.parallel.nn import TransformerEncoder, TransformerDecoder, Transformer, TransformerOpParallelConfig, \
    VocabEmbedding, CrossEntropyLoss, OpParallelConfig, EmbeddingOpParallelConfig, FixedSparseAttention
from mindspore.nn.wrap.loss_scale import DynamicLossScaleUpdateCell
from mindspore.nn.optim import AdamWeightDecay
from mindspore.nn.wrap.cell_wrapper import PipelineCell, _VirtualDatasetCell, TrainOneStepCell
from mindspore.nn.wrap.loss_scale import _TrainPipelineWithLossScaleCell
from mindspore.train import Model
from mindspore.parallel import set_algo_parameters
from tests.dataset_mock import MindData
from tests.ut.python.ops.test_math_ops import VirtualLoss

grad_all = C.GradOperation(get_all=True)


class Dataset(MindData):
    def __init__(self, *inputs, length=3):
        super(Dataset, self).__init__(size=length)
        self.inputs = inputs
        self.index = 0
        self.length = length

    def __iter__(self):
        return self

    def __next__(self):
        if self.index >= self.length:
            raise StopIteration
        self.index += 1
        return self.inputs

    def reset(self):
        self.index = 0


class TransformerNet(nn.Cell):
    def __init__(self, en_layer, de_layer, parallel_config):
        super(TransformerNet, self).__init__()
        self.network = Transformer(encoder_layers=en_layer,
                                   decoder_layers=de_layer,
                                   batch_size=2,
                                   src_seq_length=20,
                                   tgt_seq_length=10,
                                   hidden_size=64,
                                   num_heads=8,
                                   ffn_hidden_size=64,
                                   parallel_config=parallel_config)
        self.loss = CrossEntropyLoss(parallel_config=config.dp_mp_config)

    def construct(self, x1, x2, x3, x4, x5, y, mask):
        predict, _, _ = self.network(x1, x2, x3, x4, x5)
        predict = P.Reshape()(predict, (-1, F.shape(predict)[-1]))
        return self.loss(predict, y, mask)


class TransformerEncoderNet(nn.Cell):
    def __init__(self, batch_size, en_layer, de_layer, parallel_config):
        super(TransformerEncoderNet, self).__init__()
        self.embedding = VocabEmbedding(vocab_size=240, embedding_size=64,
                                        parallel_config=parallel_config.embedding_dp_mp_config)
        self.network = Transformer(encoder_layers=en_layer,
                                   decoder_layers=de_layer,
                                   batch_size=batch_size,
                                   src_seq_length=20,
                                   tgt_seq_length=10,
                                   hidden_size=64,
                                   num_heads=8,
                                   ffn_hifloat16dden_size=64,
                                   parallel_config=parallel_config)
        self.loss = CrossEntropyLoss(parallel_config=config.dp_mp_config)

    def construct(self, x, encoder_mask, label, input_mask):
        embedded, _ = self.embedding(x)
        logits, _, = self.network(embedded, encoder_mask)
        logits = P.Reshape()(logits, (-1, F.shape(logits)[-1]))
        label = P.Reshape()(label, (-1,))
        input_mask = P.Reshape()(input_mask, (-1,))
        return self.loss(logits, label, input_mask)


config = TransformerOpParallelConfig(data_parallel=1, model_parallel=8, vocab_emb_dp=False)
pipeline_config = TransformerOpParallelConfig(data_parallel=2, model_parallel=8, pipeline_stage=4,
                                              micro_batch_num=4, vocab_emb_dp=False)


class NetWithLossFiveInputs(nn.Cell):
    def __init__(self, network):
        super(NetWithLossFiveInputs, self).__init__()
        self.loss = VirtualLoss()
        self.network = network

    def construct(self, x1, x2, x3, x4, x5):
        predict, _, _ = self.network(x1, x2, x3, x4, x5)
        return self.loss(predict)


def run_network_function(dataset, pipeline_net):
    """
    Feature: Test transformer embedding shared.
    Description: a basic function for test compiling.
    Expectation: success.
    """
    params = pipeline_net.trainable_params()
    optimizer = nn.Lamb(params, learning_rate=0.01)
    model = Model(pipeline_net, optimizer=optimizer)
    model.train(2, dataset, dataset_sink_mode=False)


def run_total_transformer_model_head(e_layer,
                                     d_layer,
                                     arg_parallel_config,
                                     mode=ParallelMode.SEMI_AUTO_PARALLEL):
    dp = arg_parallel_config.data_parallel
    mp = arg_parallel_config.model_parallel
    pp = arg_parallel_config.pipeline_stage
    if dp * mp * pp != 1:
        set_auto_parallel_context(device_num=8,
                                  full_batch=True,
                                  global_rank=0, parallel_mode=mode)

    encoder_input_value = Tensor(np.ones((2, 20, 64)), mstype.float32)
    encoder_input_mask = Tensor(np.ones((2, 20, 20)), mstype.float16)
    decoder_input_value = Tensor(np.ones((2, 10, 64)), mstype.float32)
    decoder_input_mask = Tensor(np.ones((2, 10, 10)), mstype.float16)
    memory_mask = Tensor(np.ones((2, 10, 20)), mstype.float16)
    seq = 20
    if d_layer > 0:
        seq = 10
    label = Tensor(np.ones((2 * seq,)), mstype.int32)
    input_mask = Tensor(np.ones((2 * seq,)), mstype.float32)
    net = TransformerNet(en_layer=e_layer, de_layer=d_layer, parallel_config=arg_parallel_config)
    net = _VirtualDatasetCell(net)
    params = net.trainable_params()
    optimizer = AdamWeightDecay(params)
    dataset = Dataset(encoder_input_value, encoder_input_mask, decoder_input_value, decoder_input_mask,
                      memory_mask, label, input_mask)
    net_with_grad = TrainOneStepCell(net, optimizer=optimizer)
    model = Model(net_with_grad)

    model.train(1, dataset, dataset_sink_mode=False)


def test_transformer_model():
    set_auto_parallel_context(device_num=8, global_rank=0,
                              full_batch=True,
                              parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
    net = Transformer(encoder_layers=1,
                      decoder_layers=2,
                      batch_size=2,
                      src_seq_length=20,
                      tgt_seq_length=10,
                      hidden_size=64,
                      num_heads=8,
                      ffn_hidden_size=64,
                      parallel_config=config)

    encoder_input_value = Tensor(np.ones((2, 20, 64)), mstype.float32)
    encoder_input_mask = Tensor(np.ones((2, 20, 20)), mstype.float16)
    decoder_input_value = Tensor(np.ones((2, 10, 64)), mstype.float32)
    decoder_input_mask = Tensor(np.ones((2, 10, 10)), mstype.float16)
    memory_mask = Tensor(np.ones((2, 10, 20)), mstype.float16)
    net = NetWithLossFiveInputs(net)
    net = _VirtualDatasetCell(net)
    params = net.trainable_params()
    optimizer = AdamWeightDecay(params)
    dataset = Dataset(encoder_input_value, encoder_input_mask, decoder_input_value, decoder_input_mask,
                      memory_mask)
    net_with_grad = TrainOneStepCell(net, optimizer=optimizer)
    model = Model(net_with_grad)

    model.train(1, dataset, dataset_sink_mode=False)


def test_transformer_model_2d_inputs():
    set_auto_parallel_context(device_num=8, global_rank=0,
                              full_batch=True,
                              parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
    net = Transformer(encoder_layers=1,
                      decoder_layers=2,
                      batch_size=2,
                      src_seq_length=20,
                      tgt_seq_length=10,
                      hidden_size=64,
                      num_heads=8,
                      ffn_hidden_size=64,
                      parallel_config=config)

    encoder_input_value = Tensor(np.ones((40, 64)), mstype.float32)
    encoder_input_mask = Tensor(np.ones((2, 20, 20)), mstype.float16)
    decoder_input_value = Tensor(np.ones((20, 64)), mstype.float32)
    decoder_input_mask = Tensor(np.ones((2, 10, 10)), mstype.float16)
    memory_mask = Tensor(np.ones((2, 10, 20)), mstype.float16)
    net = NetWithLossFiveInputs(net)
    net = _VirtualDatasetCell(net)
    params = net.trainable_params()
    optimizer = AdamWeightDecay(params)
    dataset = Dataset(encoder_input_value, encoder_input_mask, decoder_input_value, decoder_input_mask,
                      memory_mask)
    net_with_grad = TrainOneStepCell(net, optimizer=optimizer)
    model = Model(net_with_grad)

    model.train(1, dataset, dataset_sink_mode=False)


class TestTransformerEmbeddingHead:
    def __init__(self):
        self.output_path = None

    def setup_method(self):
        self.output_path = './graphs' + self.__str__()
        context.set_context(save_graphs=True,
                            save_graphs_path=self.output_path)

    def teardown_method(self):
        shutil.rmtree(self.output_path)

    def virtual_assign_add_from_ir(self, pattern, target_count):
        """
        This function will check the assign aa count with the golden one.
        :param pattern: The match pattern for the specific count
        :param target_count: The gold float16 count in the Ir files
        """
        ir_files = glob.glob(os.path.join(self.output_path, 'rank_0', '*_validate*.ir'))
        assert len(ir_files) == 1
        appear_count = 0
        with open(ir_files[0], 'r') as fp:
            for line in fp:
                if pattern in line:
                    appear_count += 1
        assert appear_count == target_count

    def test_pipeline_with_embedding(self):
        """
        Feature: Test Transformer with embedding as shared
        Description: When do pipeline training and applied optimzier shard, the embedding which is model parallel will
                     raise the shape error. This test cast is ensure there is no error raised.
        Expectation: The number of AssignAdd is not as expected.
        """
        bs = 16
        pp = 2
        context.set_auto_parallel_context(device_num=8, global_rank=0, pipeline_stages=pp,
                                          full_batch=True,
                                          enable_parallel_optimizer=True)
        context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
        cf = TransformerOpParallelConfig(data_parallel=1, model_parallel=4, pipeline_stage=pp, vocab_emb_dp=False)
        pipeline_net = TransformerEncoderNet(batch_size=bs // pp,
                                             en_layer=2, de_layer=0, parallel_config=cf)
        pipeline_net.embedding.pipeline_stage = 0
        pipeline_net.network.encoder.blocks[0].pipeline_stage = 0
        pipeline_net.network.encoder.blocks[1].pipeline_stage = 1

        pipeline_cell_net = PipelineCell(pipeline_net, 2)
        encoder_input_value = Tensor(np.ones((bs, 20)), mstype.int32)
        encoder_input_mask = Tensor(np.ones((bs, 20, 20)), mstype.float16)
        label = Tensor(np.ones((bs, 20)), mstype.int32)
        mask = Tensor(np.ones((bs, 20)), mstype.float32)
        dataset = Dataset(encoder_input_value, encoder_input_mask, label, mask)
        run_network_function(dataset, pipeline_cell_net)
        self.virtual_assign_add_from_ir(pattern=r'AssignAdd(', target_count=35)


def test_transformer_model_int64_inputs():
    set_auto_parallel_context(device_num=8, global_rank=0,
                              full_batch=True,
                              parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
    net = Transformer(encoder_layers=1,
                      decoder_layers=2,
                      batch_size=2,
                      src_seq_length=20,
                      tgt_seq_length=10,
                      hidden_size=64,
                      num_heads=8,
                      ffn_hidden_size=64,
                      parallel_config=config)

    encoder_input_value = Tensor(np.ones((2, 20, 64)), mstype.int64)
    encoder_input_mask = Tensor(np.ones((2, 20, 20)), mstype.float16)
    decoder_input_value = Tensor(np.ones((2, 10, 64)), mstype.float32)
    decoder_input_mask = Tensor(np.ones((2, 10, 10)), mstype.float16)
    memory_mask = Tensor(np.ones((2, 10, 20)), mstype.float16)
    net = NetWithLossFiveInputs(net)
    net = _VirtualDatasetCell(net)
    params = net.trainable_params()
    optimizer = AdamWeightDecay(params)
    dataset = Dataset(encoder_input_value, encoder_input_mask, decoder_input_value, decoder_input_mask,
                      memory_mask)
    net_with_grad = TrainOneStepCell(net, optimizer=optimizer)
    model = Model(net_with_grad)

    with pytest.raises(TypeError):
        model.train(1, dataset, dataset_sink_mode=False)


def test_transformer_model_head_parallel_only_encoder():
    local_config = TransformerOpParallelConfig(data_parallel=1, model_parallel=8)
    run_total_transformer_model_head(e_layer=2, d_layer=0, arg_parallel_config=local_config)


def test_transformer_model_head_parallel():
    local_config = TransformerOpParallelConfig(data_parallel=1, model_parallel=8)
    run_total_transformer_model_head(e_layer=1, d_layer=1, arg_parallel_config=local_config)


def test_transformer_model_head_parallel_decoder():
    local_config = TransformerOpParallelConfig(data_parallel=1, model_parallel=8)
    with pytest.raises(ValueError):
        run_total_transformer_model_head(e_layer=0, d_layer=1, arg_parallel_config=local_config)


def test_transformer_model_head_stand_alone():
    local_config = TransformerOpParallelConfig(data_parallel=1, model_parallel=1)
    run_total_transformer_model_head(e_layer=2, d_layer=2, arg_parallel_config=local_config)


def test_transformer_model_auto_parallel_no_support():
    local_config = TransformerOpParallelConfig(data_parallel=8, model_parallel=1)
    with pytest.raises(RuntimeError):
        run_total_transformer_model_head(e_layer=2, d_layer=2, arg_parallel_config=local_config,
                                         mode=ParallelMode.AUTO_PARALLEL)


def pipeline_single_transformer(grad_accumulation_shard=False):
    """
    Feature: Gradient Accumulation Shard for Pipeline and Gradient Accumulation
    Description: Test a single transformer model with pipeline parallel with grad_accumulation_shard False
    Expectation: The compile passed
    """
    set_auto_parallel_context(device_num=64,
                              full_batch=True,
                              pipeline_stages=pipeline_config.pipeline_stage, global_rank=0,
                              parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
    context.set_auto_parallel_context(parallel_optimizer_config=
                                      {"gradient_accumulation_shard": grad_accumulation_shard})

    net = Transformer(batch_size=8 // pipeline_config.micro_batch_num,
                      src_seq_length=20,
                      tgt_seq_length=10,
                      encoder_layers=2,
                      decoder_layers=2,
                      hidden_size=64,
                      num_heads=8,
                      ffn_hidden_size=64,
                      parallel_config=pipeline_config)

    encoder_input_value = Tensor(np.ones((8, 20, 64)), mstype.float32)
    encoder_input_mask = Tensor(np.ones((8, 20, 20)), mstype.float16)
    decoder_input_value = Tensor(np.ones((8, 10, 64)), mstype.float32)
    decoder_input_mask = Tensor(np.ones((8, 10, 10)), mstype.float16)
    memory_mask = Tensor(np.ones((8, 10, 20)), mstype.float16)
    net = NetWithLossFiveInputs(net)
    net = PipelineCell(net, pipeline_config.micro_batch_num)
    net = _VirtualDatasetCell(net)
    params = net.infer_param_pipeline_stage()
    optimizer = AdamWeightDecay(params)
    dataset = Dataset(encoder_input_value, encoder_input_mask, decoder_input_value, decoder_input_mask,
                      memory_mask)
    update_cell = DynamicLossScaleUpdateCell(loss_scale_value=1024, scale_factor=2, scale_window=1000)
    net_with_grad = _TrainPipelineWithLossScaleCell(net, optimizer=optimizer,
                                                    scale_sense=update_cell)
    model = Model(net_with_grad)

    model.train(1, dataset, dataset_sink_mode=False)


def test_pipeline_transformer_gradient_shard_true():
    """
    Feature: Gradient Accumulation Shard for Pipeline and Gradient Accumulation
    Description: Test a single transformer model with pipeline parallel with grad_accumulation_shard True
    Expectation: The compile passed
    """
    pipeline_single_transformer(grad_accumulation_shard=True)


def test_pipeline_transformer_gradient_shard_false():
    """
    Feature: Gradient Accumulation Shard for Pipeline and Gradient Accumulation
    Description: Test a single transformer model with pipeline parallel with grad_accumulation_shard False
    Expectation: The compile passed
    """
    pipeline_single_transformer(grad_accumulation_shard=False)


def test_transformer_wrong_head():
    set_auto_parallel_context(device_num=64,
                              full_batch=True,
                              pipeline_stages=pipeline_config.pipeline_stage, global_rank=0,
                              parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
    error_test_config = TransformerOpParallelConfig(data_parallel=1, model_parallel=8, vocab_emb_dp=False)
    with pytest.raises(ValueError):
        net = Transformer(batch_size=4,
                          src_seq_length=20,
                          tgt_seq_length=10,
                          encoder_layers=2,
                          decoder_layers=2,
                          hidden_size=64,
                          num_heads=7,
                          ffn_hidden_size=64,
                          parallel_config=error_test_config)

    with pytest.raises(ValueError):
        net = Transformer(batch_size=4,
                          src_seq_length=20,
                          tgt_seq_length=10,
                          encoder_layers=2,
                          decoder_layers=2,
                          hidden_size=63,
                          num_heads=7,
                          ffn_hidden_size=64,
                          parallel_config=error_test_config)
        del net


def test_transformer_wrong_dp_no_error():
    set_auto_parallel_context(device_num=64, full_batch=False, parallel_mode=ParallelMode.DATA_PARALLEL,
                              pipeline_stages=pipeline_config.pipeline_stage, global_rank=0)
    check_config = TransformerOpParallelConfig(data_parallel=8, model_parallel=1, vocab_emb_dp=False)
    net = Transformer(batch_size=4, src_seq_length=20, tgt_seq_length=10, encoder_layers=2,
                      decoder_layers=2, hidden_size=64, num_heads=2, ffn_hidden_size=64,
                      parallel_config=check_config)
    del net


def test_transformer_wrong_semi_auto_dp_error():
    set_auto_parallel_context(device_num=64, full_batch=False, parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL,
                              pipeline_stages=pipeline_config.pipeline_stage, global_rank=0)
    check_config = TransformerOpParallelConfig(data_parallel=16, model_parallel=1, vocab_emb_dp=False)
    with pytest.raises(ValueError):
        net = Transformer(batch_size=4, src_seq_length=20, tgt_seq_length=10, encoder_layers=2,
                          decoder_layers=2, hidden_size=64, num_heads=2, ffn_hidden_size=64,
                          parallel_config=check_config)
        del net


def test_encoder():
    class NetWithLoss(nn.Cell):
        def __init__(self, network):
            super(NetWithLoss, self).__init__()
            self.loss = VirtualLoss()
            self.network = network

        def construct(self, x1, x2):
            predict, _ = self.network(x1, x2)
            return self.loss(predict)

    set_auto_parallel_context(device_num=8,
                              full_batch=True,
                              global_rank=0, parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
    net = TransformerEncoder(num_layers=2,
                             batch_size=2,
                             seq_length=16,
                             hidden_size=8,
                             ffn_hidden_size=64,
                             num_heads=8,
                             parallel_config=config)

    encoder_input_value = Tensor(np.ones((2, 16, 8)), mstype.float32)
    encoder_input_mask = Tensor(np.ones((2, 16, 16)), mstype.float16)

    net = NetWithLoss(net)

    net = _VirtualDatasetCell(net)

    dataset = Dataset(encoder_input_value, encoder_input_mask)

    model = Model(net)

    model.train(1, dataset, dataset_sink_mode=False)


def test_decoder():
    class NetWithLoss(nn.Cell):
        def __init__(self, network):
            super(NetWithLoss, self).__init__()
            self.loss = VirtualLoss()
            self.network = network

        def construct(self, x1, x2, x3, x4):
            predict, _, _ = self.network(x1, x2, x3, x4)
            return self.loss(predict)

    set_auto_parallel_context(device_num=8,
                              full_batch=True,
                              global_rank=0, parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
    net = TransformerDecoder(num_layers=1,
                             batch_size=8,
                             hidden_size=16,
                             ffn_hidden_size=8,
                             num_heads=8,
                             src_seq_length=20,
                             tgt_seq_length=10,
                             parallel_config=config)

    encoder_input_value = Tensor(np.ones((8, 20, 16)), mstype.float32)
    decoder_input_value = Tensor(np.ones((8, 10, 16)), mstype.float32)
    decoder_input_mask = Tensor(np.ones((8, 10, 10)), mstype.float16)
    memory_mask = Tensor(np.ones((8, 10, 20)), mstype.float16)

    net = NetWithLoss(net)

    net = _VirtualDatasetCell(net)

    dataset = Dataset(decoder_input_value, decoder_input_mask, encoder_input_value, memory_mask)

    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_vocabembedding_dp_true():
    set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)

    class NetWithLoss(nn.Cell):
        def __init__(self, network):
            super(NetWithLoss, self).__init__()
            self.loss = VirtualLoss()
            self.network = network

        def construct(self, x1):
            predict, _ = self.network(x1)
            return self.loss(predict)

    net = VocabEmbedding(vocab_size=160, embedding_size=16, parallel_config=config.embedding_dp_mp_config)
    net = NetWithLoss(net)
    net = _VirtualDatasetCell(net)
    encoder_input_value = Tensor(np.ones((2, 64)), mstype.int32)
    dataset = Dataset(encoder_input_value)

    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_vocabembedding_dp_false():
    set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)

    class NetWithLoss(nn.Cell):
        def __init__(self, network):
            super(NetWithLoss, self).__init__()
            self.loss = VirtualLoss()
            self.network = network

        def construct(self, x1):
            predict, _ = self.network(x1)
            return self.loss(predict)

    net = VocabEmbedding(vocab_size=160, embedding_size=16, parallel_config=config.embedding_dp_mp_config)
    net = NetWithLoss(net)
    net = _VirtualDatasetCell(net)
    encoder_input_value = Tensor(np.ones((2, 64)), mstype.int32)
    dataset = Dataset(encoder_input_value)

    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_sparse_attention_parallel_mp():
    set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode=ParallelMode.AUTO_PARALLEL)
    set_algo_parameters(fully_use_devices=False)
    sparse_attention_config = OpParallelConfig(model_parallel=8)
    net = FixedSparseAttention(batch_size=16,
                               seq_length=1024,
                               size_per_head=64,
                               num_heads=8,
                               block_size=64,
                               parallel_config=sparse_attention_config)
    q = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    k = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    v = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    mask = Tensor(np.ones((2, 1024, 1024)), mstype.float32)
    dataset = Dataset(q, k, v, mask)
    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_sparse_attention_parallel_mix():
    set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode=ParallelMode.AUTO_PARALLEL)
    set_algo_parameters(fully_use_devices=False)
    sparse_attention_config = OpParallelConfig(data_parallel=2, model_parallel=4)
    net = FixedSparseAttention(batch_size=16,
                               seq_length=1024,
                               size_per_head=64,
                               num_heads=8,
                               block_size=64,
                               parallel_config=sparse_attention_config)
    q = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    k = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    v = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    mask = Tensor(np.ones((2, 1024, 1024)), mstype.float32)
    dataset = Dataset(q, k, v, mask)
    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_sparse_attention_parallel_mix1():
    set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode=ParallelMode.AUTO_PARALLEL)
    set_algo_parameters(fully_use_devices=False)
    sparse_attention_config = OpParallelConfig(data_parallel=4, model_parallel=2)
    net = FixedSparseAttention(batch_size=16,
                               seq_length=1024,
                               size_per_head=64,
                               num_heads=8,
                               block_size=64,
                               parallel_config=sparse_attention_config)
    q = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    k = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    v = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    mask = Tensor(np.ones((2, 1024, 1024)), mstype.float32)
    dataset = Dataset(q, k, v, mask)
    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_sparse_attention_parallel_dp():
    set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode=ParallelMode.AUTO_PARALLEL)
    set_algo_parameters(fully_use_devices=False)
    sparse_attention_config = OpParallelConfig(data_parallel=8, model_parallel=1)
    net = FixedSparseAttention(batch_size=16,
                               seq_length=1024,
                               size_per_head=64,
                               num_heads=8,
                               block_size=64,
                               parallel_config=sparse_attention_config)
    net = _VirtualDatasetCell(net)
    q = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    k = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    v = Tensor(np.ones((2, 1024, 512)), mstype.float16)
    mask = Tensor(np.ones((2, 1024, 1024)), mstype.float32)
    dataset = Dataset(q, k, v, mask)
    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_parallel_cross_entroy_loss_semi_auto_parallel():
    set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode=ParallelMode.AUTO_PARALLEL)

    class NetWithLoss(nn.Cell):
        def __init__(self, network, config_setting):
            super(NetWithLoss, self).__init__()
            self.loss = CrossEntropyLoss(config_setting)
            self.network = network

        def construct(self, x1, x2, x3):
            predict, _ = self.network(x1)
            predict = P.Reshape()(predict, (-1, 16))
            return self.loss(predict, x2, x3)

    net = VocabEmbedding(vocab_size=160, embedding_size=16, parallel_config=config.embedding_dp_mp_config)
    net = NetWithLoss(net, config.dp_mp_config)
    net = _VirtualDatasetCell(net)
    embed_ids = Tensor(np.ones((2, 64)), mstype.int32)
    labels = Tensor(np.ones((2 * 64,)), mstype.int32)
    input_mask = Tensor(np.ones((2 * 64,)), mstype.float32)
    dataset = Dataset(embed_ids, labels, input_mask)

    model = Model(net)
    model.train(1, dataset, dataset_sink_mode=False)


def test_transformer_args():

    with pytest.raises(TypeError):
        Transformer(hidden_size=10, batch_size=2, ffn_hidden_size=20, src_seq_length=10,
                    tgt_seq_length=20, decoder_layers="aa")

    with pytest.raises(TypeError):
        Transformer(hidden_size=10, batch_size=2, ffn_hidden_size=20, src_seq_length=10,
                    tgt_seq_length="a")

    with pytest.raises(TypeError):
        Transformer(hidden_size=10, batch_size=2, ffn_hidden_size=20, src_seq_length=10,
                    tgt_seq_length=20, softmax_compute_type=mstype.int64)

    with pytest.raises(TypeError):
        Transformer(hidden_size=10, batch_size=2, ffn_hidden_size=20, src_seq_length=10,
                    tgt_seq_length=20, layernorm_compute_type=mstype.int64)

    with pytest.raises(TypeError):
        Transformer(hidden_size=10, batch_size=2, ffn_hidden_size=20, src_seq_length=10,
                    tgt_seq_length=20, param_init_type=mstype.int64)

    with pytest.raises(TypeError):
        Transformer(hidden_size=10, batch_size=2, ffn_hidden_size=20, src_seq_length=10,
                    tgt_seq_length=20, hidden_dropout_rate=mstype.int64)

    Transformer(hidden_size=10, batch_size=2, ffn_hidden_size=20, src_seq_length=10,
                tgt_seq_length=20, softmax_compute_type=mstype.float16)


def test_transformer_parallel_config():
    parallel_test_config = TransformerOpParallelConfig(data_parallel=1, model_parallel=3)

    with pytest.raises(TypeError):
        parallel_test_config.data_parallel = False

    with pytest.raises(ValueError):
        parallel_test_config.data_parallel = 0

    with pytest.raises(TypeError):
        parallel_test_config.model_parallel = False

    with pytest.raises(ValueError):
        parallel_test_config.model_parallel = 0

    with pytest.raises(TypeError):
        parallel_test_config.pipeline_stage = False

    with pytest.raises(ValueError):
        parallel_test_config.pipeline_stage = 0

    with pytest.raises(TypeError):
        parallel_test_config.micro_batch_num = False

    with pytest.raises(ValueError):
        parallel_test_config.micro_batch_num = 0

    with pytest.raises(TypeError):
        parallel_test_config.gradient_aggregation_group = False

    with pytest.raises(ValueError):
        parallel_test_config.gradient_aggregation_group = 0

    with pytest.raises(TypeError):
        parallel_test_config.recompute = 1

    parallel_test_config.recompute.recompute = False

    assert not parallel_test_config.recompute.recompute


def test_parallel_config():
    parallel_test_config = OpParallelConfig(data_parallel=1, model_parallel=3)

    with pytest.raises(ValueError):
        parallel_test_config.data_parallel = 0

    with pytest.raises(TypeError):
        parallel_test_config.model_parallel = False

    with pytest.raises(ValueError):
        parallel_test_config.model_parallel = 0

    assert parallel_test_config.model_parallel == 3


def test_embedding_parallel_config():
    parallel_test_config = EmbeddingOpParallelConfig(data_parallel=1, model_parallel=3, vocab_emb_dp=False)

    with pytest.raises(ValueError):
        parallel_test_config.data_parallel = 0

    with pytest.raises(TypeError):
        parallel_test_config.model_parallel = False

    with pytest.raises(ValueError):
        parallel_test_config.model_parallel = 0

    with pytest.raises(TypeError):
        parallel_test_config.vocab_emb_dp = 0

    assert not parallel_test_config.vocab_emb_dp