def local_attention1d_masked_decoder(x, kv_dim, heads_dim, feedforward_dim, hparams): """Image Transformer decoder with local1D masked layers.""" print(x) _, length_dim, model_dim = x.shape.dims for layer in range(hparams.num_decoder_layers): layer_name = "decoder_layer_%d" % layer with tf.variable_scope(layer_name): # Self attention layer length_per_split = mtf.tensor_dim_to_size_per_split( hparams.layout, hparams.mesh_shape, length_dim) x += layer_prepostprocess_dropout( mtf.layers.masked_local_attention_1d( mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"), kv_dim, heads_dim, window_size=hparams.block_length, length_per_split=length_per_split, name="self_att"), hparams) # ffn layer x += layer_prepostprocess_dropout( mtf.layers.dense_relu_dense( mtf.layers.layer_norm(x, model_dim, name="layer_norm_ffn"), feedforward_dim, hparams.dropout, dropout_broadcast_dims=[length_dim]), hparams) output = mtf.layers.layer_norm(x, model_dim, name="final_layer_norm") return output
def local_attention1d_masked_decoder(x, kv_dim, heads_dim, feedforward_dim, hparams): """Image Transformer decoder with local1D masked layers.""" print(x) _, length_dim, model_dim = x.shape.dims for layer in range(hparams.num_decoder_layers): layer_name = "decoder_layer_%d" % layer with tf.variable_scope(layer_name): # Self attention layer length_per_split = mtf.tensor_dim_to_size_per_split( hparams.layout, hparams.mesh_shape, length_dim) x += layer_prepostprocess_dropout( mtf.layers.masked_local_attention_1d( mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"), kv_dim, heads_dim, window_size=hparams.block_length, length_per_split=length_per_split, name="self_att"), hparams) # ffn layer x += layer_prepostprocess_dropout( mtf.layers.dense_relu_dense( mtf.layers.layer_norm(x, model_dim, name="layer_norm_ffn"), feedforward_dim, hparams.dropout, dropout_broadcast_dims=[length_dim]), hparams) output = mtf.layers.layer_norm(x, model_dim, name="final_layer_norm") return output
def serialize_num_microbatches(batch_dim, length_dim, mesh_shape, layout_rules, tokens_per_microbatch_per_replica=None): """Number of microbatches per batch for serialized training. We want to split each training step into multiple sequential steps to limit memory usage. Gradients are accumulated locally and reduced once. This function determines the number of microbatches per batch. If tokens_per_microbatch_per_replica=None, then the batch is not split. Args: batch_dim: a mtf.Dimension length_dim: a mtf.Dimension mesh_shape: an input to mtf.convert_to_shape() layout_rules: an input to mtf.convert_to_layout_rules() tokens_per_microbatch_per_replica: an optional integer, e.g. 2048 Returns: an integer """ if not tokens_per_microbatch_per_replica: return 1 batch_per_replica = mtf.tensor_dim_to_size_per_split( layout_rules, mesh_shape, batch_dim) # number of sequences per microbatch microbatch_size = max(1, tokens_per_microbatch_per_replica // length_dim.size) # decrease microbatch_size until it is a divisor of batch_per_replica # This is guaranteed to stop at microbatch_size=1 if not earlier. while batch_per_replica % microbatch_size: microbatch_size -= 1 num_microbatches = batch_per_replica // microbatch_size tf.logging.info( "serialize_num_microbatches: " "tokens_per_microbatch_per_replica=%d " "batch_dim=%s " "length_dim=%s " "batch_per_replica=%d " "num_microbatches=%d", tokens_per_microbatch_per_replica, batch_dim, length_dim, batch_per_replica, num_microbatches) return num_microbatches
def _layer_stack(self, x, layers, encoder_output=None, self_attention_mask=None, encdec_attention_mask=None, losses=None, step_num=None, encdec_tensors=None, states=None): """Encoder or decoder stack. Args: x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim] layers: an list of strings encoder_output: an optional mtf.Tensor with shape [<batch_dims>, encoder_length_dim, model_dim] self_attention_mask: an optional mtf.Tensor with shape [batch, length_dim, memory_length_dim] containing values 0 or -inf. encdec_attention_mask: an optional mtf.Tensor with shape [batch, length_dim, encoder_length_dim] containing values 0 or -inf. losses: a list to be appended-to step_num: an optional mtf integer Scalar (used in incrmenental mode) encdec_tensors: an optional list of num_layers tuples, each of the form (q_var, o_var, k, v), (used in incremental mode) states: an optional list of Tensors (used in incremental mode) Returns: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim] Raises: ValueError: if hparams make no sense """ hparams = self._hparams is_incremental = (step_num is not None) mode = getattr(hparams, "mode", tf.estimator.ModeKeys.TRAIN) is_training = mode == tf.estimator.ModeKeys.TRAIN def layer_prepostprocess_dropout(x): if is_incremental: return x return mtf.dropout( x, is_training, keep_prob=1.0 - hparams.layer_prepostprocess_dropout, noise_shape=mtf.Shape(self.batch_dims + [self.model_dim])) num_layers = len(layers) num_layer_norms = num_layers + 1 layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms) layer_norm_combined_var = mtf.get_variable( x.mesh, "layer_norm_scale", mtf.Shape([layer_norms_dim, self.model_dim]), initializer=tf.ones_initializer(), activation_dtype=x.dtype) layer_norm_vars = mtf.unstack(layer_norm_combined_var, layer_norms_dim) def normalize(x): scale = layer_norm_vars.pop(0) variance = mtf.reduce_mean(mtf.square(x), reduced_dim=self.model_dim) return x * mtf.rsqrt(variance + hparams.norm_epsilon) * scale if is_incremental: states = list(states) new_states = [] tf.logging.info("states = %s" % (states,)) for lnum, layer_type in enumerate(layers): with tf.variable_scope("%s_%d" % (layer_type, lnum)): if layer_type == "att": # Self attention layer if is_incremental: y, new_k, new_v = mtf.layers.multihead_self_attention_incremental( normalize(x), prev_k=states.pop(0), prev_v=states.pop(0), step_num=step_num, master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="att") new_states.append(new_k) new_states.append(new_v) x += y else: x += layer_prepostprocess_dropout( mtf.layers.multihead_attention( normalize(x), None, self_attention_mask, self.kv_dim, self.heads_dim, is_training, dropout=hparams.attention_dropout, dropout_broadcast_dims=[self.length_dim], master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="att")) elif layer_type == "enc_att": # Encoder-Decoder attention layer if is_incremental: # Encoder-Decoder attention layer q_var, o_var, k, v = encdec_tensors[lnum] x += mtf.layers.multihead_encdec_attention_incremental( normalize(x), q_var, o_var, k, v, encdec_attention_mask, name="enc_att") else: x += layer_prepostprocess_dropout( mtf.layers.multihead_attention( normalize(x), encoder_output, encdec_attention_mask, self.kv_dim, self.heads_dim, is_training, dropout=hparams.attention_dropout, dropout_broadcast_dims=[self.length_dim], master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="enc_att")) elif layer_type == "local_att": if is_incremental: y, new_k, new_v = mtf.layers.masked_local_attention_1d_incremental( normalize(x), prev_k=states.pop(0), prev_v=states.pop(0), step_num=step_num, master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="local_att") new_states.append(new_k) new_states.append(new_v) x += y else: x += layer_prepostprocess_dropout( mtf.layers.masked_local_attention_1d( normalize(x), self.kv_dim, self.heads_dim, is_training, window_size=hparams.local_attention_window_size, master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, length_per_split=mtf.tensor_dim_to_size_per_split( hparams.layout, hparams.mesh_shape, self.max_length_dim), name="local_att")) elif layer_type == "compressed_att": if is_incremental: raise ValueError("compressed_att incremental not implemented") else: x += layer_prepostprocess_dropout( mtf.layers.multihead_self_attention_memory_compressed( normalize(x), mask_right=True, compression_factor=hparams.compression_factor, kv_channels=self.kv_dim, heads=self.heads_dim, is_training=is_training, dropout=hparams.attention_dropout, dropout_broadcast_dims=[self.length_dim], master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="compressed_att")) else: if is_incremental: # insert length dimension. x_shape = x.shape shape_with_length = mtf.Shape( x_shape.dims[:-1] + [mtf.Dimension("length", 1)] + x_shape.dims[-1:]) x = mtf.reshape(x, shape_with_length) # ffn layer x += layer_prepostprocess_dropout( self._feedforward_layer(normalize(x), layer_type, losses=losses)) if is_incremental: # remove length dimension x = mtf.reshape(x, x_shape) x = layer_prepostprocess_dropout(normalize(x)) assert not layer_norm_vars if is_incremental: return x, new_states else: return x
def _layer_stack(self, x, layers, encoder_output=None, self_attention_mask=None, encdec_attention_mask=None, losses=None, step_num=None, encdec_tensors=None, states=None): """Encoder or decoder stack. Args: x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim] layers: an list of strings encoder_output: an optional mtf.Tensor with shape [<batch_dims>, encoder_length_dim, model_dim] self_attention_mask: an optional mtf.Tensor with shape [batch, length_dim, memory_length_dim] containing values 0 or -inf. encdec_attention_mask: an optional mtf.Tensor with shape [batch, length_dim, encoder_length_dim] containing values 0 or -inf. losses: a list to be appended-to step_num: an optional mtf integer Scalar (used in incrmenental mode) encdec_tensors: an optional list of num_layers tuples, each of the form (q_var, o_var, k, v), (used in incremental mode) states: an optional list of Tensors (used in incremental mode) Returns: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim] Raises: ValueError: if hparams make no sense """ hparams = self._hparams is_incremental = (step_num is not None) def layer_prepostprocess_dropout(x): if is_incremental: return x return mtf.dropout( x, keep_prob=1.0 - hparams.layer_prepostprocess_dropout, noise_shape=mtf.Shape(self.batch_dims + [self.model_dim])) num_layers = len(layers) num_layer_norms = num_layers + 1 layer_norms_dim = mtf.Dimension("layer_norms", num_layer_norms) layer_norm_combined_var = mtf.get_variable( x.mesh, "layer_norm_scale", mtf.Shape([layer_norms_dim, self.model_dim]), initializer=tf.ones_initializer(), activation_dtype=x.dtype) layer_norm_vars = mtf.unstack(layer_norm_combined_var, layer_norms_dim) def normalize(x): scale = layer_norm_vars.pop(0) variance = mtf.reduce_mean(mtf.square(x), reduced_dim=self.model_dim) return x * mtf.rsqrt(variance + hparams.norm_epsilon) * scale if is_incremental: states = list(states) new_states = [] tf.logging.info("states = %s" % (states,)) for lnum, layer_type in enumerate(layers): with tf.variable_scope("%s_%d" % (layer_type, lnum)): if layer_type == "att": # Self attention layer if is_incremental: y, new_k, new_v = mtf.layers.multihead_self_attention_incremental( normalize(x), prev_k=states.pop(0), prev_v=states.pop(0), step_num=step_num, master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="att") new_states.append(new_k) new_states.append(new_v) x += y else: x += layer_prepostprocess_dropout( mtf.layers.multihead_attention( normalize(x), None, self_attention_mask, self.kv_dim, self.heads_dim, dropout=hparams.attention_dropout, dropout_broadcast_dims=[self.length_dim], master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="att")) elif layer_type == "enc_att": # Encoder-Decoder attention layer if is_incremental: # Encoder-Decoder attention layer q_var, o_var, k, v = encdec_tensors[lnum] x += mtf.layers.multihead_encdec_attention_incremental( normalize(x), q_var, o_var, k, v, encdec_attention_mask, name="enc_att") else: x += layer_prepostprocess_dropout( mtf.layers.multihead_attention( normalize(x), encoder_output, encdec_attention_mask, self.kv_dim, self.heads_dim, dropout=hparams.attention_dropout, dropout_broadcast_dims=[self.length_dim], master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="enc_att")) elif layer_type == "local_att": if is_incremental: y, new_k, new_v = mtf.layers.masked_local_attention_1d_incremental( normalize(x), prev_k=states.pop(0), prev_v=states.pop(0), step_num=step_num, master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="local_att") new_states.append(new_k) new_states.append(new_v) x += y else: x += layer_prepostprocess_dropout( mtf.layers.masked_local_attention_1d( normalize(x), self.kv_dim, self.heads_dim, window_size=hparams.local_attention_window_size, master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, length_per_split=mtf.tensor_dim_to_size_per_split( hparams.layout, hparams.mesh_shape, self.max_length_dim), name="local_att")) elif layer_type == "compressed_att": if is_incremental: raise ValueError("compressed_att incremental not implemented") else: x += layer_prepostprocess_dropout( mtf.layers.multihead_self_attention_memory_compressed( normalize(x), mask_right=True, compression_factor=hparams.compression_factor, kv_channels=self.kv_dim, heads=self.heads_dim, dropout=hparams.attention_dropout, dropout_broadcast_dims=[self.length_dim], master_dtype=self.master_dtype, slice_dtype=self.slice_dtype, name="compressed_att")) else: if is_incremental: # insert length dimension. x_shape = x.shape shape_with_length = mtf.Shape( x_shape.dims[:-1] + [mtf.Dimension("length", 1)] + x_shape.dims[-1:]) x = mtf.reshape(x, shape_with_length) # ffn layer x += layer_prepostprocess_dropout( self._feedforward_layer(normalize(x), layer_type, losses=losses)) if is_incremental: # remove length dimension x = mtf.reshape(x, x_shape) x = layer_prepostprocess_dropout(normalize(x)) assert not layer_norm_vars if is_incremental: return x, new_states else: return x