# coding=utf-
# rewritten, Copyright (c) 2021, Ming Ding. All rights reserved.
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# 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."""
import math
import copy
import torch
import torch.nn.functional as F
from apex.normalization.fused_layer_norm import FusedLayerNorm
from .initialize import get_model_parallel_world_size
from .layers import ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding
from .mappings import gather_from_model_parallel_region, copy_to_model_parallel_region
import deepspeed
from .random import checkpoint
from .random import get_cuda_rng_tracker
from .utils import divide, sqrt, scaled_init_method, unscaled_init_method, gelu
from .utils import split_tensor_along_last_dim
class LayerNorm(FusedLayerNorm):
def __init__(self, *args, pb_relax=False, **kwargs):
super().__init__(*args, **kwargs)
self.pb_relax = pb_relax
def forward(self, x):
if not self.pb_relax:
return super().forward(x)
return super().forward(x / (x.abs().max().detach()/8))
def standard_attention(query_layer, key_layer, value_layer, attention_mask,
attention_dropout=None, log_attention_weights=None):
# We disable the PB-relax-Attention and only changes the order of computation, because it is enough for most of training.
# The implementation in the paper can be done very easily, if you really need it to train very deep transformers.
attention_scores = torch.matmul(
query_layer / math.sqrt(query_layer.shape[-1]),
key_layer.transpose(-1, -2)
)
if attention_mask.shape[-2] > 1: # if auto-regressive, skip
attention_scores = torch.mul(attention_scores, attention_mask) - \
10000.0 * (1.0 - attention_mask)
if log_attention_weights is not None:
attention_scores += log_attention_weights
attention_probs = F.softmax(attention_scores, dim=-1)
if attention_dropout is not None:
with get_cuda_rng_tracker().fork():
attention_probs = attention_dropout(attention_probs)
context_layer = torch.matmul(attention_probs, value_layer)
return context_layer
class SelfAttention(torch.nn.Module):
def __init__(self, hidden_size, num_attention_heads,
attention_dropout_prob, output_dropout_prob,
init_method, layer_id, output_layer_init_method=None,
hooks={}):
super(SelfAttention, self).__init__()
# Set output layer initialization if not provided.
if output_layer_init_method is None:
output_layer_init_method = init_method
self.hooks = hooks
self.layer_id = layer_id
# Per attention head and per partition values.
world_size = get_model_parallel_world_size()
self.hidden_size_per_partition = divide(hidden_size, world_size)
self.hidden_size_per_attention_head = divide(hidden_size, num_attention_heads)
self.num_attention_heads_per_partition = divide(num_attention_heads, world_size)
# Strided linear layer.
self.query_key_value = ColumnParallelLinear(
hidden_size,
3*hidden_size,
stride=3,
gather_output=False,
init_method=init_method
)
self.attention_dropout = torch.nn.Dropout(attention_dropout_prob)
self.dense = RowParallelLinear(
hidden_size,
hidden_size,
input_is_parallel=True,
init_method=output_layer_init_method
)
self.output_dropout = torch.nn.Dropout(output_dropout_prob)
def _transpose_for_scores(self, tensor):
"""Transpose a 3D tensor [b, s, np*hn] into a 4D tensor with
size [b, np, s, hn].
"""
new_tensor_shape = tensor.size()[:-1] + \
(self.num_attention_heads_per_partition,
self.hidden_size_per_attention_head)
tensor = tensor.view(*new_tensor_shape)
return tensor.permute(0, 2, 1, 3)
def forward(self, hidden_states, mask, *other_tensors):
if 'attention_forward' in self.hooks:
return self.hooks['attention_forward'](hidden_states, mask, *other_tensors,layer_id=self.layer_id)
else:
mixed_raw_layer = self.query_key_value(hidden_states)
(mixed_query_layer,
mixed_key_layer,
mixed_value_layer) = split_tensor_along_last_dim(mixed_raw_layer, 3)
dropout_fn = self.attention_dropout if self.training else None
query_layer = self._transpose_for_scores(mixed_query_layer)
key_layer = self._transpose_for_scores(mixed_key_layer)
value_layer = self._transpose_for_scores(mixed_value_layer)
context_layer = standard_attention(query_layer, key_layer, value_layer, mask, dropout_fn)
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
context_layer = context_layer.view(*new_context_layer_shape)
output = self.dense(context_layer)
if self.training:
output = self.output_dropout(output)
return output, None
class MLP(torch.nn.Module):
def __init__(self, hidden_size, output_dropout_prob, init_method,
output_layer_init_method=None, hooks={}):
super(MLP, self).__init__()
# Set output layer initialization if not provided.
if output_layer_init_method is None:
output_layer_init_method = init_method
self.hooks = hooks
# Project to 4h.
self.dense_h_to_4h = ColumnParallelLinear(
hidden_size,
4*hidden_size,
gather_output=False,
init_method=init_method
)
# Project back to h.
self.dense_4h_to_h = RowParallelLinear(
4*hidden_size,
hidden_size,
input_is_parallel=True,
init_method=output_layer_init_method
)
self.dropout = torch.nn.Dropout(output_dropout_prob)
def forward(self, hidden_states, *other_tensors):
if 'mlp_forward' in self.hooks:
output = self.hooks['mlp_forward'](hidden_states, *other_tensors, layer_id=self.layer_id)
else:
intermediate_parallel = self.dense_h_to_4h(hidden_states)
intermediate_parallel = gelu(intermediate_parallel)
output = self.dense_4h_to_h(intermediate_parallel)
if self.training:
output = self.dropout(output)
return output
class BaseTransformerLayer(torch.nn.Module):
"""A single layer transformer for GPT2.
We use the following notation:
h: hidden size
n: number of attention heads
b: batch size
s: sequence length
Transformore layer takes input with size [b, s, h] and returns an
output of the same size.
Arguments:
hidden_size: The hidden size of the self attention.
num_attention_heads: number of attention head in the self
attention.
attention_dropout_prob: dropout probability of the attention
score in self attention.
output_dropout_prob: dropout probability for the outputs
after self attention and final output.
layernorm_epsilon: epsilon used in layernorm to avoid
division by zero.
init_method: initialization method used for the weights. Note
that all biases are initialized to zero and
layernorm weight are initialized to one.
output_layer_init_method: output layers (attention output and
mlp output) initialization. If None,
use `init_method`.
"""
def __init__(
self,
hidden_size,
num_attention_heads,
attention_dropout_prob,
output_dropout_prob,
layernorm_epsilon,
init_method,
layer_id,
output_layer_init_method=None,
sandwich_ln=True,
hooks={}
):
super(BaseTransformerLayer, self).__init__()
# Set output layer initialization if not provided.
if output_layer_init_method is None:
output_layer_init_method = init_method
self.layer_id = layer_id
self.hooks = hooks
# Layernorm on the input data.
self.input_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon)
# Self attention.
self.attention = SelfAttention(
hidden_size,
num_attention_heads,
attention_dropout_prob,
output_dropout_prob,
init_method,
layer_id,
output_layer_init_method=output_layer_init_method,
hooks=hooks
)
# Layernorm on the input data.
self.post_attention_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon)
self.sandwich_ln = sandwich_ln
if sandwich_ln:
self.third_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon)
self.fourth_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon)
# MLP
self.mlp = MLP(
hidden_size,
output_dropout_prob,
init_method,
output_layer_init_method=output_layer_init_method,
hooks=hooks
)
def forward(self, hidden_states, mask, *other_tensors):
'''
hidden_states: [batch, seq_len, hidden_size]
mask: [(1, 1), seq_len, seq_len]
'''
# Layer norm at the begining of the transformer layer.
layernorm_output1 = self.input_layernorm(hidden_states)
# Self attention.
attention_output, output_this_layer = self.attention(layernorm_output1, mask, *other_tensors)
# Third LayerNorm
if self.sandwich_ln:
attention_output = self.third_layernorm(attention_output)
# Residual connection.
layernorm_input = hidden_states + attention_output
# Layer norm post the self attention.
layernorm_output = self.post_attention_layernorm(layernorm_input)
# MLP.
mlp_output = self.mlp(layernorm_output)
# Fourth LayerNorm
if self.sandwich_ln:
mlp_output = self.fourth_layernorm(mlp_output, *other_tensors)
# Second residual connection.
output = layernorm_input + mlp_output
return output, output_this_layer # temporally, output_this_layer is only from attention
class BaseTransformer(torch.nn.Module):
def __init__(self,
num_layers,
vocab_size,
hidden_size,
num_attention_heads,
max_sequence_length,
embedding_dropout_prob,
attention_dropout_prob,
output_dropout_prob,
checkpoint_activations,
checkpoint_num_layers=1,
layernorm_epsilon=1.0e-5,
init_method_std=0.02,
sandwich_ln=True,
parallel_output=True,
hooks={}
):
super(BaseTransformer, self).__init__()
if deepspeed.checkpointing.is_configured():
global get_cuda_rng_tracker, checkpoint
get_cuda_rng_tracker = deepspeed.checkpointing.get_cuda_rng_tracker
checkpoint = deepspeed.checkpointing.checkpoint
# recording parameters
self.parallel_output = parallel_output
self.checkpoint_activations = checkpoint_activations
self.checkpoint_num_layers = checkpoint_num_layers
self.max_sequence_length = max_sequence_length
self.hooks = copy.copy(hooks) # hooks will be updated each forward
# create embedding parameters
self.embedding_dropout = torch.nn.Dropout(embedding_dropout_prob)
self.word_embeddings = VocabParallelEmbedding(
vocab_size, hidden_size, init_method=unscaled_init_method(0.02))
self.position_embeddings = torch.nn.Embedding(max_sequence_length, hidden_size)
torch.nn.init.normal_(self.position_embeddings.weight, mean=0.0, std=init_method_std)
# create all layers
self.output_layer_init_method = scaled_init_method(init_method_std, num_layers)
self.init_method = unscaled_init_method(init_method_std)
def get_layer(layer_id):
return BaseTransformerLayer(
hidden_size,
num_attention_heads,
attention_dropout_prob,
output_dropout_prob,
layernorm_epsilon,
self.init_method,
layer_id,
output_layer_init_method=self.output_layer_init_method,
sandwich_ln=sandwich_ln,
hooks=hooks
)
self.layers = torch.nn.ModuleList(
[get_layer(layer_id) for layer_id in range(num_layers)])
# Final layer norm before output.
self.final_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon)
def forward(self, input_ids, position_ids, attention_mask, *other_tensors):
# sanity check
assert len(input_ids.shape) == 2
batch_size, query_length = input_ids.shape
assert len(position_ids.shape) <= 2
assert position_ids.shape[-1] == query_length
assert len(attention_mask.shape) == 2 or \
len(attention_mask.shape) == 4 and attention_mask.shape[1] == 1
# embedding part
if 'word_embedding_forward' in self.hooks:
hidden_states = self.hooks['word_embedding_forward'](input_ids, *other_tensors)
else: # default
hidden_states = self.word_embeddings(input_ids)
if 'position_embedding_forward' in self.hooks:
position_embeddings = self.hooks['position_embedding_forward'](position_ids, *other_tensors)
else:
position_embeddings = self.position_embeddings(position_ids)
hidden_states = hidden_states + position_embeddings
hidden_states = self.embedding_dropout(hidden_states)
# define custom_forward for checkpointing
output_per_layers = []
if self.checkpoint_activations:
def custom(start, end):
def custom_forward(*inputs):
layers_ = self.layers[start:end]
x_, mask, *other_tensors = inputs[0], inputs[1], inputs[2:]
for i, layer in enumerate(layers_):
x_, output_this_layer = layer(x_, mask, *other_tensors)
output_per_layers.append(output_this_layer)
return x_
return custom_forward
l, num_layers = 0, len(self.layers)
chunk_length = self.checkpoint_num_layers
while l < num_layers:
args = [hidden_states, attention_mask, *other_tensors]
hidden_states = checkpoint(custom(l, l + chunk_length), *args)
l += chunk_length
else:
for i, layer in enumerate(self.layers):
args = [hidden_states, attention_mask, *other_tensors]
hidden_states, output_this_layer = layer(*args, *other_tensors)
output_per_layers.append(output_this_layer)
# Final layer norm.
logits = self.final_layernorm(hidden_states)
if 'final_forward' in self.hooks:
logits_parallel = self.hooks['final_forward'](logits, *other_tensors)
else:
logits_parallel = copy_to_model_parallel_region(logits)
logits_parallel = F.linear(logits_parallel, self.word_embeddings.weight)
if self.parallel_output:
return (logits_parallel, *output_per_layers)
return (gather_from_model_parallel_region(logits_parallel), *output_per_layers)