1# Copyright 2019 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 16 17import mindspore as ms 18import mindspore.nn as nn 19from mindspore import Tensor 20from mindspore import context 21from mindspore.common.api import _cell_graph_executor 22from mindspore.ops import composite as C 23from mindspore.ops import operations as P 24from tests.ut.python.ops.test_math_ops import VirtualLoss 25 26 27grad_all = C.GradOperation(get_all=True) 28 29 30class NetWithLoss(nn.Cell): 31 def __init__(self, network): 32 super(NetWithLoss, self).__init__() 33 self.loss = VirtualLoss() 34 self.network = network 35 36 def construct(self, x, y, b): 37 predict = self.network(x, y, b) 38 return self.loss(predict) 39 40 41class GradWrap(nn.Cell): 42 def __init__(self, network): 43 super(GradWrap, self).__init__() 44 self.network = network 45 46 def construct(self, x, y, b): 47 return grad_all(self.network)(x, y, b) 48 49 50def compile_net(net, x, y, b): 51 net.set_auto_parallel() 52 net.set_train() 53 _cell_graph_executor.compile(net, x, y, b) 54 55 56# model_parallel test 57def test_two_matmul(): 58 class Net(nn.Cell): 59 def __init__(self, strategy1, strategy2): 60 super().__init__() 61 self.matmul1 = P.MatMul().shard(strategy1) 62 self.matmul2 = P.MatMul().shard(strategy2) 63 64 def construct(self, x, y, b): 65 out = self.matmul1(x, y) 66 out = self.matmul2(out, b) 67 return out 68 69 context.set_auto_parallel_context(device_num=8, global_rank=0, gradients_mean=True) 70 strategy1 = ((4, 2), (2, 1)) 71 strategy2 = ((2, 4), (4, 1)) 72 net = GradWrap(NetWithLoss(Net(strategy1, strategy2))) 73 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel") 74 75 x = Tensor(np.ones([128, 32]), dtype=ms.float32) 76 y = Tensor(np.ones([32, 64]), dtype=ms.float32) 77 b = Tensor(np.ones([64, 64]), dtype=ms.float32) 78 79 compile_net(net, x, y, b) 80 81 82def test_two_matmul_repeated_calculation1(): 83 class Net(nn.Cell): 84 def __init__(self, strategy1, strategy2): 85 super().__init__() 86 self.matmul1 = P.MatMul().shard(strategy1) 87 self.matmul2 = P.MatMul().shard(strategy2) 88 89 def construct(self, x, y, b): 90 out = self.matmul1(x, y) 91 out = self.matmul2(out, b) 92 return out 93 94 context.set_auto_parallel_context(device_num=64, global_rank=5, gradients_mean=True) 95 strategy1 = ((2, 4), (4, 8)) 96 strategy2 = ((1, 1), (1, 1)) 97 net = GradWrap(NetWithLoss(Net(strategy1, strategy2))) 98 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel") 99 100 x = Tensor(np.ones([128, 32]), dtype=ms.float32) 101 y = Tensor(np.ones([32, 64]), dtype=ms.float32) 102 b = Tensor(np.ones([64, 64]), dtype=ms.float32) 103 compile_net(net, x, y, b) 104 105 106def test_two_matmul_repeated_calculation2(): 107 class Net(nn.Cell): 108 def __init__(self, strategy1, strategy2): 109 super().__init__() 110 self.matmul1 = P.MatMul().shard(strategy1) 111 self.matmul2 = P.MatMul().shard(strategy2) 112 113 def construct(self, x, y, b): 114 out = self.matmul1(x, y) 115 out = self.matmul2(out, b) 116 return out 117 118 context.set_auto_parallel_context(device_num=64, global_rank=15) 119 strategy1 = ((2, 4), (4, 8)) 120 strategy2 = ((2, 2), (2, 1)) 121 net = GradWrap(NetWithLoss(Net(strategy1, strategy2))) 122 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel") 123 124 x = Tensor(np.ones([128, 32]), dtype=ms.float32) 125 y = Tensor(np.ones([32, 64]), dtype=ms.float32) 126 b = Tensor(np.ones([64, 64]), dtype=ms.float32) 127 compile_net(net, x, y, b) 128 129 130def test_matmul_forward_reduce_scatter(): 131 class Net(nn.Cell): 132 def __init__(self, strategy1, strategy2): 133 super().__init__() 134 self.matmul = P.MatMul().shard(strategy1) 135 self.matmul.add_prim_attr("forward_reduce_scatter", True) 136 self.mul = P.Mul().shard(strategy2) 137 138 def construct(self, x, y, b): 139 out = self.matmul(x, y) 140 out = self.mul(out, b) 141 return out 142 143 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=8, global_rank=0) 144 strategy1 = ((2, 2), (2, 2)) 145 strategy2 = ((4, 2), (4, 2)) 146 net = GradWrap(NetWithLoss(Net(strategy1, strategy2))) 147 148 x = Tensor(np.ones([128, 32]), dtype=ms.float32) 149 y = Tensor(np.ones([32, 64]), dtype=ms.float32) 150 b = Tensor(np.ones([128, 64]), dtype=ms.float32) 151 compile_net(net, x, y, b) 152 153 154def test_matmul_forward_reduce_scatter_transpose(): 155 class Net(nn.Cell): 156 def __init__(self, strategy1, strategy2): 157 super().__init__() 158 self.matmul = P.MatMul(transpose_b=True).shard(strategy1) 159 self.matmul.add_prim_attr("forward_reduce_scatter", True) 160 self.mul = P.Mul().shard(strategy2) 161 162 def construct(self, x, y, b): 163 out = self.matmul(x, y) 164 out = self.mul(out, b) 165 return out 166 167 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 168 strategy1 = ((2, 4), (2, 4)) 169 strategy2 = ((8, 2), (8, 2)) 170 net = GradWrap(NetWithLoss(Net(strategy1, strategy2))) 171 172 x = Tensor(np.ones([128, 32]), dtype=ms.float32) 173 y = Tensor(np.ones([64, 32]), dtype=ms.float32) 174 b = Tensor(np.ones([128, 64]), dtype=ms.float32) 175 compile_net(net, x, y, b) 176