1. Move the class to mindspore.parallel, support activation sharding

This commit is contained in:
huangxinjing 2021-08-19 20:08:48 +08:00
parent d08e4ec03b
commit d777742904
9 changed files with 95 additions and 48 deletions

View File

@ -332,14 +332,15 @@ class LeakyReLU(Cell):
validator.check_value_type('alpha', alpha, [float, int], self.cls_name)
self.greater_equal = P.GreaterEqual()
self.mul = P.Mul()
self.maximum = P.Maximum()
self.alpha = alpha
def construct(self, x):
alpha_array = P.Cast()(F.scalar_to_array(self.alpha), P.DType()(x))
if self.alpha <= 1:
out = P.Maximum()(alpha_array * x, x)
out = self.maximum(alpha_array * x, x)
else:
out = P.Minimum()(alpha_array * x, x)
out = self.maximum(alpha_array * x, x)
return out

View File

@ -1,23 +0,0 @@
# 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.
# ============================================================================
"""
Transformer Networks
This is an experimental interface that is subject to change and/or deletion.
"""
from .transformer import *
from .layers import *
__all__ = []
__all__.extend(transformer.__all__)

View File

@ -13,14 +13,15 @@
# limitations under the License.
# ============================================================================
"""
Parallel Networks.
Transformer Networks
This is an experimental interface that is subject to change and/or deletion.
"""
from .transformer import *
from .layers import *
from .loss import *
from .config import *
from .op_parallel_config import *
__all__ = []
__all__.extend(transformer.__all__)
__all__.extend(loss.__all__)
__all__.extend(config.__all__)
__all__.extend(op_parallel_config.__all__)

View File

@ -217,6 +217,18 @@ class _Linear(Dense):
if self.has_bias:
self.bias_add.shard(strategy_bias)
if self.activation_flag:
getattr(self.activation, self.act_name).shard(strategy_activation)
# some operations has many primitives, need to manually set the shard
if self.act_name.lower() == "leakyrelu":
self.activation.maximum.shard((strategy_activation[0], strategy_activation[0]))
elif self.act_name.lower() == "logsigmoid":
self.activation.mul.shard((strategy_activation[0], ()))
self.activation.exp.shard(strategy_activation)
self.activation.add.shard((strategy_activation[0], ()))
self.activation.rec.shard(strategy_activation)
self.activation.log.shard(strategy_activation)
elif self.act_name.lower() == "logsoftmax":
raise ValueError("logsoftmax is not supported.")
else:
getattr(self.activation, self.act_name).shard(strategy_activation)
return self

View File

@ -22,8 +22,8 @@ from mindspore.ops import operations as P
from mindspore.ops import functional as F
from mindspore.nn import Cell
from mindspore.nn.loss.loss import _check_is_tensor
from mindspore.nn.parallel.transformer.transformer import _check_input_dtype, _check_input_shape
from .config import default_dpmp_config, OpParallelConfig
from mindspore.parallel.nn.transformer import _check_input_dtype, _check_input_shape
from .op_parallel_config import default_dpmp_config, OpParallelConfig
__all__ = ["CrossEntropyLoss"]

View File

