def multihead_encdec_attention_incremental(query_antecedent,
q_var,
o_var,
k,
v,
mask,
name="multihead_attention"):
"""Incremental attention over encoder (one decode step).
In order to use only one variable containing the four weight matrices
packed together, we insist that the query and memory antecedents have the
same dimensionality (io_channels) and that the keys and values have the
same dimensionality (kv_channels).
memory_dims is a subset of query_dims
Args:
query_antecedent: a mtf.Tensor with shape query_dims + [io_channels]
q_var: a mtf.Tensor with shape [heads, io_channels, kv_channels]
o_var: a mtf.Tensor with shape [heads, io_channels, kv_channels]
k: memory_dims + [heads, memory_length, kv_channels]
v: memory_dims + [heads, memory_length, kv_channels]
mask: mask Tensor (see attention_mask())
name: an optional string.
Returns:
A mtf.Tensor with shape [batch, qlen, io_channels]
"""
heads, _, kv_channels = k.shape.dims[-3:]
query_dims = query_antecedent.shape.dims[:-1]
with tf.variable_scope(name, default_name="multihead_attention"):
q = mtf.einsum([query_antecedent, q_var],
mtf.Shape(query_dims + [heads, kv_channels]))
o = dot_product_attention(q, k, v, mask)
return mtf.einsum([o, o_var], query_antecedent.shape)
def masked_local_attention_1d_incremental(x,
prev_k,
prev_v,
step_num,
master_dtype,
slice_dtype,
name=None):
"""Incremental local self-attention (one decode step).
Incremental version of masked_local_attention_1d()
Args:
x: a mtf.Tensor with shape [batch..., io_channels]
prev_k: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
prev_v: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
step_num: mtf Scalar with dtype tf.int32
master_dtype: a tf.dtype
slice_dtype: a tf.dtype
name: an optional string.
Returns:
y: A mtf.Tensor with shape [batch..., io_channels]
new_k: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
new_v: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
Raises:
ValueError: if the dimensions do not match.
"""
batch_dims = x.shape.dims[:-1]
io_channels = x.shape.dims[-1]
heads, window_length, kv_channels = prev_k.shape.dims[-3:]
with tf.variable_scope(name, default_name="multihead_attention"):
q_var, k_var, v_var, o_var = multihead_attention_vars(
x.mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype,
x.dtype)
q = mtf.einsum([x, q_var],
mtf.Shape(batch_dims + [heads, kv_channels]))
k = mtf.einsum([x, k_var],
mtf.Shape(batch_dims + [heads, kv_channels]))
v = mtf.einsum([x, v_var],
mtf.Shape(batch_dims + [heads, kv_channels]))
current_position = mtf.equal(
mtf.range(x.mesh, window_length, dtype=tf.int32),
mtf.mod(step_num, window_length.size))
k = mtf.where(current_position, k, prev_k, output_shape=prev_k.shape)
v = mtf.where(current_position, v, prev_v, output_shape=prev_v.shape)
o = dot_product_attention(q, k, v, mask=None)
y = mtf.einsum([o, o_var], x.shape)
return y, k, v
def multihead_self_attention_incremental(query_antecedent,
prev_k,
prev_v,
step_num,
name="multihead_attention"):
"""Incremental self-attention (one decode step).
In order to use only one variable containing the four weight matrices
packed together, we insist that the query and memory antecedents have the
same dimensionality (io_channels) and that the keys and values have the
same dimensionality (kv_channels).
Args:
query_antecedent: a mtf.Tensor with shape [batch..., io_channels]
prev_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
prev_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
step_num: mtf Scalar with dtype tf.int32
name: an optional string.
Returns:
y: A mtf.Tensor with shape [batch..., io_channels]
new_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
new_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels]
Raises:
ValueError: if the dimensions do not match.
"""
batch_dims = query_antecedent.shape.dims[:-1]
io_channels = query_antecedent.shape.dims[-1]
heads, memory_length, kv_channels = prev_k.shape.dims[-3:]
with tf.variable_scope(name, default_name="multihead_attention"):
q_var, k_var, v_var, o_var = multihead_attention_vars(
query_antecedent.mesh, heads, io_channels, kv_channels,
query_antecedent.dtype)
memory_antecedent = query_antecedent
q = mtf.einsum([query_antecedent, q_var],
mtf.Shape(batch_dims + [heads, kv_channels]))
k = mtf.einsum([memory_antecedent, k_var],
mtf.Shape(batch_dims + [heads, kv_channels]))
v = mtf.einsum([memory_antecedent, v_var],
mtf.Shape(batch_dims + [heads, kv_channels]))
k = prev_k + mtf.multiply(
k, mtf.one_hot(step_num, memory_length), output_shape=prev_k.shape)
v = prev_v + mtf.multiply(
v, mtf.one_hot(step_num, memory_length), output_shape=prev_v.shape)
mask = mtf.to_float(
mtf.greater(
mtf.range(query_antecedent.mesh, memory_length,
dtype=tf.int32), step_num)) * -1e9
o = dot_product_attention(q, k, v, mask)
y = mtf.einsum([o, o_var], query_antecedent.shape)
return y, k, v
def dense(x,
output_dim,
reduced_dims=None,
expert_dims=None,
use_bias=True,
activation=None,
master_dtype=tf.float32,
slice_dtype=tf.float32,
name=None):
"""Dense layer doing (kernel*x + bias) computation.
Args:
x: a mtf.Tensor of shape [..., reduced_dims].
output_dim: a mtf.Dimension
reduced_dims: an optional list of mtf.Dimensions of x to be reduced. If
omitted, we reduce the last dimension.
expert_dims: an optional list of mtf.Dimension which represent different
experts. Different experts get different weights.
use_bias: a boolean, whether to add bias.
activation: an optional function from mtf.Tensor to mtf.Tensor
master_dtype: a tf.dtype
slice_dtype: a tf.dtype
name: a string. variable scope.
Returns:
a mtf.Tensor of shape [..., output_dim].
"""
if expert_dims is None:
expert_dims = []
if reduced_dims is None:
reduced_dims = x.shape.dims[-1:]
w_shape = mtf.Shape(expert_dims + reduced_dims + [output_dim])
output_shape = mtf.Shape(
[d for d in x.shape.dims if d not in reduced_dims] + [output_dim])
with tf.variable_scope(name, default_name="dense"):
stddev = mtf.list_product(d.size for d in reduced_dims)**-0.5
w = mtf.get_variable(
x.mesh,
"kernel",
w_shape,
initializer=tf.random_normal_initializer(stddev=stddev),
master_dtype=master_dtype,
slice_dtype=slice_dtype,
activation_dtype=x.dtype)
y = mtf.einsum([x, w], output_shape)
if use_bias:
b = mtf.get_variable(x.mesh,
"bias",
mtf.Shape(expert_dims + [output_dim]),
initializer=tf.zeros_initializer(),
activation_dtype=x.dtype)
y += b
if activation is not None:
y = activation(y)
return y
def dot_product_attention(q,
k,
v,
mask,
dropout=0.0,
dropout_broadcast_dims=None,
extra_logit=None):
"""Dot-product attention.
Args:
q: Tensor with shape [...., length_q, depth_k]. Typically leading dimensions
are [batch, heads].
k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
match with q.
v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
match with q.
mask: mask Tensor (see attention_mask())
dropout: a float.
dropout_broadcast_dims: an optional list of mtf.Dimension
extra_logit: an optional scalar or tensor
Returns:
Tensor with shape [..., length_q, depth_v].
"""
length_kv = k.shape.dims[-2]
logits_shape = mtf.Shape(q.shape.dims[:-1] + [length_kv])
logits = mtf.einsum([q, k], logits_shape)
if mask is not None:
logits += mask
weights = mtf.softmax(logits, length_kv, extra_logit=extra_logit)
if dropout != 0.0:
weights = mtf.dropout(weights,
1.0 - dropout,
noise_shape=weights.shape -
dropout_broadcast_dims)
depth_v = v.shape.dims[-1]
outputs_shape = mtf.Shape(q.shape.dims[:-1] + [depth_v])
outputs = mtf.einsum([weights, v], outputs_shape)
return outputs
def dense_relu_dense(x,
hidden_channels,
dropout=0.0,
dropout_broadcast_dims=None,
name=None):
"""Hidden layer with ReLU activation followed by linear projection.
The output has the same number of channels as the input.
Args:
x: a mtf.Tensor
hidden_channels: a mtf.Dimension - channels in the hidden layer
dropout: an optional float
dropout_broadcast_dims: an optional list of mtf.Dimension
name: an optional string
Returns:
a mtf.Tensor with the same shape as x.
"""
with tf.variable_scope(name, default_name="dense_relu_dense"):
io_channels = x.shape.dims[-1]
stddev = (hidden_channels.size * io_channels.size)**-0.25
io = mtf.Dimension("io", 2)
w = mtf.get_variable(
x.mesh,
"kernel",
mtf.Shape([io, io_channels, hidden_channels]),
initializer=tf.random_normal_initializer(stddev=stddev),
activation_dtype=x.dtype)
wi, wo = mtf.unstack(w, io)
h = mtf.relu(mtf.einsum([x, wi]))
if dropout != 0.0:
h = mtf.dropout(h,
1.0 - dropout,
noise_shape=h.shape - dropout_broadcast_dims)
return mtf.einsum([h, wo])
def multihead_attention(query_antecedent,
memory_antecedent,
mask,
kv_channels,
heads,
dropout=0.0,
dropout_broadcast_dims=None,
name="multihead_attention"):
"""Multihead scaled-dot-product attention with input/output transformations.
In order to use only one variable containing the four weight matrices
packed together, we insist that the query and memory antecedents have the
same dimensionality (io_channels) and that the keys and values have the
same dimensionality (kv_channels).
Args:
query_antecedent: a mtf.Tensor with shape
[<batch_dims>, query_length, io_channels]
memory_antecedent: a mtf.Tensor with shape
[batch, memory_length, io_channels] (optional)
mask: mask Tensor (see attention_mask())
kv_channels: a mtf.Dimension (the size of the key and value vectors)
heads: a mtf.Dimension (the number of heads)
dropout: a floating point value
dropout_broadcast_dims: an optional list of mtf.Dimension
name: an optional string.
Returns:
A mtf.Tensor with shape [batch, query_length, io_channels]
Raises:
ValueError: if the dimensions do not match.
"""
batch_dims = query_antecedent.shape.dims[:-2]
query_length, io_channels = query_antecedent.shape.dims[-2:]
with tf.variable_scope(name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
q_var, k_var, v_var, o_var = multihead_attention_vars(
query_antecedent.mesh, heads, io_channels, kv_channels,
query_antecedent.dtype)
if memory_antecedent is None:
memory_antecedent = rename_length_to_memory_length(
query_antecedent, query_length.name)
memory_batch_dims = memory_antecedent.shape.dims[:-2]
memory_length, memory_channels = memory_antecedent.shape.dims[-2:]
if memory_batch_dims != batch_dims:
raise ValueError("memory batch must equal query batch")
if memory_channels != io_channels:
raise ValueError("memory channels must equal query channels")
q = mtf.einsum([query_antecedent, q_var],
mtf.Shape(batch_dims +
[heads, query_length, kv_channels]))
k = mtf.einsum([memory_antecedent, k_var],
mtf.Shape(batch_dims +
[heads, memory_length, kv_channels]))
v = mtf.einsum([memory_antecedent, v_var],
mtf.Shape(batch_dims +
[heads, memory_length, kv_channels]))
o = dot_product_attention(q, k, v, mask, dropout,
dropout_broadcast_dims)
return mtf.einsum([o, o_var],
mtf.Shape(batch_dims + [query_length, io_channels]))
def masked_local_attention_1d(query_antecedent,
memory_antecedent,
kv_channels,
heads,
block_length=128,
name=None):
"""Attention to the source position and a neighborhood to the left of it.
The sequence is divided into blocks of length block_size.
Attention for a given query position can only see memory positions
less than or equal to the query position, in the corresponding block
and the previous block.
Args:
query_antecedent: a mtf.Tensor with shape [batch, query_length, io_channels]
memory_antecedent: a mtf.Tensor with shape
[batch, memory_length, io_channels] (optional). Currently, memory_length
must have the same size as query_length, but a different name.
kv_channels: a mtf.Dimension (the size of the key and value vectors)
heads: a mtf.Dimension (the number of heads)
block_length: an integer, representing receptive fields for attention.
name: an optional string.
Returns:
a Tensor of shape [batch, query_length, io_channels]
Raises:
ValueError: if channels or depth don't match.
"""
with tf.variable_scope(name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
batch, query_length, io_channels = query_antecedent.shape.dims
q_var, k_var, v_var, o_var = multihead_attention_vars(
query_antecedent.mesh, heads, io_channels, kv_channels,
query_antecedent.dtype)
if memory_antecedent is None:
memory_antecedent = rename_length_to_memory_length(
query_antecedent, query_length.name)
memory_batch, memory_length, memory_channels = memory_antecedent.shape.dims
if memory_batch != batch:
raise ValueError("memory batch must equal query batch")
if memory_channels != io_channels:
raise ValueError("memory channels must equal query channels")
# Get query q, keys k and values v.
q = mtf.einsum([query_antecedent, q_var],
mtf.Shape([batch, heads, query_length, kv_channels]))
k = mtf.einsum([memory_antecedent, k_var],
mtf.Shape([batch, heads, memory_length, kv_channels]))
v = mtf.einsum([memory_antecedent, v_var],
mtf.Shape([batch, heads, memory_length, kv_channels]))
# Let's assume for now we don't have padding and the block length equally
# divides the memory length.
block_length = (query_length.size if
query_length.size < block_length * 2 else block_length)
blength = mtf.Dimension("block_length", block_length)
mlength = mtf.Dimension("mem_block_length", block_length)
num_blocks = mtf.Dimension("num_blocks",
query_length.size // block_length)
q = mtf.reshape(
q, mtf.Shape([batch, heads, num_blocks, blength, kv_channels]))
k = mtf.reshape(
k, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels]))
v = mtf.reshape(
v, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels]))
# compute attention for the first query block.
def first_block_attention():
"""Compute attention for the first block."""
first_q = mtf.slice(q, 0, 1, num_blocks.name)
first_k = mtf.slice(k, 0, 1, num_blocks.name)
first_v = mtf.slice(v, 0, 1, num_blocks.name)
first_output = dot_product_attention(first_q,
first_k,
first_v,
mask=None)
return first_output
# Attention for first block, since query_length = key_length.
first_output = first_block_attention()
# Concatenate two adjacent blocks to compute the overlapping memory block.
def local(x):
"""Helper function to get memory blocks."""
prev_block = mtf.slice(x, 0, num_blocks.size - 1, num_blocks.name)
cur_block = mtf.slice(x, 1, num_blocks.size - 1, num_blocks.name)
local_block = mtf.concat([prev_block, cur_block], mlength.name)
return local_block
local_k = local(k)
local_v = local(v)
# Calculate the causal mask to avoid peeking into the future. We compute
# this once and reuse it for all blocks since the block_size is known.
mlength = local_k.shape.dims[3]
mask = attention_bias_local_block(query_antecedent.mesh, blength,
mlength)
# Remove the first block from q since we already computed that.
tail_q = mtf.slice(q, 1, num_blocks.size - 1, num_blocks.name)
tail_output = dot_product_attention(tail_q,
local_k,
local_v,
mask=mask)
# Now concatenate the first and rest of the blocks.
final_output = mtf.concat([first_output, tail_output], num_blocks.name)
final_output = mtf.reshape(
final_output, mtf.Shape([batch, heads, query_length, kv_channels]))
return mtf.einsum([final_output, o_var],
mtf.Shape([batch, query_length, io_channels]))
def local_self_attention_spatial_blocks(query_antecedent,
kv_channels,
heads,
memory_w_dim=None,
mask_right=False,
name=None):
"""Attention to the source position and a neighborhood to the left or right.
The sequence is divided into blocks of length block_size.
Attention for a given query position can only see memory positions
less than or equal to the query position, in the corresponding block
and the previous block.
Args:
query_antecedent: a mtf.Tensor with shape
[batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]
must have the same size as query_length, but a different name.
kv_channels: a mtf.Dimension (the size of the key and value vectors)
heads: a mtf.Dimension (the number of heads)
memory_w_dim: mtf Dimension, for the memory width block.
mask_right: bool, flag specifying whether we mask out attention to the right
for the decoder.
name: an optional string.
Returns:
a Tensor of shape
[batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]
Raises:
ValueError: if channels or depth don't match.
"""
with tf.variable_scope(name,
default_name="multihead_attention",
values=[query_antecedent]):
w_dim, io_channels = query_antecedent.shape.dims[-2:]
batch, num_w_blocks = query_antecedent.shape.dims[:2]
q_var, k_var, v_var, o_var = multihead_attention_vars(
query_antecedent.mesh, heads, io_channels, kv_channels,
query_antecedent.dtype)
# Rename dimensions for the memory height and width.
memory_antecedent = mtf.rename_dimension(query_antecedent, w_dim.name,
memory_w_dim.name)
# Call einsum over the query and memory to get query q, keys k and values v.
q = mtf.einsum([query_antecedent, q_var],
mtf.Shape(
[batch, heads, num_w_blocks, w_dim, kv_channels]))
k = mtf.einsum([memory_antecedent, k_var],
mtf.Shape(
[batch, heads, num_w_blocks, w_dim, kv_channels]))
v = mtf.einsum([memory_antecedent, v_var],
mtf.Shape(
[batch, heads, num_w_blocks, w_dim, kv_channels]))
# Halo exchange for memory blocks.
if memory_w_dim is not None:
k, v = local_1d_halo_exchange(k, v, num_w_blocks, w_dim,
memory_w_dim, mask_right)
# Calculate the causal mask to avoid peeking into the future. We compute
# this once and reuse it for all blocks since the block_size is known.
mask = None
if mask_right:
mask = attention_bias_local_block(query_antecedent.mesh, w_dim,
memory_w_dim)
output = dot_product_attention(q, k, v, mask=mask)
return mtf.einsum([output, o_var],
mtf.Shape([batch, num_w_blocks, w_dim, io_channels]))
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)
"""
# 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)
alive_seq, alive_log_probs, _, second_selector = grow_alive(
top2k_seq, top2k_scores, top2k_log_probs, top2k_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])
new_states = [
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) for state in new_states
]
return (i + 1, alive_seq, alive_log_probs, finished_seq,
finished_scores, finished_flags) + tuple(new_states)
def masked_local_attention_1d(x,
kv_channels,
heads,
window_size=128,
master_dtype=tf.float32,
slice_dtype=tf.float32,
length_per_split=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
slice_dtype: a tf.dtype
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.
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:]
q_var, k_var, v_var, o_var = multihead_attention_vars(
x.mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype,
x.dtype)
# Get query q, keys k and values v.
qkv_shape = mtf.Shape(batch_dims + [heads, length, kv_channels])
q = mtf.einsum([x, q_var], qkv_shape)
k = mtf.einsum([x, k_var], qkv_shape)
v = mtf.einsum([x, v_var], qkv_shape)
# 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, o_var],
mtf.Shape(batch_dims + [length, io_channels]))
def multihead_self_attention_memory_compressed(x,
mask_right,
compression_factor,
kv_channels,
heads,
dropout=0.0,
dropout_broadcast_dims=None,
master_dtype=tf.float32,
slice_dtype=tf.float32,
name="multihead_attention"):
"""Memory-compressed self-attention.
The memory is first average-pooled (strided) to make it shorter by
a factor of compression_factor.
Args:
x: a mtf.Tensor with shape
[<batch_dims>, query_length, io_channels]
mask_right: a boolean
compression_factor: an integer
kv_channels: a mtf.Dimension (the size of the key and value vectors)
heads: a mtf.Dimension (the number of heads)
dropout: a floating point value
dropout_broadcast_dims: an optional list of mtf.Dimension
master_dtype: a tf.dtype
slice_dtype: a tf.dtype
name: an optional string.
Returns:
A mtf.Tensor with shape [batch, query_length, io_channels]
Raises:
ValueError: if the dimensions do not match.
"""
batch_dims = x.shape.dims[:-2]
length, io_channels = x.shape.dims[-2:]
with tf.variable_scope(name,
default_name="compressed_attention",
values=[x]):
q_var, k_var, v_var, o_var = multihead_attention_vars(
x.mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype,
x.dtype)
memory_antecedent = compress_mean(x, length, compression_factor)
memory_antecedent = rename_length_to_memory_length(memory_antecedent)
memory_length = memory_antecedent.shape.dims[-2]
q = mtf.einsum([x, q_var],
mtf.Shape(batch_dims + [heads, length, kv_channels]))
k = mtf.einsum([memory_antecedent, k_var],
mtf.Shape(batch_dims +
[heads, memory_length, kv_channels]))
v = mtf.einsum([memory_antecedent, v_var],
mtf.Shape(batch_dims +
[heads, memory_length, kv_channels]))
if mask_right:
query_pos = mtf.range(x.mesh, length, dtype=tf.int32)
memory_pos = (mtf.range(x.mesh, memory_length, dtype=tf.int32) *
compression_factor + (compression_factor - 1))
mask = mtf.cast(mtf.greater(memory_pos, query_pos), x.dtype) * -1e9
else:
mask = None
o = dot_product_attention(q,
k,
v,
mask,
dropout,
dropout_broadcast_dims,
extra_logit=0.0)
return mtf.einsum([o, o_var],
mtf.Shape(batch_dims + [length, io_channels]))