def attention_bias_local_2d_block(mesh, h_dim, w_dim, memory_h_dim, memory_w_dim, dtype=tf.int32): """Bias for attention for local blocks where attention to right is disallowed. Create the bias matrix by using two separate masks, one for the memory part which doesn't overlap with the query and second which interacts with the query and should be disallowed to look to the right of the current query position. Args: mesh: a MeshTensorflow object h_dim: a mtf.Dimension w_dim: a mtf.Dimension memory_h_dim: a mtf.Dimension memory_w_dim: a mtf.Dimension dtype: a tf.dtype Returns: a mtf.Tensor with shape [block_length, memory_length] """ memory_height = mtf.Dimension(memory_h_dim.name, h_dim.size) memory_width = mtf.Dimension(memory_w_dim.name, w_dim.size) mask_top_visible = mtf.zeros(mesh, [h_dim, memory_height], dtype=dtype) mask_left_visible = mtf.zeros(mesh, [w_dim, memory_width], dtype=dtype) mask_query = mtf.greater( mtf.range(mesh, memory_height, dtype=tf.int32), mtf.range(mesh, memory_width, dtype=dtype)) width_mask = mtf.concat([mask_left_visible, mask_query], memory_width.name) mask = mtf.cast( mtf.concat([mask_top_visible, width_mask], memory_height.name), dtype=tf.float32) * -1e9 return mask
def tile_2d(physical_shape, tile_shape, outer_name="outer", inner_name="inner"): """2D tiling of a 3d physical mesh. The "inner" mesh dimension corresponds to the position within a tile of processors. The "outer" mesh dimension corresponds to which processor. Example: tile_2d(physical_shape=[8, 16, 2], tile_shape=[4, 4]) The "inner" dimension has size 4x4x2=32 and corresponds to the position within a 4x4 tile of processors. The "outer" dimension has size 8/4 * 16/4 = 8, and corresponds to the 8 tiles in the mesh. Args: physical_shape: a triple tile_shape: a pair outer_name: a string inner_name: a string Returns: mesh_shape: a mtf.Shape logical_to_physical: a list """ logical_to_physical = [] p0, p1, p2 = physical_shape t0, t1 = tile_shape tile_ring = _ring_2d(t0, t1) tiles_ring = _ring_2d(p0 // t0, p1 // t1) for logical_pnum in range(p0 * p1 * p2): core_on_chip = logical_pnum % p2 logical_chip_num = logical_pnum // p2 logical_pos_in_tile = logical_chip_num % (t0 * t1) logical_tile_num = logical_chip_num // (t0 * t1) tile_i, tile_j = tile_ring[logical_pos_in_tile] tiles_i, tiles_j = tiles_ring[logical_tile_num] physical_pnum = core_on_chip + p2 * (tile_i * p1 + tile_j + tiles_i * p1 * t0 + tiles_j * t1) logical_to_physical.append(physical_pnum) assert sorted(logical_to_physical) == list(range(p0 * p1 * p2)) tile_size = t0 * t1 * p2 num_tiles = p0 * p1 // (t0 * t1) mesh_shape = mtf.Shape([ mtf.Dimension(outer_name, int(num_tiles)), mtf.Dimension(inner_name, int(tile_size)) ]) return mesh_shape, logical_to_physical
def attention_bias_local_block(mesh, block_length, memory_length, dtype=tf.int32): """Bias for attention for local blocks where attention to right is disallowed. Create the bias matrix by using two separate masks, one for the memory part which doesn't overlap with the query and second which interacts with the query and should be disallowed to look to the right of the current query position. Args: mesh: a MeshTensorflow object block_length: a mtf.Dimension memory_length: a mtf.Dimension dtype: a tf.dtype Returns: a mtf.Tensor with shape [block_length, memory_length] """ memory_length = mtf.Dimension(memory_length.name, block_length.size) memory_mask = mtf.zeros(mesh, [block_length, memory_length], dtype=dtype) mask = mtf.cast(mtf.less(mtf.range(mesh, block_length, dtype=dtype), mtf.range(mesh, memory_length, dtype=dtype)), dtype=dtype) mask = mtf.cast( mtf.concat([memory_mask, mask], memory_length.name), dtype=tf.float32) * -1e9 return mask
def compress_mean(x, dim, compression_factor): """Compress by taking group means. Args: x: a Tensor dim: a dimension in x.shape compression_factor: an integer Returns: a Tensor """ dims = x.shape.dims pos = dims.index(dim) compressed_dim = mtf.Dimension(dim.name, dim.size // compression_factor) compression_factor_dim = mtf.Dimension( "compression_factor", compression_factor) new_shape = ( dims[:pos] + [compressed_dim, compression_factor_dim] + dims[pos + 1:]) x = mtf.reshape(x, new_shape) x = mtf.reduce_mean(x, reduced_dim=compression_factor_dim) return x
def conv2d(x, output_dim, filter_size=(3, 3), strides=(1, 1), padding="SAME", filter_initializer=None, variable_dtype=None, name=None): """2D Convolution. Args: x: a mtf.Tensor of format NHWC. output_dim: a mtf.Dimension, indicating the output channel dimension. filter_size: a list or tuple in format [filter_height, filter_width]. strides: a list or tuple in format [stride_height, stride_width]. padding: either "SAME" or "VALID". filter_initializer: the initializer for tf.get_variable. variable_dtype: a mtf.VariableDType name: a string used for tf.variable_scope. Returns: a mtf.Tensor. """ fh_dim = mtf.Dimension("fh", filter_size[0]) fw_dim = mtf.Dimension("fw", filter_size[1]) input_dim = x.shape[-1] with tf.variable_scope(name, default_name="conv2d"): if variable_dtype is None: variable_dtype = mtf.VariableDType(activation_dtype=x.dtype) conv_filter = mtf.get_variable(x.mesh, "kernel", [fh_dim, fw_dim, input_dim, output_dim], initializer=filter_initializer, dtype=variable_dtype) # Pad stride in batch and channel dimensions. strides = [1] + list(strides) + [1] return mtf.Conv2dOperation(x, conv_filter, strides, padding).outputs[0]
def __init__(self, spec, physical_shape): """Constructs a HierarchicalTiling. spec is a list corresponding to the logical dimensions. spec[i] corresponds to the i-th logical dimension and consists of a name and a list of integers, the list being the shape of logical axis i when it is physically projected to the physical mesh and then compacted. Striding information is omitted. By convention, the earlier dimensions get more strided. so the axis corresponding to the last dimension always gets projected to the tile specified by its shape. Args: spec: a list of (string, list-of-integers) pairs physical_shape: a list of integers """ self._names = [p[0] for p in spec] logical_ndims = len(spec) physical_ndims = len(physical_shape) projected_shapes = [p[1] for p in spec] if logical_ndims > 0 and projected_shapes[0] is None: # fill in missing value projected_shapes[0] = list(physical_shape) for s in projected_shapes[1:]: for i, x in enumerate(s): projected_shapes[0][i] //= x # compute strides, and verify that the spec is valid. products = [1] * physical_ndims sizes_and_strides = [] for s in reversed(projected_shapes): sizes_and_strides.append([(size, stride) for size, stride in zip(s, products)]) for i, x in enumerate(s): products[i] *= x if products != physical_shape: raise ValueError("mesh spec multiplies to the wrong size" "spec=%s physical_shape=%s products=%s" % (spec, physical_shape, products)) sizes_and_strides.reverse() self._physical_coordinates = _logical_to_physical_v1( sizes_and_strides, physical_shape) self._logical_to_physical = [ mtf.processor_coordinates_to_pnum(physical_shape, c) for c in self._physical_coordinates ] self._mesh_shape = mtf.Shape([ mtf.Dimension(name, mtf.list_product(s)) for name, s in zip(self._names, projected_shapes) ])
def multihead_attention_params(mesh, heads, io_channels, kv_channels, variable_dtype, combine=False): """Create Parameters for Multihead Attention. If the combine flag is set to True, then we create only one variable which stacks together all of the parameters. Otherwise, we create four separate variables. Args: mesh: a Mesh heads: a Dimension io_channels: a Dimension kv_channels: a Dimension variable_dtype: a mtf.VariableDType combine: a boolean Returns: wq: a Tensor with shape [heads, io_channels, kv_channels] wk: a Tensor with shape [heads, io_channels, kv_channels] wv: a Tensor with shape [heads, io_channels, kv_channels] wo: a Tensor with shape [heads, io_channels, kv_channels] """ qkvo = mtf.Dimension("qkvo", 4) qk_stddev = (io_channels.size ** -0.5) * (kv_channels.size ** -0.25) v_stddev = io_channels.size ** -0.5 # TODO(noam): should be: o_stddev = (kv_channels.size * heads.size) ** -0.5 # verify that this still works and change it. o_stddev = (io_channels.size * heads.size) ** -0.5 if combine: def qkvo_initializer(shape, dtype=None, partition_info=None, verify_shape=None): del partition_info, verify_shape return tf.random_normal(shape, dtype=dtype) * tf.reshape( tf.cast([qk_stddev, qk_stddev, v_stddev, o_stddev], dtype or tf.float32), [4, 1, 1, 1]) var = mtf.get_variable( mesh, "qkvo", mtf.Shape([qkvo, heads, io_channels, kv_channels]), initializer=qkvo_initializer, dtype=variable_dtype) return mtf.unstack(var, qkvo) else: return [mtf.get_variable( mesh, name, mtf.Shape([heads, io_channels, kv_channels]), initializer=tf.random_normal_initializer(stddev=stddev), dtype=variable_dtype) for name, stddev in zip( ["q", "k", "v", "o"], [qk_stddev, qk_stddev, v_stddev, o_stddev])]
def spec_to_mesh_shape(cls, spec, num_processors): """Compute mesh shape even without knowing the physical shape. This is useful in cases where the mesh shape must be computed before you know the physical_shape. Args: spec: a list of (string, list-of-integers) pairs num_processors: an integer Returns: a mtf.Shape """ logical_ndims = len(spec) names = [p[0] for p in spec] sizes = [p[1] for p in spec] sizes = [None if s is None else mtf.list_product(s) for s in sizes] if logical_ndims > 0 and sizes[0] is None: sizes[0] = num_processors // mtf.list_product(sizes[1:]) if mtf.list_product(sizes) != num_processors: raise ValueError("product of spec must be num_processors" " spec=%s num_processors=%s" % (spec, num_processors)) return mtf.Shape( [mtf.Dimension(name, s) for name, s in zip(names, sizes)])
def beam_search(logits_fn, initial_ids, alpha, states=None, eos_id=EOS_ID, stop_early=True, decode_length=None, use_tpu=True, dtype=tf.float32, layout=None, mesh_shape=None, num_prefilter=2): """Beam search with length penalties. Requires a function that can take the currently decoded symbols and return the logits for the next symbol. The implementation is inspired by https://arxiv.org/abs/1609.08144. When running, the beam search steps can be visualized by using tfdbg to watch the operations generating the output ids for each beam step. These operations have the pattern: (alive|finished)_topk_(seq,scores) Operations marked `alive` represent the new beam sequences that will be processed in the next step. Operations marked `finished` represent the completed beam sequences, which may be padded with 0s if no beams finished. Operations marked `seq` store the full beam sequence for the time step. Operations marked `scores` store the sequence's final log scores. The beam search steps will be processed sequentially in order, so when capturing observed from these operations, tensors, clients can make assumptions about which step is being recorded. num_prefilter is a theoretically lossy shortcut around slow performance of top_k on TPU on large Tensors and large k. This option should be removed once better top_k implementations on TPU are avialable. If num_prefilter is set to a nonzero value, then at each step we first compute the top num_prefilter sequences per beam and then compute the top k sequences overall from among those. Empirically, there seems to be no quality difference in setting num_prefilter to 2. Args: logits_fn: Interface to the model, to provide logits. Should take: step_num - mtf Scalar ids - mtf Tensor with shape [batch, beam, length] Should return: logits - [batch, beam, vocab_size], dtype=dtype initial_ids: a mtf.Tensor with shape [batch_dim, beam_dim, length_dim]) alpha: alpha for length penalty. states: list of mtf.Tensor eos_id: ID for end of sentence. stop_early: a boolean - stop once best sequence is provably determined. decode_length: a mtf Scalar of dtype tf.int32 - maximum length of decodes use_tpu: a boolean dtype: a tf.dtype layout: an optional string mesh_shape: an optional string num_prefilter: an optional integer Returns: Tuple of (decoded beams [batch, beam, length] decoding probabilities [batch, beam_size]) """ batch_dim, beam_dim, length_dim = initial_ids.shape.dims batch_and_beam_dim = mtf.Dimension(batch_dim.name, batch_dim.size * beam_dim.size) mesh = initial_ids.mesh batch_by_beam = mtf.Shape([batch_dim, beam_dim]) initial_log_probs = mtf.broadcast( mtf.one_hot(mtf.constant(mesh, 0, dtype=tf.int32), beam_dim, on_value=0.0, off_value=-INF, dtype=dtype), batch_by_beam) length_scalar = mtf.constant(mesh, length_dim.size, dtype=tf.int32) if decode_length is None: decode_length = length_scalar else: decode_length = mtf.minimum(decode_length, length_scalar) alive_log_probs = initial_log_probs alive_seq = initial_ids # Finished will keep track of all the sequences that have finished so far # Finished log probs will be negative infinity in the beginning # finished_flags will keep track of booleans finished_seq = initial_ids finished_scores = mtf.constant(mesh, -INF, batch_by_beam, dtype=dtype) # Setting the scores of the initial to negative infinity. finished_flags = mtf.constant(mesh, False, batch_by_beam, tf.bool) def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq, curr_scores, curr_finished): """Given sequences and scores, will gather the top k=beam size sequences. Args: finished_seq: Current finished sequences. [batch, beam, length] finished_scores: scores for each of these sequences. [batch, beam] finished_flags: finished bools for each of these sequences. [batch, beam] curr_seq: current topk sequence that has been grown by one position. [batch, beam, length] curr_scores: scores for each of these sequences. [batch, beam] curr_finished: Finished flags for each of these sequences. [batch, beam] Returns: Tuple of (Topk sequences based on scores, log probs of these sequences, Finished flags of these sequences, None (no states)) """ # Set the scores of the unfinished seq in curr_seq to large negative # values curr_scores += (1. - mtf.cast(curr_finished, curr_scores.dtype)) * -INF unused_batch_dim, beam_dim, unused_length_dim = finished_seq.shape.dims # concatenating the sequences and scores along beam axis def _my_concat(a, b): a = mtf.rename_dimension(a, "beam", "triple_beam") b = mtf.rename_dimension(b, "double_beam", "triple_beam") return mtf.concat([a, b], "triple_beam") curr_finished_seq = _my_concat(finished_seq, curr_seq) curr_finished_scores = _my_concat(finished_scores, curr_scores) curr_finished_flags = _my_concat(finished_flags, curr_finished) return compute_topk_scores_and_seq(curr_finished_seq, curr_finished_scores, curr_finished_scores, curr_finished_flags, beam_dim, "grow_finished") def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished): """Given sequences and scores, will gather the top k=beam size sequences. Args: curr_seq: current topk sequence that has been grown by one position. [batch, beam, length] curr_scores: scores for each of these sequences. [batch_size, beam_size] curr_log_probs: log probs for each of these sequences. [batch, beam] curr_finished: Finished flags for each of these sequences. [batch, beam] Returns: Tuple of (Topk sequences based on scores, log probs of these sequences, Finished flags of these sequences) """ # Set the scores of the finished seq in curr_seq to large negative # values curr_scores += mtf.cast(curr_finished, curr_scores.dtype) * -INF return compute_topk_scores_and_seq(curr_seq, curr_scores, curr_log_probs, curr_finished, beam_dim, "grow_alive") def grow_topk(i, alive_seq, alive_log_probs, states=None): r"""Inner beam search loop. This function takes the current alive sequences, and grows them to topk sequences where k = 2*beam. We use 2*beam because, we could have beam_size number of sequences that might hit <EOS> and there will be no alive sequences to continue. With 2*beam_size, this will not happen. This relies on the assumption the vocab size is > beam size. If this is true, we'll have at least beam_size non <EOS> extensions if we extract the next top 2*beam words. Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to https://arxiv.org/abs/1609.08144. Args: i: loop index alive_seq: Topk sequences decoded so far [batch, beam, length] alive_log_probs: probabilities of these sequences. [batch, beam] states: optional list of mtf.Tensor Returns: Tuple of (Topk sequences extended by the next word, The log probs of these sequences, The scores with length penalty of these sequences, Flags indicating which of these sequences have finished decoding, list of transformed decoding states) """ logits, new_states = logits_fn(i, alive_seq, states) batch_dim, beam_dim, vocab_dim = logits.shape.dims # Convert logits to normalized log probs candidate_log_probs = mtf.log_softmax(logits, vocab_dim) # Multiply the probabilities by the current probabilities of the beam. # (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1) log_probs = candidate_log_probs + alive_log_probs length_penalty = mtf.pow(((5. + mtf.cast(i + 1, logits.dtype)) / 6.), alpha) # scores have shape [batch, beam, vocab] curr_scores = log_probs / length_penalty # We find the top 2k sequences to make sure we get k alive sequences. # # TODO(noam): This is inefficient. We should separately compute the k # finished sequences (previously alive sequences + EOS), and the top k new # alive sequences. double_beam = mtf.Dimension("double_beam", beam_dim.size * 2) if use_tpu and layout is not None and mesh_shape is not None: # Do some partial top-k-ing first locally to avoid communication. # We reshape the logits from: # [batch, beam, vocab] to # [batch, beam, major_vocab, minor_vocab] # We first reduce (locally) across the minor_vocab dimension. This makes # the thing we need to broadcast smaller. # This also enables our shortcut of only picking the top num_prefilter # sequences per beam per major_vocab in the first pass. major_vocab_size = mtf.tensor_dim_to_mesh_dim_size( layout, mesh_shape, vocab_dim) major_vocab = mtf.Dimension(vocab_dim.name, major_vocab_size) minor_vocab = mtf.Dimension("minor_vocab", vocab_dim.size // major_vocab_size) curr_scores = mtf.reshape( curr_scores, [batch_dim, beam_dim, major_vocab, minor_vocab]) prefilter = mtf.Dimension("prefilter", num_prefilter or double_beam.size) # shape = [batch_dim, beam_dim, major_vocab, prefilter] top_scores, top_minor_vocab_ids = mtf.top_k( curr_scores, reduced_dim=minor_vocab, k_dim=prefilter) combined = mtf.Dimension( "combined", beam_dim.size * major_vocab.size * prefilter.size) top_scores = mtf.reshape(top_scores, [batch_dim, combined]) top_minor_vocab_ids = mtf.reshape(top_minor_vocab_ids, [batch_dim, combined]) # shpae = [batch_dim, double_beam] # ids are indices representing (beam, major_vocab, prefilter) top_scores, top_combined_ids = mtf.top_k(top_scores, reduced_dim=combined, k_dim=double_beam) top_minor_vocab_ids = mtf.gather( top_minor_vocab_ids, top_combined_ids, combined, output_shape=[batch_dim, double_beam]) top_beam_index = top_combined_ids // (major_vocab.size * prefilter.size) top_combined_ids -= top_beam_index * (major_vocab.size * prefilter.size) top_major_vocab_ids = top_combined_ids // prefilter.size top_combined_ids -= top_major_vocab_ids * prefilter.size top_ids = top_major_vocab_ids * minor_vocab.size + top_minor_vocab_ids else: beam_and_vocab_dim = mtf.Dimension("beam_and_vocab", beam_dim.size * vocab_dim.size) flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim]) # Flatten out (beam_size, vocab_size) probs into a list of possibilities flat_curr_scores = mtf.reshape(curr_scores, flat_shape, name="flatten_scores") top_scores, top_ids = mtf.top_k(flat_curr_scores, reduced_dim=beam_and_vocab_dim, k_dim=double_beam) # Work out what beam the top probs are in. top_beam_index = top_ids // vocab_dim.size top_ids %= vocab_dim.size # Unflatten the ids # Recovering the log probs because we will need to send them back top_log_probs = top_scores * length_penalty selector = mtf.one_hot(top_beam_index, beam_dim, dtype=tf.float32) def my_gather(tensor): return mtf.gather(tensor, top_beam_index, beam_dim, output_shape=mtf.Shape([ double_beam if d == beam_dim else d for d in tensor.shape.dims ])) # Gather up the most probable 2*beams both for the ids and finished_in_alive # bools top_seq = my_gather(alive_seq) # Append the most probable alive top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32) top_finished = mtf.equal(top_ids, eos_id) return (top_seq, top_log_probs, top_scores, top_finished, new_states, selector) def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags, *states): """Inner beam search loop. There are three groups of tensors, alive, finished, and topk. The alive group contains information about the current alive sequences The topk group contains information about alive + topk current decoded words the finished group contains information about finished sentences, that is, the ones that have decoded to <EOS>. These are what we return. The general beam search algorithm is as follows: While we haven't terminated (pls look at termination condition) 1. Grow the current alive to get beam*2 topk sequences 2. Among the topk, keep the top beam_size ones that haven't reached EOS into alive 3. Among the topk, keep the top beam_size ones have reached EOS into finished Repeat To make things simple with using fixed size tensors, we will end up inserting unfinished sequences into finished in the beginning. To stop that we add -ve INF to the score of the unfinished sequence so that when a true finished sequence does appear, it will have a higher score than all the unfinished ones. Args: i: loop index alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1] alive_log_probs: probabilities of the beams. [batch_size, beam_size] finished_seq: Current finished sequences. [batch_size, beam_size, i+1] finished_scores: scores for each of these sequences. [batch_size, beam_size] finished_flags: finished bools for each of these sequences. [batch_size, beam_size] *states: mtf Tensors Returns: Tuple of (Incremented loop index New alive sequences, Log probs of the alive sequences, New finished sequences, Scores of the new finished sequences, Flags indicating which sequence in finished as reached EOS, dict of final decoding states) """ states = [ mtf.replace_dimensions(state, batch_and_beam_dim, [batch_dim, beam_dim]) for state in states ] # Each inner loop, we carry out three steps: # 1. Get the current topk items. # 2. Extract the ones that have finished and haven't finished # 3. Recompute the contents of finished based on scores. (top2k_seq, top2k_log_probs, top2k_scores, top2k_finished, new_states, first_selector) = grow_topk(i, alive_seq, alive_log_probs, states) with tf.variable_scope("grow_alive"): alive_seq, alive_log_probs, _, second_selector = grow_alive( top2k_seq, top2k_scores, top2k_log_probs, top2k_finished) with tf.variable_scope("grow_finished"): finished_seq, finished_scores, finished_flags, _ = grow_finished( finished_seq, finished_scores, finished_flags, top2k_seq, top2k_scores, top2k_finished) old_beam_dim = mtf.Dimension("old_beam", beam_dim.size) selector = mtf.einsum([ mtf.rename_dimension(first_selector, beam_dim.name, old_beam_dim.name), second_selector ], output_shape=[batch_dim, old_beam_dim, beam_dim]) gathered_states = [] if use_tpu and layout is not None and mesh_shape is not None: # This hack combines the beam dimension with some of the batch dimension. # It makes gathering faster on TPU. # # Instead of multiplying by a [beam, beam] selector matrix, we instead # multiply by a [minor_batch*beam, minor_batch*beam] selector matrix. # This is theoretically more FLOPs, but it brings the matrix size closer # to the magic optimal value of 128. # # TODO(noam): file a bug with the XLA team to do this automatically major_batch_size = mtf.tensor_dim_to_mesh_dim_size( layout, mesh_shape, batch_dim) major_batch = mtf.Dimension(batch_dim.name, major_batch_size) minor_batch = mtf.Dimension("minor_batch", batch_dim.size // major_batch.size) old_minor_batch = mtf.Dimension("old_minor_batch", minor_batch.size) old_combined = mtf.Dimension("old_combined", minor_batch.size * beam_dim.size) combined = mtf.Dimension("new_combined", old_combined.size) same_minor_batch = mtf.to_float( mtf.equal(mtf.range(mesh, old_minor_batch, tf.float32), mtf.range(mesh, minor_batch, tf.float32))) selector = mtf.reshape( selector, [major_batch, minor_batch, old_beam_dim, beam_dim]) selector = mtf.einsum([selector, same_minor_batch], output_shape=[ major_batch, old_minor_batch, old_beam_dim, minor_batch, beam_dim ], reduced_dims=[]) selector = mtf.reshape(selector, [major_batch, old_combined, combined]) for state in new_states: s = mtf.replace_dimensions(state, [batch_dim, beam_dim], [major_batch, old_combined]) s = mtf.einsum([s, mtf.cast(selector, state.dtype)], reduced_dims=[old_combined], output_shape=mtf.replace_dimensions( state.shape, [batch_dim, beam_dim], [major_batch, combined])) gathered_states.append( mtf.replace_dimensions(s, [major_batch, combined], batch_and_beam_dim)) else: for state in new_states: state = mtf.einsum([ mtf.rename_dimension(state, beam_dim.name, old_beam_dim.name), mtf.cast(selector, state.dtype) ], reduced_dims=[old_beam_dim], output_shape=state.shape) state = mtf.replace_dimensions(state, [batch_dim, beam_dim], batch_and_beam_dim) gathered_states.append(state) return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags) + tuple(gathered_states) def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq, finished_scores, finished_in_finished, *unused_states): """Checking termination condition. We terminate when we decoded up to decode_length or the lowest scoring item in finished has a greater score that the highest prob item in alive divided by the max length penalty Args: i: loop index alive_log_probs: probabilities of the beams. [batch_size, beam_size] finished_scores: scores for each of these sequences. [batch_size, beam_size] finished_in_finished: finished bools for each of these sequences. [batch_size, beam_size] Returns: Bool. """ # TODO(noam): support a different decode length... # decode_length = mtf.constant(mesh, length_dim.size, dtype=tf.int32) # del alive_log_probs, finished_scores, finished_in_finished # return mtf.less(i, length_dim.size) if not stop_early: return mtf.less(i, decode_length) max_length_penalty = mtf.pow( ((5. + mtf.cast(decode_length, finished_scores.dtype)) / 6.), alpha) # The best possible score of the most likely alive sequence. lower_bound_alive_scores = mtf.gather( alive_log_probs, mtf.constant(mesh, 0, dtype=tf.int32), beam_dim) / max_length_penalty # Now to compute the lowest score of a finished sequence in finished # If the sequence isn't finished, we multiply it's score by 0. since # scores are all -ve, taking the min will give us the score of the lowest # finished item. lowest_score_of_finished_in_finished = mtf.reduce_min( finished_scores * mtf.cast(finished_in_finished, finished_scores.dtype), reduced_dim=beam_dim) # If none of the sequences have finished, then the min will be 0 and # we have to replace it by -ve INF if it is. The score of any seq in alive # will be much higher than -ve INF and the termination condition will not # be met. lowest_score_of_finished_in_finished += ((1. - mtf.cast( mtf.reduce_any(finished_in_finished, reduced_dim=beam_dim), finished_scores.dtype)) * -INF) bound_is_met = mtf.reduce_all( mtf.greater(lowest_score_of_finished_in_finished, lower_bound_alive_scores)) return mtf.logical_and(mtf.less(i, decode_length), mtf.logical_not(bound_is_met)) initial_step_num = mtf.constant(mesh, 0, dtype=tf.int32) states = [ mtf.replace_dimensions(state, [batch_dim, beam_dim], batch_and_beam_dim) for state in states ] while_loop_inputs = [ initial_step_num, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags ] + states (_, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags) = mtf.while_loop(_is_finished, inner_loop, while_loop_inputs, num_loop_vars=None if use_tpu else 6)[:6] # Accounting for corner case: It's possible that no sequence in alive for a # particular batch item ever reached EOS. In that case, we should just copy # the contents of alive for that batch item. tf.reduce_any(finished_flags, 1) # if 0, means that no sequence for that batch index had reached EOS. We need # to do the same for the scores as well. finished_seq = mtf.where( mtf.reduce_any(finished_flags, reduced_dim=beam_dim), finished_seq, alive_seq) finished_scores = mtf.where( mtf.reduce_any(finished_flags, reduced_dim=beam_dim), finished_scores, alive_log_probs) return finished_seq, finished_scores
def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags, *states): """Inner beam search loop. There are three groups of tensors, alive, finished, and topk. The alive group contains information about the current alive sequences The topk group contains information about alive + topk current decoded words the finished group contains information about finished sentences, that is, the ones that have decoded to <EOS>. These are what we return. The general beam search algorithm is as follows: While we haven't terminated (pls look at termination condition) 1. Grow the current alive to get beam*2 topk sequences 2. Among the topk, keep the top beam_size ones that haven't reached EOS into alive 3. Among the topk, keep the top beam_size ones have reached EOS into finished Repeat To make things simple with using fixed size tensors, we will end up inserting unfinished sequences into finished in the beginning. To stop that we add -ve INF to the score of the unfinished sequence so that when a true finished sequence does appear, it will have a higher score than all the unfinished ones. Args: i: loop index alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1] alive_log_probs: probabilities of the beams. [batch_size, beam_size] finished_seq: Current finished sequences. [batch_size, beam_size, i+1] finished_scores: scores for each of these sequences. [batch_size, beam_size] finished_flags: finished bools for each of these sequences. [batch_size, beam_size] *states: mtf Tensors Returns: Tuple of (Incremented loop index New alive sequences, Log probs of the alive sequences, New finished sequences, Scores of the new finished sequences, Flags indicating which sequence in finished as reached EOS, dict of final decoding states) """ states = [ mtf.replace_dimensions(state, batch_and_beam_dim, [batch_dim, beam_dim]) for state in states ] # Each inner loop, we carry out three steps: # 1. Get the current topk items. # 2. Extract the ones that have finished and haven't finished # 3. Recompute the contents of finished based on scores. (top2k_seq, top2k_log_probs, top2k_scores, top2k_finished, new_states, first_selector) = grow_topk(i, alive_seq, alive_log_probs, states) with tf.variable_scope("grow_alive"): alive_seq, alive_log_probs, _, second_selector = grow_alive( top2k_seq, top2k_scores, top2k_log_probs, top2k_finished) with tf.variable_scope("grow_finished"): finished_seq, finished_scores, finished_flags, _ = grow_finished( finished_seq, finished_scores, finished_flags, top2k_seq, top2k_scores, top2k_finished) old_beam_dim = mtf.Dimension("old_beam", beam_dim.size) selector = mtf.einsum([ mtf.rename_dimension(first_selector, beam_dim.name, old_beam_dim.name), second_selector ], output_shape=[batch_dim, old_beam_dim, beam_dim]) gathered_states = [] if use_tpu and layout is not None and mesh_shape is not None: # This hack combines the beam dimension with some of the batch dimension. # It makes gathering faster on TPU. # # Instead of multiplying by a [beam, beam] selector matrix, we instead # multiply by a [minor_batch*beam, minor_batch*beam] selector matrix. # This is theoretically more FLOPs, but it brings the matrix size closer # to the magic optimal value of 128. # # TODO(noam): file a bug with the XLA team to do this automatically major_batch_size = mtf.tensor_dim_to_mesh_dim_size( layout, mesh_shape, batch_dim) major_batch = mtf.Dimension(batch_dim.name, major_batch_size) minor_batch = mtf.Dimension("minor_batch", batch_dim.size // major_batch.size) old_minor_batch = mtf.Dimension("old_minor_batch", minor_batch.size) old_combined = mtf.Dimension("old_combined", minor_batch.size * beam_dim.size) combined = mtf.Dimension("new_combined", old_combined.size) same_minor_batch = mtf.to_float( mtf.equal(mtf.range(mesh, old_minor_batch, tf.float32), mtf.range(mesh, minor_batch, tf.float32))) selector = mtf.reshape( selector, [major_batch, minor_batch, old_beam_dim, beam_dim]) selector = mtf.einsum([selector, same_minor_batch], output_shape=[ major_batch, old_minor_batch, old_beam_dim, minor_batch, beam_dim ], reduced_dims=[]) selector = mtf.reshape(selector, [major_batch, old_combined, combined]) for state in new_states: s = mtf.replace_dimensions(state, [batch_dim, beam_dim], [major_batch, old_combined]) s = mtf.einsum([s, mtf.cast(selector, state.dtype)], reduced_dims=[old_combined], output_shape=mtf.replace_dimensions( state.shape, [batch_dim, beam_dim], [major_batch, combined])) gathered_states.append( mtf.replace_dimensions(s, [major_batch, combined], batch_and_beam_dim)) else: for state in new_states: state = mtf.einsum([ mtf.rename_dimension(state, beam_dim.name, old_beam_dim.name), mtf.cast(selector, state.dtype) ], reduced_dims=[old_beam_dim], output_shape=state.shape) state = mtf.replace_dimensions(state, [batch_dim, beam_dim], batch_and_beam_dim) gathered_states.append(state) return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags) + tuple(gathered_states)
def grow_topk(i, alive_seq, alive_log_probs, states=None): r"""Inner beam search loop. This function takes the current alive sequences, and grows them to topk sequences where k = 2*beam. We use 2*beam because, we could have beam_size number of sequences that might hit <EOS> and there will be no alive sequences to continue. With 2*beam_size, this will not happen. This relies on the assumption the vocab size is > beam size. If this is true, we'll have at least beam_size non <EOS> extensions if we extract the next top 2*beam words. Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to https://arxiv.org/abs/1609.08144. Args: i: loop index alive_seq: Topk sequences decoded so far [batch, beam, length] alive_log_probs: probabilities of these sequences. [batch, beam] states: optional list of mtf.Tensor Returns: Tuple of (Topk sequences extended by the next word, The log probs of these sequences, The scores with length penalty of these sequences, Flags indicating which of these sequences have finished decoding, list of transformed decoding states) """ logits, new_states = logits_fn(i, alive_seq, states) batch_dim, beam_dim, vocab_dim = logits.shape.dims # Convert logits to normalized log probs candidate_log_probs = mtf.log_softmax(logits, vocab_dim) # Multiply the probabilities by the current probabilities of the beam. # (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1) log_probs = candidate_log_probs + alive_log_probs length_penalty = mtf.pow(((5. + mtf.cast(i + 1, logits.dtype)) / 6.), alpha) # scores have shape [batch, beam, vocab] curr_scores = log_probs / length_penalty # We find the top 2k sequences to make sure we get k alive sequences. # # TODO(noam): This is inefficient. We should separately compute the k # finished sequences (previously alive sequences + EOS), and the top k new # alive sequences. double_beam = mtf.Dimension("double_beam", beam_dim.size * 2) if use_tpu and layout is not None and mesh_shape is not None: # Do some partial top-k-ing first locally to avoid communication. # We reshape the logits from: # [batch, beam, vocab] to # [batch, beam, major_vocab, minor_vocab] # We first reduce (locally) across the minor_vocab dimension. This makes # the thing we need to broadcast smaller. # This also enables our shortcut of only picking the top num_prefilter # sequences per beam per major_vocab in the first pass. major_vocab_size = mtf.tensor_dim_to_mesh_dim_size( layout, mesh_shape, vocab_dim) major_vocab = mtf.Dimension(vocab_dim.name, major_vocab_size) minor_vocab = mtf.Dimension("minor_vocab", vocab_dim.size // major_vocab_size) curr_scores = mtf.reshape( curr_scores, [batch_dim, beam_dim, major_vocab, minor_vocab]) prefilter = mtf.Dimension("prefilter", num_prefilter or double_beam.size) # shape = [batch_dim, beam_dim, major_vocab, prefilter] top_scores, top_minor_vocab_ids = mtf.top_k( curr_scores, reduced_dim=minor_vocab, k_dim=prefilter) combined = mtf.Dimension( "combined", beam_dim.size * major_vocab.size * prefilter.size) top_scores = mtf.reshape(top_scores, [batch_dim, combined]) top_minor_vocab_ids = mtf.reshape(top_minor_vocab_ids, [batch_dim, combined]) # shpae = [batch_dim, double_beam] # ids are indices representing (beam, major_vocab, prefilter) top_scores, top_combined_ids = mtf.top_k(top_scores, reduced_dim=combined, k_dim=double_beam) top_minor_vocab_ids = mtf.gather( top_minor_vocab_ids, top_combined_ids, combined, output_shape=[batch_dim, double_beam]) top_beam_index = top_combined_ids // (major_vocab.size * prefilter.size) top_combined_ids -= top_beam_index * (major_vocab.size * prefilter.size) top_major_vocab_ids = top_combined_ids // prefilter.size top_combined_ids -= top_major_vocab_ids * prefilter.size top_ids = top_major_vocab_ids * minor_vocab.size + top_minor_vocab_ids else: beam_and_vocab_dim = mtf.Dimension("beam_and_vocab", beam_dim.size * vocab_dim.size) flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim]) # Flatten out (beam_size, vocab_size) probs into a list of possibilities flat_curr_scores = mtf.reshape(curr_scores, flat_shape, name="flatten_scores") top_scores, top_ids = mtf.top_k(flat_curr_scores, reduced_dim=beam_and_vocab_dim, k_dim=double_beam) # Work out what beam the top probs are in. top_beam_index = top_ids // vocab_dim.size top_ids %= vocab_dim.size # Unflatten the ids # Recovering the log probs because we will need to send them back top_log_probs = top_scores * length_penalty selector = mtf.one_hot(top_beam_index, beam_dim, dtype=tf.float32) def my_gather(tensor): return mtf.gather(tensor, top_beam_index, beam_dim, output_shape=mtf.Shape([ double_beam if d == beam_dim else d for d in tensor.shape.dims ])) # Gather up the most probable 2*beams both for the ids and finished_in_alive # bools top_seq = my_gather(alive_seq) # Append the most probable alive top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32) top_finished = mtf.equal(top_ids, eos_id) return (top_seq, top_log_probs, top_scores, top_finished, new_states, selector)
def tile_2d(physical_shape, tile_shape, outer_name="outer", inner_name="inner", cores_name=None): """2D tiling of a 3d physical mesh. The "outer" mesh dimension corresponds to which tile. The "inner" mesh dimension corresponds to the position within a tile of processors. Optionally, if cores_name is specified, then a 3 dimensional logical mesh is returned, with the third dimension representing the two different cores within a chip. If cores_name is not specified, then the cores-in-a-chip dimension is folded into the inner dimension. TODO(noam): explain this better. Example: tile_2d(physical_shape=[8, 16, 2], tile_shape=[4, 4]) The "inner" dimension has size 4x4x2=32 and corresponds to the position within a 4x4 tile of processors. The "outer" dimension has size 8/4 * 16/4 = 8, and corresponds to the 8 tiles in the mesh. Args: physical_shape: a triple of integers [X, Y, cores] tile_shape: a pair outer_name: a string inner_name: a string cores_name: an optional string Returns: mesh_shape: a mtf.Shape logical_to_physical: a list """ logical_to_physical = [] p0, p1, p2 = physical_shape t0, t1 = tile_shape tile_ring = _ring_2d(t0, t1) tiles_ring = _ring_2d(p0 // t0, p1 // t1) for logical_pnum in range(p0 * p1 * p2): core_on_chip = logical_pnum % p2 logical_chip_num = logical_pnum // p2 logical_pos_in_tile = logical_chip_num % (t0 * t1) logical_tile_num = logical_chip_num // (t0 * t1) tile_i, tile_j = tile_ring[logical_pos_in_tile] tiles_i, tiles_j = tiles_ring[logical_tile_num] physical_pnum = core_on_chip + p2 * (tile_i * p1 + tile_j + tiles_i * p1 * t0 + tiles_j * t1) logical_to_physical.append(physical_pnum) assert sorted(logical_to_physical) == list(range(p0 * p1 * p2)) tile_size = t0 * t1 * p2 num_tiles = p0 * p1 // (t0 * t1) if cores_name: mesh_shape = mtf.Shape([ mtf.Dimension(outer_name, int(num_tiles)), mtf.Dimension(inner_name, int(t0 * t1)), mtf.Dimension(cores_name, int(p2)) ]) else: mesh_shape = mtf.Shape([ mtf.Dimension(outer_name, int(num_tiles)), mtf.Dimension(inner_name, int(tile_size)) ]) return mesh_shape, logical_to_physical
def masked_local_attention_1d(x, kv_channels, heads, window_size=128, master_dtype=tf.float32, slice_dtype=tf.float32, length_per_split=None, return_kv=None, params=None, name=None): """Attention to the source position and a neighborhood to the left of it. Attention for a given query position p can only see memory positions in the range (p - window_size, p]. Args: x: a mtf.Tensor with shape batch_dims + [length, io_channels] kv_channels: a mtf.Dimension (the size of the key and value vectors) heads: a mtf.Dimension (the number of heads) window_size: an integer master_dtype: a tf.dtype (deprecated - use params arg) slice_dtype: a tf.dtype (deprecated - use params arg) length_per_split: an optional integer indicating the part of the length dimension per processor. You can omit if the length dimension is not split. return_kv: an optional list onto which to append the computed k and v. params: an optional quadruple of Tensors (see multihead_attention_params()) name: an optional string. Returns: a Tensor with the same shape as x Raises: ValueError: if channels or depth don't match. """ with tf.variable_scope( name, default_name="masked_local_attention_1d", values=[x]): batch_dims = x.shape.dims[:-2] length, io_channels = x.shape.dims[-2:] if params is None: wq, wk, wv, wo = multihead_attention_vars( x.mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype, x.dtype) else: wq, wk, wv, wo = params # Get query q, keys k and values v. qkv_shape = mtf.Shape(batch_dims + [heads, length, kv_channels]) q = mtf.einsum([x, wq], qkv_shape) k = mtf.einsum([x, wk], qkv_shape) v = mtf.einsum([x, wv], qkv_shape) if return_kv is not None: return_kv.extend([k, v]) # Choose a suitable block size. # We choose the greatest divisor of length_per_split less than or equal # to max(window_size, 128) if length_per_split is None: length_per_split = length.size block_length = max(window_size, 128) while length_per_split % block_length != 0: block_length -= 1 query_block_length = mtf.Dimension("query_block_length", block_length) memory_block_length = mtf.Dimension("memory_block_length", block_length) # The num_blocks dimension gets the same name as the length dimension, # so it will be split in the same way. num_blocks = mtf.Dimension(length.name, length.size // block_length) q_shape = batch_dims + [heads, num_blocks, query_block_length, kv_channels] kv_shape = batch_dims + [ heads, num_blocks, memory_block_length, kv_channels] q = mtf.reshape(q, q_shape) k = mtf.reshape(k, kv_shape) v = mtf.reshape(v, kv_shape) # augment the keys and values for each block with keys and values for # the previous window_size timesteps. k = mtf.left_halo_exchange(k, num_blocks, memory_block_length, window_size) v = mtf.left_halo_exchange(v, num_blocks, memory_block_length, window_size) padded_memory_block_length = mtf.Dimension( "memory_block_length", window_size + block_length) mpos = mtf.range(x.mesh, padded_memory_block_length, tf.float32) qpos = mtf.range(x.mesh, query_block_length, tf.float32) + window_size # prevent looking forward mask = mtf.cast(mtf.greater(mpos, qpos), x.dtype) * -1e9 # prevent looking >=block_length timesteps backward mask += mtf.cast(mtf.less_equal(mpos, qpos - block_length), x.dtype) * -1e9 # Note: The first window_size-1 positions can see back into pre-time # where all the keys and values are zero. We could mask this out, but we # don't. o = dot_product_attention(q, k, v, mask=mask) o = mtf.reshape(o, batch_dims + [heads, length, kv_channels]) return mtf.einsum([o, wo], mtf.Shape(batch_dims + [length, io_channels]))