# Copyright (c) 2020 PaddlePaddle Authors. 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. import collections import copy import math import numpy as np import numbers import json import sys import os import paddle import paddle.nn as nn import paddle.tensor as tensor import paddle.nn.functional as F from paddle.nn import TransformerEncoder, Linear, Layer, Embedding, LayerNorm, Tanh from paddle.nn import Layer, LayerList from paddle.fluid.initializer import Constant import utility from bert_padding import generate_mask #try: # from custom_setup_ops import custom_fmha, custom_fused_dropout_residual_ln, custom_fused_dense #except ImportError as e: # print('custom_setup_ops import error: {}'.format(e)) from .load_tf_checkpoint import load_pickled_tf_checkpoint, save_pickled_tf_checkpoint from .mlperf_logging_helper import paddle_bert_print_event __all__ = [ 'BertConfig', 'BertModel', 'BertForPretraining', 'BertPretrainingCriterion', 'BertPretrainingHeads', ] use_nv_input = utility.use_nv_input() GELU_APPROXIMATE = True def get_activation(): return nn.GELU(approximate=GELU_APPROXIMATE) # return nn.ReLU() def mask_gather(var, mask): return paddle.gather_nd(var, paddle.fluid.layers.where(mask)) def gen_pos_id(input_ids): ones = paddle.ones_like(input_ids) seq_length = paddle.cumsum(ones, axis=-1) position_ids = seq_length - ones position_ids.stop_gradient = True return position_ids def fuse_dense(x, y, bias, transx=False, transy=False, with_gelu=False, use_addto=False): #out = custom_fused_dense( # x=x, y=y, bias=bias, transx=transx, transy=transy, use_addto=use_addto) # paddle.static.Print(x, message="print matmul x:") # paddle.static.Print(y, message="print matmul y:") #paddle.static.Print(transx, message="print matmul transx:") #paddle.static.Print(transy, message="print matmul transy:") # print("print matmul transx: ", transx) # print("print matmul transy", transy) # print("B.shape", bias.shape) # os.environ["ROCBLAS_LAYER"] = "3" # setenv_status = int(os.environ.get('ROCBLAS_LAYER', '3')) out = paddle.matmul(x, y, transpose_x=transx, transpose_y=transy) # os.environ["ROCBLAS_LAYER"] = "0" # paddle.static.Print(out, message="print matmul out:") # paddle.static.Print(bias, message="print matmul bias:") #paddle.static.Print(out, message="print matmul out before bias:") if bias is not None: out = out + bias # paddle.static.Print(out, message="print matmul out after bias:") if with_gelu: out = get_activation()(out) # paddle.static.Print(out, message="print matmul out after bias:") return out class BertConfig(object): """Configuration class to store the configuration of a `BertModel`. """ def __init__(self, vocab_size_or_config_json_file, **kwargs): """Constructs BertConfig. Args: vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`. hidden_size: Size of the encoder layers and the pooler layer. num_hidden_layers: Number of hidden layers in the Transformer encoder. num_attention_heads: Number of attention heads for each attention layer in the Transformer encoder. intermediate_size: The size of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. hidden_act: The non-linear activation function (function or string) in the encoder and pooler. If string, "gelu", "relu" and "swish" are supported. hidden_dropout_prob: The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. attention_probs_dropout_prob: The dropout ratio for the attention probabilities. max_position_embeddings: The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048). type_vocab_size: The vocabulary size of the `token_type_ids` passed into `BertModel`. initializer_range: The sttdev of the truncated_normal_initializer for initializing all weight matrices. """ if isinstance(vocab_size_or_config_json_file, str) or ( sys.version_info[0] == 2 and isinstance(vocab_size_or_config_json_file, unicode)): with open( vocab_size_or_config_json_file, "r", encoding='utf-8') as reader: json_config = json.loads(reader.read()) assert not kwargs, "kwargs should be empty if config json file is provided" self._fill_dict(json_config) elif isinstance(vocab_size_or_config_json_file, int): self.vocab_size = vocab_size_or_config_json_file self._fill_dict(kwargs) else: raise ValueError( "First argument must be either a vocabulary size (int)" "or the path to a pretrained model config file (str)") @classmethod def from_dict(cls, json_object): """Constructs a `BertConfig` from a Python dictionary of parameters.""" config = BertConfig(vocab_size_or_config_json_file=-1) config._fill_dict(json_object) return config @classmethod def from_json_file(cls, json_file): """Constructs a `BertConfig` from a json file of parameters.""" with open(json_file, "r", encoding='utf-8') as reader: text = reader.read() return cls.from_dict(json.loads(text)) def __repr__(self): return str(self.to_json_string()) def to_dict(self): """Serializes this instance to a Python dictionary.""" output = copy.deepcopy(self.__dict__) return output def to_json_string(self): """Serializes this instance to a JSON string.""" return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n" def _fill_dict(self, kwargs=None): defaults = { "vocab_size": 30522, "hidden_size": 768, "num_hidden_layers": 12, "num_attention_heads": 12, "intermediate_size": 3072, "hidden_act": "gelu", "hidden_dropout_prob": 0.1, "attention_probs_dropout_prob": 0.1, "max_position_embeddings": 512, "type_vocab_size": 2, "initializer_range": 0.02, "pool_act": "tanh", "pad_token_id": 0, } # fill defaults for key, value in defaults.items(): if key not in self.__dict__: self.__dict__[key] = value # fill other values if kwargs: for key, value in kwargs.items(): self.__dict__[key] = value assert self.pad_token_id == 0, "pad_token_id must be 0" def transpose_2d(x): assert len(x.shape) == 2 return np.ascontiguousarray(np.transpose(x, (1, 0))) class TFCkptHelper: def __init__(self, args, config, checkpoint_path, place): self.pd_vars_to_tf_vars = collections.OrderedDict() self.transpose_vars = set() self.args = args self.config = config self.place = place self.checkpoint_path = checkpoint_path self._tf_vars = None self.fuse_attn_qkv = self.args.unpad_fmha self.attn_fused_qkv_weights = [ None for _ in range(self.config.num_hidden_layers) ] self.attn_fused_qkv_biases = [ None for _ in range(self.config.num_hidden_layers) ] @property def tf_vars(self): if self._tf_vars is None: self._tf_vars = load_pickled_tf_checkpoint(self.checkpoint_path) return self._tf_vars def save(self, output_path, get_parameter_func): tf_vars = {} if self.fuse_attn_qkv: for idx in range(self.config.num_hidden_layers): prefix = self._enc_prefix(idx) + "attention/self/" pd_weight_name = self.attn_fused_qkv_weights[idx] pd_weight_var = self._get_fp32_param(pd_weight_name, get_parameter_func) assert len(pd_weight_var.shape) == 2 pd_bias_name = self.attn_fused_qkv_biases[idx] pd_bias_var = self._get_fp32_param(pd_bias_name, get_parameter_func) assert len(pd_bias_var.shape) == 1 need_transpose = pd_weight_name in self.transpose_vars if need_transpose: # split along dim 0 and transpose assert pd_weight_var.shape[0] == 3 * pd_weight_var.shape[1] assert pd_weight_var.shape[0] == pd_bias_var.shape[0] tf_weights = np.split(pd_weight_var, 3, axis=0) tf_weights = [transpose_2d(w) for w in tf_weights] else: # split along dim 1 assert pd_weight_var.shape[1] == 3 * pd_weight_var.shape[0] assert pd_weight_var.shape[1] == pd_bias_var.shape[0] tf_weights = np.split(pd_weight_var, 3, axis=1) tf_biases = np.split(pd_bias_var, 3, axis=0) assert len(tf_weights) == 3 assert len(tf_biases) == 3 for i, name in enumerate(["query", "key", "value"]): tf_var_name = prefix + name + "/" weight_name = tf_var_name + "kernel" bias_name = tf_var_name + "bias" tf_vars[weight_name] = tf_weights[i] tf_vars[bias_name] = tf_biases[i] for pd_var_name, tf_var_name in self.pd_vars_to_tf_vars.items(): pd_var = self._get_fp32_param(pd_var_name, get_parameter_func) if "output_weights" in tf_var_name or pd_var_name in self.transpose_vars: pd_var = transpose_2d(pd_var) tf_vars[tf_var_name] = pd_var for key, tf_value in self.tf_vars.items(): if key not in tf_vars: tf_vars[key] = tf_value continue pd_value = tf_vars[key] if tf_value.shape == pd_value.shape: continue if key == 'bert/embeddings/word_embeddings': assert len(tf_value.shape) == 2 assert len(pd_value.shape) == 2 assert tf_value.shape[1] == pd_value.shape[1] pd_value = pd_value[0:tf_value.shape[0]] elif key == 'cls/predictions/output_bias': assert len(tf_value.shape) == 1 assert len(pd_value.shape) == 1 pd_value = pd_value[0:tf_value.shape[0]] else: raise ValueError("unsupported key {}".format(key)) tf_vars[key] = pd_value return save_pickled_tf_checkpoint(tf_vars, output_path) def load(self, get_parameter_func): tf_vars = self.tf_vars loaded_var_names = set() if self.fuse_attn_qkv: for idx in range(self.config.num_hidden_layers): weights = [] biases = [] prefix = self._enc_prefix(idx) + "attention/self/" pd_weight_name = self.attn_fused_qkv_weights[idx] pd_bias_name = self.attn_fused_qkv_biases[idx] need_transpose = pd_weight_name in self.transpose_vars for name in ["query", "key", "value"]: tf_var_name = prefix + name + "/" weight_name = tf_var_name + "kernel" bias_name = tf_var_name + "bias" if utility.get_trainer_id() == 0: paddle_bert_print_event(key='weights_initialization', metadata={'tensor':weight_name}) paddle_bert_print_event(key='weights_initialization', metadata={'tensor':bias_name}) if need_transpose: weights.append(transpose_2d(tf_vars[weight_name])) else: weights.append(tf_vars[weight_name]) biases.append(tf_vars[bias_name]) loaded_var_names.add(weight_name) loaded_var_names.add(bias_name) weight = np.concatenate( weights, axis=0 if need_transpose else 1) bias = np.concatenate(biases) weight_pd_vars = get_parameter_func(pd_weight_name) bias_pd_vars = get_parameter_func(pd_bias_name) self._set_var_value( weight_pd_vars, weight, self.attn_fused_qkv_weights[idx] + "/qkv/kernel", prefix + "qkv/kernel") self._set_var_value( bias_pd_vars, bias, self.attn_fused_qkv_biases[idx] + "/qkv/bias", prefix + "qkv/bias") for idx, (pd_var_name, tf_var_name) in enumerate(self.pd_vars_to_tf_vars.items()): if utility.get_trainer_id() == 0: paddle_bert_print_event(key='weights_initialization', metadata={'tensor':tf_var_name}) var_value = tf_vars[tf_var_name] if "output_weights" in tf_var_name or pd_var_name in self.transpose_vars: if utility.get_trainer_id() == 0: print('{} needs to transpose'.format(tf_var_name)) var_value = transpose_2d(var_value) pd_vars = get_parameter_func(pd_var_name) self._set_var_value(pd_vars, var_value, pd_var_name, tf_var_name) loaded_var_names.add(tf_var_name) left_var_names = set() for var_name in self.pd_vars_to_tf_vars.values(): if var_name not in loaded_var_names: left_var_names.add(var_name) else: loaded_var_names.remove(var_name) if self.fuse_attn_qkv: assert len(loaded_var_names ) == 6 * self.config.num_hidden_layers, loaded_var_names else: assert len(loaded_var_names) == 0, loaded_var_names assert len(left_var_names) == 0, left_var_names def _set_var_value(self, pd_vars, var_value, pd_var_name, tf_var_name): if isinstance(pd_vars, (list, tuple)): assert len(pd_vars) == 2 pd_var, master_pd_var = pd_vars else: pd_var = pd_vars master_pd_var = None pd_var_shape = tuple(pd_var.shape()) tf_var_shape = tuple(var_value.shape) if pd_var_shape != tf_var_shape: if utility.get_trainer_id() == 0: print("{} vs {} shape differs: {} vs {}".format( pd_var_name, tf_var_name, pd_var_shape, tf_var_shape)) assert len(pd_var_shape) == len(tf_var_shape) slices = [] n = len(pd_var_shape) for i in range(n): assert pd_var_shape[i] >= tf_var_shape[i] slices.append(slice(0, tf_var_shape[i], 1)) new_var_value = np.zeros(pd_var_shape, dtype=var_value.dtype) new_var_value[slices] = var_value var_value = new_var_value if pd_var._dtype() == paddle.float16: assert var_value.dtype == np.float32 if master_pd_var is not None: assert master_pd_var._dtype() == paddle.float32 if utility.get_trainer_id() == 0: print("Set master weight for {} {}".format(pd_var_name, tf_var_name)) self._inplace_set_tensor(master_pd_var, var_value) self._inplace_set_tensor(pd_var, var_value.astype(np.float16)) elif pd_var._dtype() == paddle.float32: assert var_value.dtype == np.float32 assert master_pd_var is None self._inplace_set_tensor(pd_var, var_value) else: raise TypeError("unsupported data type {}".format(pd_var._dtype())) def _inplace_set_tensor(self, tensor, value): old_ptr = tensor._ptr() tensor.set(value, self.place) new_ptr = tensor._ptr() assert old_ptr == new_ptr def _get_fp32_param(self, pd_var_name, get_parameter_func): pd_var_name = self._to_pd_var_name(pd_var_name) pd_vars = get_parameter_func(pd_var_name) assert isinstance(pd_vars, (list, tuple)) assert len(pd_vars) == 2 pd_var, master_pd_var = pd_vars if master_pd_var is not None: assert pd_var._dtype() == paddle.float16 assert master_pd_var._dtype() == paddle.float32 assert pd_var.shape() == master_pd_var.shape() return np.array(master_pd_var) else: assert pd_var._dtype() == paddle.float32 return np.array(pd_var) def _enc_prefix(self, idx): return "bert/encoder/layer_{}/".format(idx) def _to_pd_var_name(self, var): return var if isinstance(var, (str, bytes)) else var.name def _record_pd_vars(self, pd_vars, tf_vars, tf_var_prefix="", weight_transpose=None): if not isinstance(pd_vars, (list, tuple)): pd_vars = [pd_vars] pd_vars = [self._to_pd_var_name(v) for v in pd_vars] if not isinstance(tf_vars, (list, tuple)): tf_vars = [tf_vars] tf_vars = [tf_var_prefix + v for v in tf_vars] assert len(pd_vars) == len(tf_vars) for pd_var, tf_var in zip(pd_vars, tf_vars): assert pd_var not in self.pd_vars_to_tf_vars, pd_var self.pd_vars_to_tf_vars[pd_var] = tf_var if weight_transpose: assert len(pd_vars) == 2 self.transpose_vars.add(pd_vars[0]) def embeddings(self, pd_vars): return self._record_pd_vars(pd_vars, [ "word_embeddings", "position_embeddings", "token_type_embeddings" ], "bert/embeddings/") def norm_after_embeddings(self, pd_vars): return self._record_pd_vars(pd_vars, ["gamma", "beta"], "bert/embeddings/LayerNorm/") def enc_attn_query_fc(self, pd_vars, idx, weight_transpose=None): assert not self.fuse_attn_qkv prefix = self._enc_prefix(idx) + "attention/self/query/" return self._record_pd_vars(pd_vars, ["kernel", "bias"], prefix, weight_transpose) def enc_attn_key_fc(self, pd_vars, idx, weight_transpose=False): assert not self.fuse_attn_qkv prefix = self._enc_prefix(idx) + "attention/self/key/" return self._record_pd_vars(pd_vars, ["kernel", "bias"], prefix, weight_transpose) def enc_attn_value_fc(self, pd_vars, idx, weight_transpose=False): assert not self.fuse_attn_qkv prefix = self._enc_prefix(idx) + "attention/self/value/" return self._record_pd_vars(pd_vars, ["kernel", "bias"], prefix, weight_transpose) def enc_fused_attn_qkv_fc(self, pd_vars, idx, weight_transpose=False): assert self.fuse_attn_qkv assert self.attn_fused_qkv_weights[idx] is None assert self.attn_fused_qkv_biases[idx] is None weight, bias = pd_vars self.attn_fused_qkv_weights[idx] = self._to_pd_var_name(weight) self.attn_fused_qkv_biases[idx] = self._to_pd_var_name(bias) if weight_transpose: self.transpose_vars.add(self.attn_fused_qkv_weights[idx]) def enc_attn_proj_fc(self, pd_vars, idx, weight_transpose=False): prefix = self._enc_prefix(idx) + "attention/output/dense/" return self._record_pd_vars(pd_vars, ["kernel", "bias"], prefix, weight_transpose) def enc_attn_norm(self, pd_vars, idx): prefix = self._enc_prefix(idx) + "attention/output/LayerNorm/" return self._record_pd_vars(pd_vars, ["gamma", "beta"], prefix) def enc_intermediate_fc(self, pd_vars, idx, weight_transpose=False): prefix = self._enc_prefix(idx) + "intermediate/dense/" return self._record_pd_vars(pd_vars, ["kernel", "bias"], prefix, weight_transpose) def enc_output_fc(self, pd_vars, idx, weight_transpose=False): prefix = self._enc_prefix(idx) + "output/dense/" return self._record_pd_vars(pd_vars, ["kernel", "bias"], prefix, weight_transpose) def enc_output_norm(self, pd_vars, idx): prefix = self._enc_prefix(idx) + "output/LayerNorm/" return self._record_pd_vars(pd_vars, ["gamma", "beta"], prefix) def pooler_fc(self, pd_vars, weight_transpose=False): return self._record_pd_vars(pd_vars, ["kernel", "bias"], "bert/pooler/dense/", weight_transpose) def cls_pred_trans_fc(self, pd_vars, weight_transpose=False): return self._record_pd_vars(pd_vars, ["kernel", "bias"], "cls/predictions/transform/dense/", weight_transpose) def cls_pred_trans_norm(self, pd_vars): return self._record_pd_vars(pd_vars, ["gamma", "beta"], "cls/predictions/transform/LayerNorm/") def cls_pred_fc_bias(self, pd_vars): return self._record_pd_vars(pd_vars, "output_bias", "cls/predictions/") def cls_seq_relation_fc(self, pd_vars, weight_transpose=False): return self._record_pd_vars(pd_vars, ["output_weights", "output_bias"], "cls/seq_relationship/", weight_transpose) class FMHA(Layer): def __init__(self, config): super(FMHA, self).__init__() self.p_dropout = config.attention_probs_dropout_prob self.h = config.num_attention_heads self.hidden_size = config.hidden_size self.d = self.hidden_size // self.h self.fused_qkv_bias = config.fused_bias_mha self.weight_transpose = True self.use_unpad_fmha_mke_opt = config.unpad_fmha_mke_opt assert self.d * self.h == self.hidden_size, "Invalid hidden size/num_heads" # create_parameter self._dtype = self._helper.get_default_dtype() if self.weight_transpose: Wqkv_shape = [3 * config.hidden_size, config.hidden_size] else: Wqkv_shape = [config.hidden_size, 3 * config.hidden_size] self.Wqkv = self.create_parameter( shape=Wqkv_shape, attr=None, dtype=self._dtype, is_bias=False) self.Bqkv = self.create_parameter( shape=[3 * config.hidden_size], attr=None, dtype=self._dtype, is_bias=True) def forward(self, hidden_states, cu_seqlens, host_cu_seqlens, max_s, is_training=True): #paddle.static.Print(hidden_states, message="fmha layer, input shape") #paddle.static.Print(self.Wqkv, message="Wqkv.shape") #paddle.static.Print(self.Bqkv, message="Bqkv.shape") #print("fmha layer, input shape: ", hidden_states.shape) #print("Wqkv.shape", self.Wqkv.shape) #print("Bqkv.shape", self.Bqkv.shape) if not self.fused_qkv_bias: qkv = paddle.matmul( hidden_states, self.Wqkv, transpose_x=False, transpose_y=self.weight_transpose) qkv = qkv + self.Bqkv else: qkv = fuse_dense( hidden_states, self.Wqkv, self.Bqkv, transx=False, transy=self.weight_transpose) qkv = paddle.reshape(qkv, [-1, 3, self.h, self.d]) #print("qkv.shape", qkv.shape) # FMHA: max_s = var memcpy_0.tmp_0 : LOD_TENSOR.shape(1,).dtype(int32).stop_gradient(False) #print("FMHA: max_s = ", max_s) out, _ = custom_fmha( qkv, cu_seqlens, host_cu_seqlens, not is_training, self.p_dropout, zero_tensors=False, use_fmha_mke_opt=self.use_unpad_fmha_mke_opt) return paddle.reshape(out, [-1, self.hidden_size]) class BertSelfAttention(Layer): def __init__(self, config): super(BertSelfAttention, self).__init__() assert config.hidden_size % config.num_attention_heads == 0 self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = nn.Linear(config.hidden_size, self.all_head_size) self.key = nn.Linear(config.hidden_size, self.all_head_size) self.value = nn.Linear(config.hidden_size, self.all_head_size) self.weight_transpose = False self.dropout = nn.Dropout(config.attention_probs_dropout_prob) self.softmax = nn.Softmax(axis=-1) def transpose_for_scores(self, x): new_x_shape = [0] * ( len(x.shape) - 1 ) + [self.num_attention_heads, self.attention_head_size] x = paddle.reshape(x, new_x_shape) return paddle.transpose(x, [0, 2, 1, 3]) def transpose_key_for_scores(self, x): new_x_shape = [0] * ( len(x.shape) - 1 ) + [self.num_attention_heads, self.attention_head_size] x = paddle.reshape(x, new_x_shape) return paddle.transpose(x, [0, 2, 3, 1]) def forward(self, hidden_states, attention_mask): mixed_query_layer = self.query(hidden_states) mixed_key_layer = self.key(hidden_states) mixed_value_layer = self.value(hidden_states) query_layer = self.transpose_for_scores(mixed_query_layer) key_layer = self.transpose_key_for_scores(mixed_key_layer) value_layer = self.transpose_for_scores(mixed_value_layer) # paddle.static.Print(query_layer, message="query_layer") # paddle.static.Print(key_layer, message="key_layer") attention_scores = paddle.matmul(query_layer, key_layer) attention_scores = attention_scores / math.sqrt( self.attention_head_size) attention_scores = attention_scores + attention_mask.unsqueeze( 1).unsqueeze(2) attention_probs = self.softmax(attention_scores) attention_probs = self.dropout(attention_probs) # paddle.static.Print(attention_probs, message="attention_probs") # paddle.static.Print(value_layer, message="value_layer") context_layer = paddle.matmul(attention_probs, value_layer) #paddle.static.Print(context_layer, message="context_layer") context_layer = paddle.transpose(context_layer, [0, 2, 1, 3]) new_context_layer_shape = [0] * (len(context_layer.shape) - 2 ) + [self.all_head_size] context_layer = paddle.reshape(context_layer, new_context_layer_shape) return context_layer class BertAttention(Layer): def __init__(self, config): super(BertAttention, self).__init__() self.self = BertSelfAttention(config) self.output = BertSelfOutput(config) def forward(self, input_tensor, attention_mask): self_output = self.self(input_tensor, attention_mask) attention_output = self.output(self_output, input_tensor) return attention_output class BertLayer(Layer): def __init__(self, config): super(BertLayer, self).__init__() assert use_nv_input assert not config.pad_fmha self.unpad = config.unpad if config.unpad_fmha: assert self.unpad self.attention = UnpadFMHABertAttention(config) else: self.attention = BertAttention(config) self.intermediate = BertIntermediate(config) self.output = BertOutput(config) def forward(self, hidden_states, attention_mask, seqlen=None, host_seqlen=None, batch=None): if self.unpad: attention_output = self.attention(hidden_states, attention_mask, seqlen, host_seqlen, batch) else: attention_output = self.attention(hidden_states, attention_mask) intermediate_output = self.intermediate(attention_output) layer_output = self.output(intermediate_output, attention_output) return layer_output class BertEncoder(Layer): def __init__(self, config): super(BertEncoder, self).__init__() if use_nv_input: self.layers = nn.LayerList( [BertLayer(config) for _ in range(config.num_hidden_layers)]) else: encoder_layer = nn.TransformerEncoderLayer( config.hidden_size, config.num_attention_heads, config.intermediate_size, dropout=config.hidden_dropout_prob, activation=config.hidden_act, attn_dropout=config.attention_probs_dropout_prob, act_dropout=0) self.encoder = nn.TransformerEncoder(encoder_layer, config.num_hidden_layers) self.num_attention_heads = config.num_attention_heads self.unpad = config.unpad self.unpad_embed = config.unpad_embed self.unpad_fmha = config.unpad_fmha self.pad_fmha = config.pad_fmha self.hidden_size = config.hidden_size self.maxseqlen = config.max_seq_length def record_ckpt_vars(self, ckpt, idx): if use_nv_input: layer = self.layers[idx] attn = layer.attention if isinstance(attn, UnpadFMHABertAttention): ckpt.enc_fused_attn_qkv_fc([attn.fmha.Wqkv, attn.fmha.Bqkv], idx, attn.fmha.weight_transpose) else: ckpt.enc_attn_query_fc( [attn.self.query.weight, attn.self.query.bias], idx, attn.self.weight_transpose) ckpt.enc_attn_key_fc( [attn.self.key.weight, attn.self.key.bias], idx, attn.self.weight_transpose) ckpt.enc_attn_value_fc( [attn.self.value.weight, attn.self.value.bias], idx, attn.self.weight_transpose) ckpt.enc_attn_proj_fc( [attn.output.dense.weight, attn.output.dense.bias], idx, attn.output.weight_transpose) if attn.output.fused_dropout: ckpt.enc_attn_norm([ attn.output.fused_dropout_add_ln.weight, attn.output.fused_dropout_add_ln.bias ], idx) else: ckpt.enc_attn_norm([ attn.output.layer_norm.weight, attn.output.layer_norm.bias ], idx) intermediate = layer.intermediate last_output = layer.output ckpt.enc_intermediate_fc( [intermediate.dense.weight, intermediate.dense.bias], idx, intermediate.weight_transpose) ckpt.enc_output_fc( [last_output.dense.weight, last_output.dense.bias], idx, last_output.weight_transpose) if last_output.fused_dropout: ckpt.enc_output_norm([ last_output.fused_dropout_add_ln.weight, last_output.fused_dropout_add_ln.bias ], idx) else: ckpt.enc_output_norm([ last_output.layer_norm.weight, last_output.layer_norm.bias ], idx) else: layer = self.encoder.layers[idx] attn = layer.self_attn ckpt.enc_attn_query_fc([attn.q_proj.weight, attn.q_proj.bias], idx, False) ckpt.enc_attn_key_fc([attn.k_proj.weight, attn.k_proj.bias], idx, False) ckpt.enc_attn_value_fc([attn.v_proj.weight, attn.v_proj.bias], idx, False) ckpt.enc_attn_proj_fc([attn.out_proj.weight, attn.out_proj.bias], idx, False) ckpt.enc_attn_norm([layer.norm1.weight, layer.norm1.bias], idx) ckpt.enc_intermediate_fc( [layer.linear1.weight, layer.linear1.bias], idx, False) ckpt.enc_output_fc([layer.linear2.weight, layer.linear2.bias], idx, False) ckpt.enc_output_norm([layer.norm2.weight, layer.norm2.bias], idx) def forward(self, hidden_states, attention_mask, output_all_encoded_layers=False, batch=56, maxseqlen=512, hidden_size=1024, zero_tensor=None, attention_indices=None, new_attention_mask=None, seqlen=None, cu_seqlens=None, host_cu_seqlens=None, actual_seqlens=None, maxseqlen_in_batch=None): if use_nv_input: return self.forward_with_nv_input( hidden_states, attention_mask, output_all_encoded_layers, batch, maxseqlen, hidden_size, zero_tensor, attention_indices, new_attention_mask, seqlen, cu_seqlens, host_cu_seqlens, actual_seqlens, maxseqlen_in_batch) if output_all_encoded_layers: output = hidden_states encoder_outputs = [] for mod in self.encoder.layers: output = mod(output, src_mask=attention_mask) encoder_outputs.append(output) if self.encoder.norm is not None: encoder_outputs[-1] = self.encoder.norm(encoder_outputs[-1]) return encoder_outputs else: sequence_output = self.encoder(hidden_states, attention_mask) return [sequence_output] def forward_with_nv_input(self, hidden_states, attention_mask, output_all_encoded_layers=False, batch=56, maxseqlen=512, hidden_size=1024, zero_tensor=None, attention_indices=None, new_attention_mask=None, seqlen=None, cu_seqlens=None, host_cu_seqlens=None, actual_seqlens=None, maxseqlen_in_batch=None): # Unpad inputs and mask. It will remove tokens that are padded. Assume ntokens is total number of tokens (padded and non-padded) # and ntokens_unpad is total number of non-padded tokens. Then unpadding performs the following compression of the inputs: # hidden_states[ntokens,hidden] -> hidden_states[ntokens_unpad,hidden] if not self.unpad_embed and self.unpad_fmha: #batch = None #seqlen = None if self.unpad_fmha: batch = hidden_states.shape[0] maxseqlen = hidden_states.shape[1] assert maxseqlen == self.maxseqlen hidden_size = hidden_states.shape[2] assert hidden_size == self.hidden_size zero_tensor = paddle.zeros_like(hidden_states) zero_tensor = paddle.reshape(zero_tensor, [-1, self.hidden_size]) # attention_indices: 把attn_mask flatten后，提取非零元的下标。 # seqlen： 对[bs, max_seq_len]，每一行求和，代表获取每一行的实际seq_len（一维）。 # ntokens: 代表当前batch的所有实际seq_len之和。 # cu_seqlens: 对seq_len求prefix_sum的结果。 # actual_seqlens: 与seqlen相同的值。 # maxseqlen_in_batch: the max seqlen in a batch. #attention_indices, attention_mask, seqlen, ntokens, cu_seqlens, actual_seqlens, maxseqlen_in_batch = generate_mask( # attention_mask, unpad_fmha=self.unpad_fmha) print("maxseqlen_in_batch = ", maxseqlen_in_batch) hidden_states = paddle.reshape(hidden_states, [-1, self.hidden_size]) hidden_states = paddle.gather(hidden_states, attention_indices) # print("unpad after shape: ", hidden_states) elif self.unpad_fmha: attention_mask = new_attention_mask all_encoder_layers = [] def custom(start, end): def custom_forward(*inputs): layers = self.layers[start:end] x_ = inputs[0] for layer in layers: x_ = layer(x_, inputs[1]) return x_ return custom_forward for i, layer_module in enumerate(self.layers): if seqlen is None and batch is None: hidden_states = layer_module(hidden_states, attention_mask) else: assert seqlen is not None assert batch is not None if self.unpad_fmha: hidden_states = layer_module(hidden_states, cu_seqlens, host_cu_seqlens, maxseqlen_in_batch) print("hidden_states:", hidden_states) else: hidden_states = layer_module(hidden_states, attention_mask, seqlen, batch) if output_all_encoded_layers: all_encoder_layers.append(hidden_states) # Pad inputs and mask. It will insert back zero-padded tokens. Assume ntokens is total number of tokens (padded and non-padded) # and ntokens_unpad is total number of non-padded tokens. Then padding performs the following de-compression: # hidden_states[ntokens_unpad,hidden] -> hidden_states[ntokens,hidden] if self.unpad_fmha: hidden_states = paddle.scatter(zero_tensor, attention_indices, hidden_states) # todo: is self.maxseqlen same as maxseqlen? hidden_states = paddle.reshape( hidden_states, [batch, self.maxseqlen, self.hidden_size]) #print("hidden_states.shape:", hidden_states.shape) #print("hidden_states:", hidden_states) all_encoder_layers.append(hidden_states) return all_encoder_layers class UnpadFMHABertAttention(Layer): def __init__(self, config): super(UnpadFMHABertAttention, self).__init__() self.fmha = FMHA(config) self.output = BertSelfOutput(config) def forward(self, input_tensor, cu_seqlens, host_cu_seqlens, max_s, batch_size=None): self_output = self.fmha( input_tensor, cu_seqlens, host_cu_seqlens, max_s, is_training=self.training) attention_output = self.output(self_output, input_tensor) return attention_output class FusedDropoutResidualLn(Layer): def __init__(self, config, normalized_shape, epsilon=1e-12): super(FusedDropoutResidualLn, self).__init__() if isinstance(normalized_shape, numbers.Integral): normalized_shape = [normalized_shape] self._normalized_shape = list(normalized_shape) param_shape = [np.prod(self._normalized_shape)] self._weight_attr = None self._bias_attr = None self.weight = self.create_parameter( attr=self._weight_attr, shape=param_shape, default_initializer=Constant(1.0)) self.bias = self.create_parameter( attr=self._bias_attr, shape=param_shape, is_bias=True) self.p = config.hidden_dropout_prob self.epsilon = epsilon self.is_test = not self.training # todo: use default configs. self.fix_seed = False self.is_upscale_in_train = True self.seed_val = 0 def forward(self, hidden_states, input_tensor): out, dropout_mask, ln_mean, ln_var, dropout_residual_out = custom_fused_dropout_residual_ln( hidden_states, input_tensor, self.weight, self.bias, self.epsilon, self.is_test, self.fix_seed, self.seed_val, self.is_upscale_in_train, self.p) return out def extra_repr(self): return 'normalized_shape={}, epsilon={}'.format(self._normalized_shape, self.epsilon) # support nn(weight is not transposed) and nt(weight is transposed) class FusedDense(Layer): def __init__(self, in_features, out_features, weight_transpose=False, weight_attr=None, bias_attr=None, with_gelu=False, name=None): super(FusedDense, self).__init__() self._dtype = self._helper.get_default_dtype() self._weight_attr = weight_attr self._bias_attr = bias_attr self.weight_transpose = weight_transpose if weight_transpose: self.weight = self.create_parameter( shape=[out_features, in_features], attr=self._weight_attr, dtype=self._dtype, is_bias=False) else: self.weight = self.create_parameter( shape=[in_features, out_features], attr=self._weight_attr, dtype=self._dtype, is_bias=False) self.bias = self.create_parameter( shape=[out_features], attr=self._bias_attr, dtype=self._dtype, is_bias=True) self.with_gelu = with_gelu self.name = name def forward(self, hidden_states): out = fuse_dense( hidden_states, self.weight, self.bias, transx=False, transy=self.weight_transpose, with_gelu=self.with_gelu) return out class BertSelfOutput(Layer): def __init__(self, config): super(BertSelfOutput, self).__init__() self.fused_fc_bias = config.fused_bias_fc self.fused_dropout = config.fused_dropout_add_ln if not self.fused_fc_bias: self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.weight_transpose = False else: self.dense = FusedDense( config.hidden_size, config.hidden_size, weight_transpose=config.weight_transpose) self.weight_transpose = config.weight_transpose if self.fused_dropout: self.fused_dropout_add_ln = FusedDropoutResidualLn( config, config.hidden_size, epsilon=1e-12) else: self.layer_norm = nn.LayerNorm(config.hidden_size, epsilon=1e-12) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, hidden_states, input_tensor): #print("selfoutput input: hidden_states.shape = ", hidden_states.shape) hidden_states = self.dense(hidden_states) if not self.fused_dropout: hidden_states = self.dropout(hidden_states) hidden_states = hidden_states + input_tensor hidden_states = self.layer_norm(hidden_states) else: hidden_states = self.fused_dropout_add_ln(hidden_states, input_tensor) return hidden_states class BertIntermediate(Layer): def __init__(self, config): super(BertIntermediate, self).__init__() self.fused_fc_bias = config.fused_bias_fc if not self.fused_fc_bias: self.dense = nn.Linear(config.hidden_size, config.intermediate_size) self.weight_transpose = False self.intermediate_act_fn = get_activation() else: self.weight_transpose = config.weight_transpose self.dense = FusedDense( config.hidden_size, config.intermediate_size, weight_transpose=self.weight_transpose, with_gelu=True) def forward(self, hidden_states): if not self.fused_fc_bias: hidden_states = self.dense(hidden_states) hidden_states = self.intermediate_act_fn(hidden_states) else: hidden_states = self.dense(hidden_states) return hidden_states class BertOutput(Layer): def __init__(self, config): super(BertOutput, self).__init__() self.fused_fc_bias = config.fused_bias_fc self.fused_dropout = config.fused_dropout_add_ln if not self.fused_fc_bias: self.dense = nn.Linear(config.intermediate_size, config.hidden_size) self.weight_transpose = False else: self.dense = FusedDense( config.intermediate_size, config.hidden_size, weight_transpose=config.weight_transpose) self.weight_transpose = config.weight_transpose if self.fused_dropout: self.fused_dropout_add_ln = FusedDropoutResidualLn( config, config.hidden_size, epsilon=1e-12) else: self.dropout = nn.Dropout(config.hidden_dropout_prob) self.layer_norm = nn.LayerNorm(config.hidden_size, epsilon=1e-12) # todo: add fused_dropout opt. self.p = config.hidden_dropout_prob def forward(self, hidden_states, input_tensor): #print("BertOutput, input.shape = ", hidden_states.shape) hidden_states = self.dense(hidden_states) if not self.fused_dropout: hidden_states = self.dropout(hidden_states) hidden_states = hidden_states + input_tensor hidden_states = self.layer_norm(hidden_states) else: hidden_states = self.fused_dropout_add_ln(hidden_states, input_tensor) return hidden_states class BertEmbeddings(Layer): """ Include embeddings from word, position and token_type embeddings """ def __init__(self, vocab_size, hidden_size=768, hidden_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=16, config=None): super(BertEmbeddings, self).__init__() self.unpad_embed = False if config is not None: self.unpad_embed = config.unpad_embed self.unpad_fmha = config.unpad_fmha if self.unpad_embed: self.batch_size = config.batch_size self.max_seq_length = config.max_seq_length self.hidden_size = hidden_size self.word_embeddings = nn.Embedding( vocab_size, hidden_size, padding_idx=0) self.position_embeddings = nn.Embedding(max_position_embeddings, hidden_size) self.token_type_embeddings = nn.Embedding(type_vocab_size, hidden_size) self.layer_norm = nn.LayerNorm(hidden_size, epsilon=1e-12) self.dropout = nn.Dropout(hidden_dropout_prob) def forward(self, input_ids, token_type_ids=None, position_ids=None, attention_mask=None, attention_indices=None, seqlen=None, cu_seqlens=None, actual_seqlens=None, maxseqlen_in_batch=None): if position_ids is None: position_ids = gen_pos_id(input_ids) if token_type_ids is None: token_type_ids = paddle.zeros_like(input_ids) # todo(@limin29): in order to construct the shape of zero_tensor, we use pad method to compute token_type_embeddings. token_type_embeddings = self.token_type_embeddings(token_type_ids) if self.unpad_embed: assert self.unpad_fmha assert attention_mask is not None assert attention_indices is not None assert seqlen is not None assert cu_seqlens is not None assert maxseqlen_in_batch is not None cur_batch_size = input_ids.shape[0] zero_tensor = paddle.zeros_like(token_type_embeddings) zero_tensor = paddle.reshape(zero_tensor, [-1, self.hidden_size]) # attention_indices: 把attn_mask flatten后，提取非零元的下标。 # seqlen： 对[bs, max_seq_len]，每一行求和，代表获取每一行的实际seq_len（一维）。 # ntokens: 代表当前batch的所有实际seq_len之和。 # cu_seqlens: 对seq_len求prefix_sum的结果。 # actual_seqlens: 与seqlen相同的值。 # maxseqlen_in_batch: the max seqlen in a batch. #attention_indices, attention_mask, seqlen, ntokens, cu_seqlens, actual_seqlens, maxseqlen_in_batch = generate_mask( # attention_mask, unpad_fmha=self.unpad_fmha) print("maxseqlen_in_batch = ", maxseqlen_in_batch) input_ids = paddle.reshape(input_ids, [-1]) input_ids = paddle.gather(input_ids, attention_indices) position_ids = paddle.reshape(position_ids, [-1]) position_ids = paddle.gather(position_ids, attention_indices) #token_type_ids = paddle.reshape(token_type_ids, [-1]) #token_type_ids = paddle.gather(token_type_ids, attention_indices) token_type_embeddings = paddle.reshape(token_type_embeddings, [-1, self.hidden_size]) token_type_embeddings = paddle.gather(token_type_embeddings, attention_indices) input_embedings = self.word_embeddings(input_ids) position_embeddings = self.position_embeddings(position_ids) #token_type_embeddings = self.token_type_embeddings(token_type_ids) embeddings = input_embedings + position_embeddings + token_type_embeddings embeddings = self.layer_norm(embeddings) embeddings = self.dropout(embeddings) if not self.unpad_embed: return embeddings else: return embeddings, cur_batch_size, zero_tensor class BertPooler(Layer): """ Pool the result of BertEncoder. """ def __init__(self, hidden_size, pool_act="tanh"): super(BertPooler, self).__init__() #self.dense = nn.Linear(hidden_size, hidden_size) #self.weight_transpose = False self.dense = FusedDense(hidden_size, hidden_size) self.weight_transpose = self.dense.weight_transpose self.activation = nn.Tanh() self.pool_act = pool_act assert self.pool_act == "tanh" def forward(self, hidden_states): # We "pool" the model by simply taking the hidden state corresponding # to the first token. first_token_tensor = hidden_states[:, 0] pooled_output = self.dense(first_token_tensor) if self.pool_act == "tanh": pooled_output = self.activation(pooled_output) return pooled_output class BertModel(nn.Layer): """ The bare BERT Model transformer outputting raw hidden-states without any specific head on top. This model inherits from :class:`~paddlenlp.transformers.model_utils.PretrainedModel`. Refer to the superclass documentation for the generic methods. This model is also a Paddle `paddle.nn.Layer `__ subclass. Use it as a regular Paddle Layer and refer to the Paddle documentation for all matter related to general usage and behavior. Args: vocab_size (int): Vocabulary size of `inputs_ids` in `BertModel`. Also is the vocab size of token embedding matrix. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling `BertModel`. hidden_size (int, optional): Dimensionality of the embedding layer, encoder layer and pooler layer. Defaults to `768`. num_hidden_layers (int, optional): Number of hidden layers in the Transformer encoder. Defaults to `12`. num_attention_heads (int, optional): Number of attention heads for each attention layer in the Transformer encoder. Defaults to `12`. intermediate_size (int, optional): Dimensionality of the feed-forward (ff) layer in the encoder. Input tensors to ff layers are firstly projected from `hidden_size` to `intermediate_size`, and then projected back to `hidden_size`. Typically `intermediate_size` is larger than `hidden_size`. Defaults to `3072`. hidden_act (str, optional): The non-linear activation function in the feed-forward layer. ``"gelu"``, ``"relu"`` and any other paddle supported activation functions are supported. Defaults to `"gelu"`. hidden_dropout_prob (float, optional): The dropout probability for all fully connected layers in the embeddings and encoder. Defaults to `0.1`. attention_probs_dropout_prob (float, optional): The dropout probability used in MultiHeadAttention in all encoder layers to drop some attention target. Defaults to `0.1`. max_position_embeddings (int, optional): The maximum value of the dimensionality of position encoding, which dictates the maximum supported length of an input sequence. Defaults to `512`. type_vocab_size (int, optional): The vocabulary size of `token_type_ids`. Defaults to `16`. initializer_range (float, optional): The standard deviation of the normal initializer. Defaults to 0.02. pad_token_id (int, optional): The index of padding token in the token vocabulary. Defaults to `0`. pool_act (str, optional): The non-linear activation function in the pooling layer. Defaults to `"tanh"`. """ def __init__(self, config): super(BertModel, self).__init__() self.unpad = config.unpad self.pad_fmha = config.pad_fmha self.unpad_embed = config.unpad_embed self.unpad_fmha = config.unpad_fmha self.pad_token_id = config.pad_token_id self.initializer_range = config.initializer_range self.num_hidden_layers = config.num_hidden_layers self.maxseqlen = config.max_seq_length self.hidden_size = config.hidden_size # todo: self.embeddings = BertEmbeddings( config.vocab_size, config.hidden_size, config.hidden_dropout_prob, config.max_position_embeddings, config.type_vocab_size, config) self.encoder = BertEncoder(config) self.pooler = BertPooler(config.hidden_size, config.pool_act) def record_ckpt_vars(self, ckpt): emb = self.embeddings ckpt.embeddings([ emb.word_embeddings.weight, emb.position_embeddings.weight, emb.token_type_embeddings.weight, ]) ckpt.norm_after_embeddings([ emb.layer_norm.weight, emb.layer_norm.bias, ]) for idx in range(self.num_hidden_layers): self.encoder.record_ckpt_vars(ckpt, idx) pooler = self.pooler ckpt.pooler_fc([pooler.dense.weight, pooler.dense.bias], pooler.weight_transpose) def forward(self, input_ids, token_type_ids=None, position_ids=None, attention_mask=None, seq_len=None, prefix_sum_seq_len=None, host_prefix_sum_seq_len=None, max_seq_len=None, nonzeros_indices=None, output_hidden_states=False): if use_nv_input: if attention_mask is None: attention_mask = paddle.ones_like(input_ids) if token_type_ids is None: token_type_ids = paddle.zeros_like(input_ids) extended_attention_mask = attention_mask if not self.unpad and not self.pad_fmha: extended_attention_mask = extended_attention_mask.cast( dtype=paddle.float32) extended_attention_mask = ( 1.0 - extended_attention_mask) * -10000.0 attention_mask = extended_attention_mask else: if attention_mask is None: attention_mask = paddle.unsqueeze( (input_ids == self.pad_token_id ).astype(self.pooler.dense.weight.dtype) * -1e9, axis=[1, 2]) new_attention_mask = attention_mask attention_indices = nonzeros_indices seqlen = seq_len cu_seqlens = prefix_sum_seq_len host_cu_seqlens = host_prefix_sum_seq_len maxseqlen_in_batch = max_seq_len actual_seqlens = seqlen if self.unpad_embed or self.unpad_fmha: assert attention_indices is not None assert seqlen is not None assert cu_seqlens is not None assert host_cu_seqlens is not None assert maxseqlen_in_batch is not None if not self.unpad_embed: embedding_output = self.embeddings( input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) cur_batch_size = None zero_tensor = None attention_indices = None new_attention_mask = None seqlen = None cu_seqlens = None host_cu_seqlens = None actual_seqlens = None maxseqlen_in_batch = None else: embedding_output, cur_batch_size, zero_tensor = self.embeddings( input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, attention_indices=attention_indices, seqlen=seqlen, cu_seqlens=cu_seqlens, actual_seqlens=seqlen, maxseqlen_in_batch=maxseqlen_in_batch) encoder_outputs = self.encoder( embedding_output, attention_mask, output_hidden_states, cur_batch_size, self.maxseqlen, self.hidden_size, zero_tensor, attention_indices, new_attention_mask, seqlen, cu_seqlens, host_cu_seqlens, actual_seqlens, maxseqlen_in_batch) pooled_output = self.pooler(encoder_outputs[-1]) print("encoder_outputs[-1]:", encoder_outputs[-1]) print("pooled_output:", pooled_output) if output_hidden_states: return encoder_outputs, pooled_output else: return encoder_outputs[-1], pooled_output class BertLMPredictionHead(Layer): """ Bert Model with a `language modeling` head on top for CLM fine-tuning. """ def __init__(self, config, embedding_weights=None): super(BertLMPredictionHead, self).__init__() self.transform = FusedDense( config.hidden_size, config.hidden_size, with_gelu=True) self.weight_transpose = self.transform.weight_transpose assert config.hidden_act == "gelu" self.layer_norm = nn.LayerNorm(config.hidden_size, epsilon=1e-12) self.decoder_weight = self.create_parameter( shape=[config.vocab_size, config.hidden_size], dtype=self.transform.weight.dtype, is_bias=False) if embedding_weights is None else embedding_weights self.decoder_bias = self.create_parameter( shape=[config.vocab_size], dtype=self.decoder_weight.dtype, is_bias=True) def forward(self, hidden_states, masked_positions=None): if masked_positions is not None: hidden_states = paddle.reshape(hidden_states, [-1, hidden_states.shape[-1]]) hidden_states = paddle.tensor.gather(hidden_states, masked_positions) # gather masked tokens might be more quick hidden_states = self.transform(hidden_states) hidden_states = self.layer_norm(hidden_states) hidden_states = fuse_dense( hidden_states, self.decoder_weight, self.decoder_bias, transx=False, transy=True) return hidden_states class BertPretrainingHeads(Layer): """ Perform language modeling task and next sentence classification task. Args: hidden_size (int): See :class:`BertModel`. vocab_size (int): See :class:`BertModel`. activation (str): Activation function used in the language modeling task. embedding_weights (Tensor, optional): Decoding weights used to map hidden_states to logits of the masked token prediction. Its data type should be float32 and its shape is [vocab_size, hidden_size]. Defaults to `None`, which means use the same weights of the embedding layer. """ def __init__(self, config, embedding_weights=None): super(BertPretrainingHeads, self).__init__() self.predictions = BertLMPredictionHead(config, embedding_weights) #self.seq_relationship = nn.Linear(config.hidden_size, 2) #self.seq_relationship_weight_transpose = False self.seq_relationship = FusedDense(config.hidden_size, 2) self.seq_relationship_weight_transpose = self.seq_relationship.weight_transpose self.dense_seq_output = config.dense_seq_output self.share_weight = embedding_weights is not None def record_ckpt_vars(self, ckpt): pred_trans = self.predictions.transform ckpt.cls_pred_trans_fc([pred_trans.weight, pred_trans.bias], self.predictions.weight_transpose) norm = self.predictions.layer_norm ckpt.cls_pred_trans_norm([norm.weight, norm.bias]) assert self.share_weight ckpt.cls_pred_fc_bias(self.predictions.decoder_bias) seq_relation_fc = self.seq_relationship ckpt.cls_seq_relation_fc( [seq_relation_fc.weight, seq_relation_fc.bias], self.seq_relationship_weight_transpose) def forward(self, sequence_output, pooled_output, masked_positions=None): """ Args: sequence_output(Tensor): Sequence of hidden-states at the last layer of the model. It's data type should be float32 and its shape is [batch_size, sequence_length, hidden_size]. pooled_output(Tensor): The output of first token (`[CLS]`) in sequence. We "pool" the model by simply taking the hidden state corresponding to the first token. Its data type should be float32 and its shape is [batch_size, hidden_size]. masked_positions(Tensor, optional): A tensor indicates positions to be masked in the position embedding. Its data type should be int64 and its shape is [batch_size, mask_token_num]. `mask_token_num` is the number of masked tokens. It should be no bigger than `sequence_length`. Defaults to `None`, which means we output hidden-states of all tokens in masked token prediction. Returns: tuple: Returns tuple (``prediction_scores``, ``seq_relationship_score``). With the fields: - `prediction_scores` (Tensor): The scores of masked token prediction. Its data type should be float32. If `masked_positions` is None, its shape is [batch_size, sequence_length, vocab_size]. Otherwise, its shape is [batch_size, mask_token_num, vocab_size]. - `seq_relationship_score` (Tensor): The scores of next sentence prediction. Its data type should be float32 and its shape is [batch_size, 2]. """ prediction_scores = self.predictions(sequence_output, masked_positions) seq_relationship_score = self.seq_relationship(pooled_output) return prediction_scores, seq_relationship_score def nv_forward(self, sequence_output, pooled_output, masked_lm_labels, num_valid=None, masked_lm_ids=None, masked_lm_positions=None): print("sequence_output: ", sequence_output) if self.dense_seq_output: # nonzero indices index = masked_lm_positions sequence_flattened = paddle.index_select( sequence_output.reshape((-1, sequence_output.shape[-1])), index=index, axis=0) sequence_output = sequence_flattened prediction_scores = self.predictions(sequence_output) seq_relationship_score = self.seq_relationship(pooled_output) return prediction_scores, seq_relationship_score class BertForPretraining(nn.Layer): """ Bert Model with pretraining tasks on top. Args: bert (:class:`BertModel`): An instance of :class:`BertModel`. """ def __init__(self, bert, config): super(BertForPretraining, self).__init__() self.config = config self.bert = bert self.cls = BertPretrainingHeads( config, embedding_weights=self.bert.embeddings.word_embeddings.weight) def load_tf_ckpt(self, args, get_parameter_func): place = utility.get_place() ckpt = TFCkptHelper(args, self.config, args.tf_ckpt_path, place) self.bert.record_ckpt_vars(ckpt) self.cls.record_ckpt_vars(ckpt) ckpt.load(get_parameter_func) return ckpt def forward(self, input_ids, token_type_ids=None, position_ids=None, attention_mask=None, masked_positions=None, seq_len=None, prefix_sum_seq_len=None, host_prefix_sum_seq_len=None, max_seq_len=None, nonzeros_indices=None, num_valid=None, masked_lm_ids=None, masked_lm_positions=None): r""" Args: input_ids (Tensor): See :class:`BertModel`. token_type_ids (Tensor, optional): See :class:`BertModel`. position_ids (Tensor, optional): See :class:`BertModel`. attention_mask (Tensor, optional): See :class:`BertModel`. masked_positions(Tensor, optional): See :class:`BertPretrainingHeads`. Returns: tuple: Returns tuple (``prediction_scores``, ``seq_relationship_score``). With the fields: - `prediction_scores` (Tensor): The scores of masked token prediction. Its data type should be float32. If `masked_positions` is None, its shape is [batch_size, sequence_length, vocab_size]. Otherwise, its shape is [batch_size, mask_token_num, vocab_size]. - `seq_relationship_score` (Tensor): The scores of next sentence prediction. Its data type should be float32 and its shape is [batch_size, 2]. """ outputs = self.bert( input_ids, token_type_ids=token_type_ids, position_ids=position_ids, attention_mask=attention_mask, seq_len=seq_len, prefix_sum_seq_len=prefix_sum_seq_len, host_prefix_sum_seq_len=host_prefix_sum_seq_len, max_seq_len=max_seq_len, nonzeros_indices=nonzeros_indices) sequence_output, pooled_output = outputs prediction_scores, seq_relationship_score = self.cls( sequence_output, pooled_output, masked_positions, num_valid, masked_lm_ids, masked_lm_positions) return prediction_scores, seq_relationship_score class BertPretrainingCriterion(paddle.nn.Layer): """ Args: vocab_size(int): Vocabulary size of `inputs_ids` in `BertModel`. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling `BertModel`. """ def __init__(self, config): super(BertPretrainingCriterion, self).__init__() # CrossEntropyLoss is expensive since the inner reshape (copy) self.loss_fn = paddle.nn.loss.CrossEntropyLoss(ignore_index=-1) self.vocab_size = config.vocab_size self.dense_seq_output = config.dense_seq_output def forward(self, prediction_scores, seq_relationship_score, masked_lm_labels, next_sentence_labels, masked_lm_scale): """ Args: prediction_scores(Tensor): The scores of masked token prediction. Its data type should be float32. If `masked_positions` is None, its shape is [batch_size, sequence_length, vocab_size]. Otherwise, its shape is [batch_size, mask_token_num, vocab_size] seq_relationship_score(Tensor): The scores of next sentence prediction. Its data type should be float32 and its shape is [batch_size, 2] masked_lm_labels(Tensor): The labels of the masked language modeling, its dimensionality is equal to `prediction_scores`. Its data type should be int64. If `masked_positions` is None, its shape is [batch_size, sequence_length, 1]. Otherwise, its shape is [batch_size, mask_token_num, 1] next_sentence_labels(Tensor): The labels of the next sentence prediction task, the dimensionality of `next_sentence_labels` is equal to `seq_relation_labels`. Its data type should be int64 and its shape is [batch_size, 1] masked_lm_scale(Tensor or int): The scale of masked tokens. Used for the normalization of masked language modeling loss. If it is a `Tensor`, its data type should be int64 and its shape is equal to `prediction_scores`. Returns: Tensor: The pretraining loss, equals to the sum of `masked_lm_loss` plus the mean of `next_sentence_loss`. Its data type should be float32 and its shape is [1]. """ # paddle.static.Print(prediction_scores, message="print prediction_scores:") # paddle.static.Print(masked_lm_labels, message="print masked_lm_labels:") masked_lm_loss = F.cross_entropy( prediction_scores, masked_lm_labels, reduction='none', ignore_index=-1) # paddle.static.Print(masked_lm_loss, message="print masked_lm_loss:") masked_lm_loss = masked_lm_loss / masked_lm_scale next_sentence_loss = F.cross_entropy( seq_relationship_score, next_sentence_labels, reduction='none') valid_mask = masked_lm_labels != -1 total_loss_before_cast = paddle.sum(masked_lm_loss) + paddle.mean( next_sentence_loss) def func(): total_loss = total_loss_before_cast.astype('float32') mlm_acc = paddle.cast( paddle.sum((paddle.argmax( prediction_scores, axis=-1, keepdim=True) == masked_lm_labels) * valid_mask), dtype=masked_lm_scale.dtype) / masked_lm_scale return total_loss, mlm_acc, masked_lm_scale return func def nv_forward(self, prediction_scores, seq_relationship_score, masked_lm_labels, next_sentence_labels, num_valid=None, masked_lm_ids=None, masked_lm_positions=None): if self.dense_seq_output: loss_fct = paddle.nn.loss.CrossEntropyLoss(ignore_index=-1) else: loss_fct = paddle.nn.loss.CrossEntropyLoss(ignore_index=0) if self.dense_seq_output: masked_lm_labels_dense = masked_lm_ids #print("masked_lm_labels_dense = ", masked_lm_labels_dense) #paddle.static.Print(prediction_scores.reshape((-1, self.vocab_size)), message="print prediction_scores.reshape((-1, self.vocab_size)):",summarize=1024,print_phase="forward") #paddle.static.Print(masked_lm_labels_dense, message="print masked_lm_labels_dense:") masked_lm_loss = loss_fct( prediction_scores.reshape((-1, self.vocab_size)), masked_lm_labels_dense) #paddle.static.Print(masked_lm_loss, message="print masked_lm_loss:",print_phase="forward") #with paddle.static.device_guard('cpu'): # num_valid = paddle.numel(masked_lm_labels_dense) else: masked_lm_loss = loss_fct( prediction_scores.reshape((-1, self.vocab_size)), masked_lm_labels.reshape((-1, ))) nsp_loss_fct = paddle.nn.loss.CrossEntropyLoss(ignore_index=-1) next_sentence_loss = nsp_loss_fct( seq_relationship_score.reshape([-1, 2]), next_sentence_labels.reshape([-1])) #paddle.static.Print(next_sentence_loss, message="print next_sentence_loss:") total_loss = masked_lm_loss + next_sentence_loss # Masked Language Model Accuracy # NOTE: total_loss and mlm_acc use float32 in NV if not self.dense_seq_output: def func(): valid_mask = masked_lm_labels_flat != 0 num_valid_cnt = valid_mask.astype('int32').sum(dtype='float32') mlm_labels = masked_lm_labels_flat prediction_scores_flat = prediction_scores.reshape( (-1, prediction_scores.shape[-1])) mlm_predictions_scores = prediction_scores_flat mlm_predictions = mlm_predictions_scores.argmax(axis=-1) mlm_acc = ((mlm_predictions == mlm_labels) * valid_mask).sum(dtype='float32') / num_valid_cnt return total_loss, mlm_acc, num_valid_cnt else: mlm_labels = masked_lm_ids dtype = masked_lm_labels_dense.dtype if dtype == paddle.int32: dtype = 'int32' elif dtype == paddle.int64: dtype = 'int64' else: assert False mlm_predictions = prediction_scores.argmax( axis=-1, dtype=dtype, keepdim=False) assert len(mlm_predictions.shape) == 1 num_valid_cnt = num_valid def func(): mlm_acc = paddle.cast(mlm_predictions == masked_lm_labels_dense, 'float32').mean() return total_loss, mlm_acc, num_valid_cnt return func if use_nv_input: BertPretrainingHeads.forward = BertPretrainingHeads.nv_forward BertPretrainingCriterion.forward = BertPretrainingCriterion.nv_forward