def apply_grad(self, grad, var):
"""See base class."""
if grad is None:
tf.logging.warning("Gradient is None for variable %s" % var.name)
return []
grad = mtf.to_float(grad)
assignments = []
m = mtf.get_variable(var.mesh,
var.name + "/adam_m",
var.shape,
initializer=tf.zeros_initializer(),
trainable=False)
v = mtf.get_variable(var.mesh,
var.name + "/adam_v",
var.shape,
initializer=tf.zeros_initializer(),
trainable=False)
# Standard Adam update.
next_m = self.beta_1 * m + (1.0 - self.beta_1) * grad
next_v = self.beta_2 * v + (1.0 - self.beta_2) * mtf.square(grad)
update = next_m / (mtf.sqrt(next_v) + self.epsilon)
# Just adding the square of the weights to the loss function is *not*
# the correct way of using L2 regularization/weight decay with Adam,
# since that will interact with the m and v parameters in strange ways.
#
# Instead we want ot decay the weights in a manner that doesn't interact
# with the m/v parameters. This is equivalent to adding the square
# of the weights to the loss with plain (non-momentum) SGD.
if self._do_use_weight_decay(var.name):
update += self.weight_decay_rate * var.value
update_with_lr = self.learning_rate * update
var_update = mtf.assign_sub(var, update_with_lr)
assignments.extend(
[var_update,
mtf.assign(m, next_m),
mtf.assign(v, next_v)])
return assignments
def apply_grad(self, grad, var):
if grad is None:
tf.logging.warning("Gradient is None for variable %s" % var)
return []
# create slots
grad = mtf.to_float(grad)
factored_dims = self._factored_dims(var.shape)
if factored_dims:
d0, d1 = factored_dims
vr_shape = var.shape - d0
vc_shape = var.shape - d1
vr = mtf.get_variable(var.mesh,
var.name + "_slot_vr",
vr_shape,
initializer=tf.zeros_initializer(),
trainable=False)
vc = mtf.get_variable(var.mesh,
var.name + "_slot_vc",
vc_shape,
initializer=tf.zeros_initializer(),
trainable=False)
else:
v = mtf.get_variable(var.mesh,
var.name + "_slot_v",
var.shape,
initializer=tf.zeros_initializer(),
trainable=False)
if self._beta1:
m = mtf.get_variable(var.mesh,
var.name + "_slot_m",
var.shape,
initializer=tf.zeros_initializer(),
trainable=False)
with tf.variable_scope(var.name + "/adafactor"):
grad_squared = mtf.square(grad) + self._epsilon1
decay_rate = self._decay_rate
old_val = mtf.to_float(var.value)
if self._multiply_by_parameter_scale:
update_scale = self._parameter_scale(
old_val) * self._learning_rate
else:
update_scale = self._learning_rate
mixing_rate = 1.0 - decay_rate
updates = []
if factored_dims:
grad_squared_row_mean = mtf.reduce_mean(grad_squared,
output_shape=vr_shape)
grad_squared_col_mean = mtf.reduce_mean(grad_squared,
output_shape=vc_shape)
new_vr = vr * decay_rate + grad_squared_row_mean * mixing_rate
new_vc = vc * decay_rate + grad_squared_col_mean * mixing_rate
vr_update = mtf.assign(vr, new_vr)
vc_update = mtf.assign(vc, new_vc)
updates.extend([vr_update, vc_update])
long_term_mean = mtf.reduce_mean(new_vr, reduced_dim=d1)
r_factor = mtf.rsqrt(new_vr / long_term_mean)
c_factor = mtf.rsqrt(new_vc)
x = grad * r_factor * c_factor
else:
new_v = v * decay_rate + grad_squared * mixing_rate
v_update = mtf.assign(v, new_v)
updates.append(v_update)
x = grad * mtf.rsqrt(new_v)
if self._clipping_threshold is not None:
clipping_denom = mtf.maximum(
1.0,
reduce_rms(x) / self._clipping_threshold)
x /= clipping_denom
subtrahend = x * update_scale
if self._beta1:
new_m = (m * tf.constant(self._beta1) +
subtrahend * tf.constant(1.0 - self._beta1))
subtrahend = new_m
updates.append(mtf.assign(m, new_m))
# It is critical to use assign_sub instead of mtf.assign(var - subtrahend)
# for the case of bfloat16 activations, so as to avoid repeatedly
# rounding the slice value, which results in poor quality.
var_update = mtf.assign_sub(var, subtrahend)
updates.append(var_update)
return updates
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)