@ -28,7 +28,7 @@ from mindspore._checkparam import Validator
from mindspore.ops.primitive import constexpr
from mindspore import log as logger
from .layers import _LayerNorm, _Linear
from ..config import default_dpmp_config, _PipeLineConfig, OpParallelConfig, _Config, _check_config
from .op_parallel_config import default_dpmp_config, _PipeLineConfig, OpParallelConfig, _Config, _check_config
__all__ = [
"AttentionMask",
@ -303,8 +303,9 @@ class FeedForward(Cell):
hidden_size (int): The dimension of the inputs.
ffn_hidden_size (int): The intermediate hidden size.
dropout_rate (float): The dropout rate for the second linear's output.
hidden_act (str): The activate type of the first linear. Support `gelu`, `relu`, `sigmpid` and so on.
Default: gelu.
hidden_act (str): The activation of the internal feedforward layer. Supports 'relu',
'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
'hsigmoid', 'logsigmoid' and so on. Default: gelu.
param_init_type (dtype.Number): The parameter initialization type. Can be dtype.float32 or dtype.float16.
parallel_config(OpParallelConfig): the config of parallel setting, see `OpParallelConfig`
Inputs:
@ -903,8 +904,9 @@ class TransformerEncoderLayer(Cell):
hidden_dropout_rate(float): The dropout rate of the final output of the layer. Default:0.1
attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1
post_layernorm_residual(bool): Do residuals adds before the layernorm. Default False.
hidden_act(str): The activation of the internal feedforward layer. Support `gelu`, `relu`, `sigmpid` and so on.
Default: gelu.
hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
'hsigmoid', 'logsigmoid' and so on. Default: gelu.
layernorm_compute_type(dtype.Number): The computation type of the layernorm.
Can be dtype.float32 or dtype.float16. Default dtype.float16.
softmax_comptue_type(dtype.Number): The computation type of the softmax in the attention.
@ -1100,8 +1102,9 @@ class TransformerDecoderLayer(Cell):
hidden_dropout_rate(float): The dropout rate of the final output of the layer. Default:0.1.
attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1.
post_layernorm_residual(bool): Do residuals adds before the layernorm. Default False.
hidden_act(str): The activation of the internal feedforward layer. Support `gelu`, `relu`, `sigmpid` and so on.
Default: gelu.
hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
'hsigmoid', 'logsigmoid' and so on. Default: gelu.
layernorm_compute_type(dtype.Number): The computation type of the layernorm.
Can be dtype.float32 or dtype.float16. Default dtype.float16.
softmax_comptue_type(dtype.Number): The computation type of the softmax in the attention.
@ -1390,8 +1393,9 @@ class TransformerEncoder(Cell):
hidden_dropout_rate(float): The dropout rate of the final output of the layer. Default:0.1
attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1
post_layernorm_residual(bool): Do residuals adds before the layernorm. Default False.
hidden_act(str): The activation of the internal feedforward layer. Support `gelu`, `relu`, `sigmpid` and so on.
Default: gelu.
hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
'hsigmoid', 'logsigmoid' and so on. Default: gelu.
layernorm_compute_type(dtype.Number): The computation type of the layernorm.
Can be dtype.float32 or dtype.float16. Default dtype.float16.
softmax_comptue_type(dtype.Number): The computation type of the softmax in the attention.
@ -1527,8 +1531,9 @@ class TransformerDecoder(Cell):
hidden_dropout_rate(float): The dropout rate of the final output of the layer. Default:0.1.
attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1.
post_layernorm_residual(bool): Do residuals adds before the layernorm. Default False.
hidden_act(str): The activation of the internal feedforward layer. Support `gelu`, `relu`, `sigmpid` and so on.
Default: gelu.
hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
'hsigmoid', 'logsigmoid' and so on. Default: gelu.
layernorm_compute_type(dtype.Number): The computation type of the layernorm.
Can be dtype.float32 or dtype.float16. Default dtype.float16.
softmax_comptue_type(dtype.Number): The computation type of the softmax in the attention.
@ -1668,8 +1673,8 @@ class TransformerDecoder(Cell):
class Transformer(Cell):
r"""
Transformer module. The difference is the module use the residual addition before the layernormalization. And the
default hidden act is `gelu`.
Transformer module including encoder and decoder. The difference with the original implements is the module use
the residual addition before the layernormalization. And the default hidden act is `gelu`.
The detials can be found in `Attention is all you need
<https://arxiv.org/pdf/1706.03762v5.pdf>`.
@ -1689,8 +1694,9 @@ class Transformer(Cell):
hidden_dropout_rate(float): The dropout rate of the final output of the layer. Default:0.1
attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1
post_layernorm_residual(bool): Do residuals adds before the layernorm. Default False.
hidden_act(str): The activation of the internal feedforward layer. Support `gelu`, `relu`, `sigmpid` and so on.
Default: gelu.
hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
'hsigmoid', 'logsigmoid' and so on. Default: gelu.
lambda_func: A function can specific the fusion index, pipeline stages and recompute attribute. If the user
wants to specific the pipeline stage and gradient aggregation fusion, the user can pass a function
that accepts `network`, `layer_id`, `offset`, `parallel_config`, `layers`. The `network(Cell)`

View File

@ -14,9 +14,10 @@
# ============================================================================
""" test transformer"""
import numpy as np
import pytest
from mindspore import Tensor
from mindspore.common import dtype
from mindspore.nn.parallel import MultiHeadAttention, FeedForward, TransformerEncoderLayer, TransformerEncoder, \
from mindspore.parallel.nn import MultiHeadAttention, FeedForward, TransformerEncoderLayer, TransformerEncoder, \
TransformerDecoder, TransformerDecoderLayer, Transformer, CrossEntropyLoss, AttentionMask
from mindspore.common.api import _executor
@ -36,6 +37,55 @@ def test_transformer_encoder_only():
_executor.compile(model, encoder_input_value, encoder_input_mask)
def test_transformer_encoder_log_softmax():
with pytest.raises(ValueError):
model = Transformer(batch_size=2,
src_seq_length=20,
tgt_seq_length=0,
encoder_layers=2,
decoder_layers=0,
hidden_act='logsoftmax',
hidden_size=64,
ffn_hidden_size=64)
encoder_input_value = Tensor(np.ones((2, 20, 64)), dtype.float32)
encoder_input_mask = Tensor(np.ones((2, 20, 20)), dtype.float16)
_executor.compile(model, encoder_input_value, encoder_input_mask)
def test_transformer_encoder_leakyrelu():
model = Transformer(batch_size=2,
src_seq_length=20,
tgt_seq_length=0,
encoder_layers=2,
decoder_layers=0,
hidden_act='leakyrelu',
hidden_size=64,
ffn_hidden_size=64)
encoder_input_value = Tensor(np.ones((2, 20, 64)), dtype.float32)
encoder_input_mask = Tensor(np.ones((2, 20, 20)), dtype.float16)
_executor.compile(model, encoder_input_value, encoder_input_mask)
def test_transformer_encoder_logsigmoid():
model = Transformer(batch_size=2,
src_seq_length=20,
tgt_seq_length=0,
encoder_layers=2,
decoder_layers=0,
hidden_act='logsigmoid',
hidden_size=64,
ffn_hidden_size=64)
encoder_input_value = Tensor(np.ones((2, 20, 64)), dtype.float32)
encoder_input_mask = Tensor(np.ones((2, 20, 20)), dtype.float16)
_executor.compile(model, encoder_input_value, encoder_input_mask)
def test_encoder_and_decoder():
model = Transformer(batch_size=2,
src_seq_length=20,

View File

@ -21,7 +21,7 @@ from mindspore.context import set_auto_parallel_context, ParallelMode
from mindspore.ops import composite as C
from mindspore.ops import functional as F
import mindspore.ops as P
from mindspore.nn.parallel import TransformerEncoder, TransformerDecoder, Transformer, TransformerOpParallelConfig, \
from mindspore.parallel.nn import TransformerEncoder, TransformerDecoder, Transformer, TransformerOpParallelConfig, \
VocabEmbedding, CrossEntropyLoss, OpParallelConfig, EmbeddingOpParallelConfig
from mindspore.nn import Dense as Linear
from mindspore.nn.wrap.loss_scale import DynamicLossScaleUpdateCell