1# Copyright 2020 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15import numpy as np 16import pytest 17 18import mindspore as ms 19from mindspore import context, Tensor, Parameter 20from mindspore.common.api import _cell_graph_executor 21from mindspore.common.initializer import initializer 22from mindspore.nn import Cell, TrainOneStepCell, Momentum 23from mindspore.ops import operations as P 24 25 26class Net(Cell): 27 def __init__(self, mul_weight, strategy1=None, strategy2=None, strategy3=None): 28 super().__init__() 29 self.begin_norm_axis = 2 30 self.begin_params_axis = 1 31 self.mul = P.Mul().shard(strategy1) 32 self.layer_norm = P.LayerNorm(self.begin_norm_axis, self.begin_params_axis).shard(strategy2) 33 self.mul2 = P.Mul().shard(strategy3) 34 self.mul_weight = Parameter(mul_weight, "w1") 35 self.normalized_shape = [64, 32, 16] 36 self.gamma = Parameter(initializer('ones', self.normalized_shape), name="gamma") 37 self.beta = Parameter(initializer('zeros', self.normalized_shape), name="beta") 38 39 def construct(self, x, b): 40 out = self.mul(x, self.mul_weight) 41 out, _, _ = self.layer_norm(out, self.gamma, self.beta) 42 out = self.mul2(out, b) 43 return out 44 45 46_x = Tensor(np.ones([16, 64, 32, 16]), dtype=ms.float32) 47_w = Tensor(np.ones([16, 64, 32, 16]), dtype=ms.float32) 48_b = Tensor(np.ones([16, 64, 32, 16]), dtype=ms.float32) 49 50 51def compile_net(net): 52 optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9) 53 train_net = TrainOneStepCell(net, optimizer) 54 train_net.set_auto_parallel() 55 train_net.set_train() 56 _cell_graph_executor.compile(train_net, _x, _b) 57 context.reset_auto_parallel_context() 58 59 60def test_layer_norm_data_parallel(): 61 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 62 strategy1 = ((16, 1, 1, 1), (16, 1, 1, 1)) 63 strategy2 = ((16, 1, 1, 1), (1, 1, 1), (1, 1, 1)) 64 strategy3 = ((16, 1, 1, 1), (16, 1, 1, 1)) 65 net = Net(_w, strategy1, strategy2, strategy3) 66 compile_net(net) 67 68 69def test_layer_norm_model_parallel(): 70 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 71 strategy1 = ((1, 16, 1, 1), (1, 16, 1, 1)) 72 strategy2 = ((1, 16, 1, 1), (16, 1, 1), (16, 1, 1)) 73 strategy3 = ((1, 16, 1, 1), (1, 16, 1, 1)) 74 net = Net(_w, strategy1, strategy2, strategy3) 75 compile_net(net) 76 77 78def test_layer_norm_hybrid_parallel(): 79 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 80 strategy1 = ((2, 8, 1, 1), (2, 8, 1, 1)) 81 strategy2 = ((2, 8, 1, 1), (8, 1, 1), (8, 1, 1)) 82 strategy3 = ((2, 8, 1, 1), (2, 8, 1, 1)) 83 net = Net(_w, strategy1, strategy2, strategy3) 84 compile_net(net) 85 86 87def test_layer_norm_auto_parallel(): 88 context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=16, global_rank=0) 89 net = Net(_w) 90 compile_net(net) 91 92 93def test_layer_norm_repeat_calc(): 94 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 95 strategy1 = ((2, 2, 4, 1), (2, 2, 4, 1)) 96 strategy2 = ((2, 2, 1, 1), (2, 1, 1), (2, 1, 1)) 97 strategy3 = ((2, 2, 4, 1), (2, 2, 4, 1)) 98 net = Net(_w, strategy1, strategy2, strategy3) 99 compile_net(net) 100 101 102def test_layer_norm_wrong_strategy(): 103 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 104 strategy1 = ((2, 2, 4, 1), (2, 2, 4, 1)) 105 strategy2 = ((1, 2, 1, 2), (2, 1, 2), (2, 1, 2)) 106 strategy3 = ((2, 2, 4, 1), (2, 2, 4, 1)) 107 net = Net(_w, strategy1, strategy2, strategy3) 108 with pytest.raises(RuntimeError): 109 compile_net(net) 110