mindspore-Ecosystem
/
mindspore
21
0
Fork 27
Code Issues Pull Requests 6 Projects Releases Wiki Activity

!10389 Fix the Backward of the AllSwap Operation 

From: @huangxinjing
Reviewed-by: @stsuteng,@stsuteng,@zhunaipan
Signed-off-by: @stsuteng,@stsuteng
This commit is contained in:
mindspore-ci-bot 2020-12-24 09:38:29 +08:00 committed by Gitee
parent 1db198f3a6 565ce81b29
commit e1248a7246
2 changed files with 14 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mindspore/ops/_grad/grad_comm_ops.py
View File

@ -183,11 +183,11 @@ def get_bprop_allswap(self):
all_swap_grad = AllSwap(self.group)
if self.instance_name:
instance_name = "grad" + self.instance_name
all_to_all_grad.set_prim_instance_name(instance_name)
all_swap_grad.set_prim_instance_name(instance_name)
def bprop(x, send_size, recv_size, out, dout):
dx = all_swap_grad(dout, recv_size, send_size)
return (dx,)
return (dx, zeros_like(send_size), zeros_like(recv_size))
return bprop

16
tests/ut/python/communication/test_comm.py
View File

@ -27,7 +27,7 @@ from mindspore.nn import ReLU
from mindspore.nn import TrainOneStepCell, WithLossCell
from mindspore.ops.operations.comm_ops import AllReduce, AllGather, _AlltoAll, ReduceOp, ReduceScatter
from mindspore.ops.operations.comm_ops import Broadcast, AllSwap
from mindspore.ops.operations.math_ops import ReduceSum
from mindspore.ops.operations.array_ops import GatherV2
import mindspore
# pylint: disable=W0212
@ -127,14 +127,15 @@ class AllSwapNet(nn.Cell):
self.dense = Dense(input_channel, out_channel)
self.allswap = AllSwap()
self.relu = ReLU()
self.reduce = ReduceSum()
part_slice = batch_size / 2
self.send_size = Tensor([0, part_slice*out_channel, part_slice*out_channel], mindspore.int64)
self.recv_size = Tensor([part_slice*out_channel, part_slice*out_channel, 0], mindspore.int64)
self.gatherv2 = GatherV2()
self.input = Tensor(np.ones([1]), mindspore.int32)
def construct(self, x):
x = self.dense(x)
x = self.allswap(x, self.send_size, self.recv_size)
x = self.relu(x)
x = self.gatherv2(x, self.input, 0)
return x
@ -180,8 +181,15 @@ def test_allswap():
"""run_allswap"""
context.set_context(mode=context.GRAPH_MODE)
input_tensor = Tensor(np.ones((100, 20)), dtype=mindspore.float32)
label_tensor = Tensor(np.ones((1, 20)), dtype=mindspore.float32)
network = AllSwapNet(100, 20, 20)
_executor.compile(network, input_tensor)
loss_fn = nn.SoftmaxCrossEntropyWithLogits()
optimizer = Momentum(filter(lambda x: x.requires_grad, network.get_parameters()),
learning_rate=0.1,
momentum=0.9)
network = WithLossCell(network, loss_fn)
network = TrainOneStepCell(network, optimizer)
_executor.compile(network, input_tensor, label_tensor)
def run_reducescatter(op):