qsdy
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!25160 Fix Dropout Shard Interface 

Merge pull request !25160 from huangxinjing/dropout_v3
This commit is contained in:
i-robot 2021-10-20 07:27:38 +00:00 committed by Gitee
parent f0d3b4bbc7 75a704e101
commit fb756154fe
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mindspore/parallel/nn/transformer.py
View File

@ -370,6 +370,8 @@ class FeedForward(Cell):
self.projection.bias.parallel_optimizer = False
self.dropout = nn.Dropout(1 - dropout_rate)
self.dropout.dropout.shard(((dp, 1),))
self.dropout_3d = nn.Dropout(1 - dropout_rate)
self.dropout_3d.dropout.shard(((dp, 1, 1),))
self.cast = P.Cast()
def construct(self, x):
@ -380,7 +382,10 @@ class FeedForward(Cell):
hidden = self.mapping(x)
output = self.projection(hidden)
# returned shape is [bs, seq_length, ffn_hidden_size] or [bs * seq_length, ffn_hidden_size]
output = self.dropout(output)
if len(F.shape(output)) == 3:
output = self.dropout_3d(output)
else:
output = self.dropout(output)
return output