forked from huawei/mindspore2022
1. Move the class to mindspore.parallel, support activation sharding
This commit is contained in:
parent
d08e4ec03b
commit
d777742904
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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__)
|
||||
|
|
@ -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__)
|
||||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
||||
|
|
@ -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)`
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue