def test_transformer_relative_bias(self, use_relative_bias):
p = transformers.Transformer.Params().Set(name='jax_transformer_layer',
input_dims=32,
hidden_dims=128,
num_heads=8,
mask_self_attention=True,
packed_input=True,
cross_attention=False)
seq_len = np.random.randint(10, 32)
batch_size = 10
if use_relative_bias:
p.tr_atten_tpl.relative_bias_tpl = attentions.RelativeBias.Params(
).Set(relative_attention_num_buckets=2,
relative_attention_max_distance=8)
transformer_layer = p.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = transformer_layer.instantiate_variables(prng_key)
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, p.input_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
attention_mask = attentions.convert_paddings_to_mask(paddings)
causal_mask = attentions.causal_mask(inputs)
attention_mask = jnp.minimum(attention_mask, causal_mask)
segment_ids = np.random.random_integers(0, 2, [batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids, dtype=np.float32)
attention_mask = jnp.minimum(attention_mask, segment_mask)
if use_relative_bias:
segment_pos = np.random.randint(
0, seq_len, [batch_size, seq_len]).astype('int32')
segment_pos = jnp.asarray(segment_pos)
else:
segment_pos = None
with base_layer.JaxContext.new_context(
prng_key=prng_key,
global_step=jnp.array(0, dtype=jnp.uint32)) as jax_context:
jax_context.bind(transformer_layer,
transformer_layer.vars_to_flax_vars(initial_vars))
outputs, _ = transformer_layer.fprop(inputs,
paddings,
attention_mask=attention_mask,
segment_pos=segment_pos)
logging.info('initial_vars in transformer layer = %s', initial_vars)
np_outputs = test_utils.to_np(outputs)
logging.info('np_outputs: %s', np_outputs)
if use_relative_bias:
self.assertAlmostEqual(np_outputs[0, 0, 1], 0.79015386, places=5)
self.assertAlmostEqual(np_outputs[0, 1, 0], 0.48336178, places=5)
# Plumbing test.
self.assertAllClose(np_outputs, np_outputs, atol=1e-5)
def test_transformer_layer_norm_policies(self, norm_policy):
p = transformers.Transformer.Params().Set(name='jax_transformer_layer',
input_dims=32,
hidden_dims=128,
num_heads=8,
mask_self_attention=True,
packed_input=True,
cross_attention=False,
norm_policy=norm_policy)
seq_len = np.random.randint(10, 32)
batch_size = 10
transformer_layer = p.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = transformer_layer.instantiate_variables(prng_key)
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, p.input_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
attention_mask = attentions.convert_paddings_to_mask(paddings)
causal_mask = attentions.causal_mask(inputs)
attention_mask = jnp.minimum(attention_mask, causal_mask)
segment_ids = np.random.random_integers(0, 2, [batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids, dtype=np.float32)
attention_mask = jnp.minimum(attention_mask, segment_mask)
with base_layer.JaxContext.new_context(
prng_key=prng_key,
global_step=jnp.array(0, dtype=jnp.uint32)) as jax_context:
jax_context.bind(transformer_layer,
transformer_layer.vars_to_flax_vars(initial_vars))
outputs, _ = transformer_layer.fprop(inputs,
paddings,
attention_mask=attention_mask)
logging.info('initial_vars in transformer layer = %s', initial_vars)
np_outputs = test_utils.to_np(outputs)
# Plumbing test.
self.assertAllClose(np_outputs, np_outputs, atol=1e-5)
def test_stacked_transformer_layer_extendstep(self, packed_input,
cross_attention, combine_qkv,
dconv_qkv,
use_rotary_position_emb):
if cross_attention and combine_qkv:
self.skipTest(
'combine_qkv optimization only works for self-attention.')
layer_params = transformers.StackedTransformer.Params()
num_layers = 2
model_dims = 8
p = layer_params.Set(name='jax_transformer_layer',
model_dims=model_dims,
hidden_dims=32,
num_heads=2,
mask_self_attention=True,
packed_input=packed_input,
cross_attention=cross_attention,
num_layers=num_layers)
p.transformer_layer_params_tpl.tr_atten_tpl.combine_qkv = combine_qkv
p.transformer_layer_params_tpl.tr_atten_tpl.dconv_qkv = dconv_qkv
p.transformer_layer_params_tpl.tr_atten_tpl.use_rotary_position_emb = (
use_rotary_position_emb)
if cross_attention:
p.transformer_layer_params_tpl.cross_atten_tpl = (
p.transformer_layer_params_tpl.tr_atten_tpl.Copy())
# Cross attention should not have depth-wise convolution.
p.transformer_layer_params_tpl.cross_atten_tpl.dconv_qkv = False
# Cross attention should not have rotary position embedding.
p.transformer_layer_params_tpl.cross_atten_tpl.use_rotary_position_emb = (
False)
p_copy = p.Copy()
p_copy.num_layers = 1
p = transformers.StackedTransformerRepeated.Params()
p.name = 'jax_transformer_repeated_layer'
p.block = p_copy
p.x_times = num_layers
seq_len = 4
batch_size = 4
stacked_transformer_layer = p.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = stacked_transformer_layer.instantiate_variables(
prng_key)
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, model_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
attention_mask = attentions.convert_paddings_to_mask(paddings)
segment_mask = None
if packed_input:
segment_ids = np.random.random_integers(0, 2,
[batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids,
dtype=np.float32)
cross_inputs = None
cross_paddings = None
cross_segment_mask = None
if cross_attention:
cross_seq_len = np.random.randint(10, 32)
npy_cross_inputs = np.random.normal(
1.0, 0.5,
[batch_size, cross_seq_len, model_dims]).astype('float32')
cross_inputs = jnp.asarray(npy_cross_inputs)
npy_cross_paddings = np.random.randint(
0, 1, [batch_size, cross_seq_len]).astype('float32')
cross_paddings = jnp.asarray(npy_cross_paddings)
if packed_input:
source_segment_ids = np.random.random_integers(
0, 2, [batch_size, cross_seq_len])
cross_segment_mask = attentions.segment_mask(
segment_ids, source_segment_ids, dtype=np.float32)
prng_key = jax.random.PRNGKey(seed=123)
global_step = jnp.array(0, dtype=jnp.uint64)
with base_layer.JaxContext.new_context(
prng_key=prng_key, global_step=global_step) as jax_context:
jax_context.bind(
stacked_transformer_layer,
stacked_transformer_layer.vars_to_flax_vars(initial_vars))
fprop_outputs = stacked_transformer_layer.fprop(
inputs,
paddings,
segment_mask=segment_mask,
cross_inputs=cross_inputs,
cross_paddings=cross_paddings,
cross_segment_mask=cross_segment_mask)
decoder_outputs = jnp.zeros(
shape=[seq_len, batch_size, model_dims])
initial_states = stacked_transformer_layer.init_states(
batch_size, seq_len)
atten_states = initial_states
for t in range(seq_len):
segment_mask_t = attention_mask[:, :, t, :]
cross_segment_mask_t = cross_segment_mask
if segment_mask is not None:
segment_mask_t = jnp.minimum(segment_mask_t,
segment_mask[:, :, t, :])
if cross_segment_mask is not None:
cross_segment_mask_t = cross_segment_mask[:, :, t, :]
atten_states, encoded = stacked_transformer_layer.extend_step(
atten_states,
inputs=inputs[:, t, :],
time_step=t,
segment_mask=segment_mask_t,
cross_inputs=cross_inputs,
cross_paddings=cross_paddings,
cross_segment_mask=cross_segment_mask_t)
decoder_outputs = decoder_outputs.at[t].set(encoded)
decoder_out_transposed = jnp.transpose(decoder_outputs, [1, 0, 2])
# TODO(lepikhin): remove noisy test logging
# logging.info('initial_vars in transformer layer = %s', initial_vars)
np_fprop_outputs = test_utils.to_np(fprop_outputs)
np_decoder_outputs = test_utils.to_np(decoder_out_transposed)
self.assertAllClose(np_fprop_outputs, np_decoder_outputs, atol=1e-5)
def test_stacked_transformer_layer(self, mask_self_attention, packed_input,
cross_attention):
p = transformers.StackedTransformer.Params().Set(
name='jax_stacked_transformer_layer',
model_dims=16,
hidden_dims=64,
num_heads=8,
mask_self_attention=mask_self_attention,
num_layers=4,
packed_input=packed_input,
cross_attention=cross_attention)
seq_len = np.random.randint(10, 32)
batch_size = 10
stacked_transformer_layer = p.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = stacked_transformer_layer.instantiate_variables(
prng_key)
# test conversion between vars and flax vars.
pax_vars = stacked_transformer_layer.vars
flax_vars = stacked_transformer_layer.flax_vars
tf.nest.assert_same_structure(
pax_vars, stacked_transformer_layer.flax_vars_to_vars(flax_vars))
tf.nest.assert_same_structure(
flax_vars, stacked_transformer_layer.vars_to_flax_vars(pax_vars))
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, p.model_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
segment_mask = None
tf_segment_mask = None
if packed_input:
segment_ids = np.random.random_integers(0, 2,
[batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids,
dtype=np.float32)
if mask_self_attention:
tf_segment_mask = batch_major_attention.CausalSegmentMask(
segment_ids, tf.float32)
else:
tf_segment_mask = batch_major_attention.SegmentMask(
segment_ids, segment_ids)
cross_inputs = None
cross_paddings = None
cross_segment_mask = None
tf_cross_inputs = None
tf_cross_paddings = None
tf_cross_segment_mask = None
if cross_attention:
cross_seq_len = np.random.randint(10, 64)
npy_cross_inputs = np.random.normal(
1.0, 0.5,
[batch_size, cross_seq_len, p.model_dims]).astype('float32')
cross_inputs = jnp.asarray(npy_cross_inputs)
tf_cross_inputs = tf.constant(npy_cross_inputs, dtype=tf.float32)
npy_cross_paddings = np.random.randint(
0, 1, [batch_size, cross_seq_len]).astype('float32')
cross_paddings = jnp.asarray(npy_cross_paddings)
tf_cross_paddings = tf.constant(npy_cross_paddings,
dtype=tf.float32)
if packed_input:
source_segment_ids = np.random.random_integers(
0, 2, [batch_size, cross_seq_len])
cross_segment_mask = attentions.segment_mask(
segment_ids, source_segment_ids, dtype=np.float32)
tf_cross_segment_mask = batch_major_attention.SegmentMask(
segment_ids, source_segment_ids)
outputs = test_utils.apply(stacked_transformer_layer,
initial_vars,
stacked_transformer_layer.fprop,
inputs,
paddings,
context_p=None,
segment_mask=segment_mask,
cross_inputs=cross_inputs,
cross_paddings=cross_paddings,
cross_segment_mask=cross_segment_mask)
logging.info('initial_vars in transformer layer = %s', initial_vars)
# Test whether tf Transformer layer returns same output
# Modify initial_vars to use TF compatible params
tf_initial_vars = py_utils.NestedMap()
tf_initial_vars.x_layers = []
for jax_initial_vars in initial_vars.x_layers:
tf_layer_vars = test_utils.replace_jax_attention_vars_to_tf(
jax_initial_vars, cross_attention)
tf_initial_vars.x_layers.append(tf_layer_vars)
tf_initial_vars = test_utils.to_tf_nmap(tf_initial_vars)
logging.info('tf_initial_vars in transformer layer = %s', initial_vars)
tf_p = batch_major_attention.StackedTransformerLayers.Params().Set(
name='tf_transformer_layer',
mdl_dim=p.model_dims,
hidden_dim=p.hidden_dims,
num_atten_heads=p.num_heads,
mask_self_atten=mask_self_attention,
num_layers=p.num_layers,
packed_input=packed_input,
has_aux_atten=cross_attention)
tf_p.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.batch_norm = (
False)
tf_p.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.has_bias = True
tf_stacked_transformer_layer = tf_p.Instantiate()
tf_output, _ = tf_stacked_transformer_layer.FProp(
tf_initial_vars,
test_utils.to_tf_nmap(npy_inputs),
paddings=test_utils.to_tf_nmap(npy_paddings),
segment_mask=test_utils.to_tf_nmap(tf_segment_mask),
aux_vec=test_utils.to_tf_nmap(tf_cross_inputs),
aux_paddings=test_utils.to_tf_nmap(tf_cross_paddings),
aux_segment_mask=test_utils.to_tf_nmap(tf_cross_segment_mask))
np_outputs = test_utils.to_np(outputs)
tf_np_outputs = test_utils.to_np(tf_output)
self.assertAllClose(tf_np_outputs, np_outputs, atol=1e-5)
def test_repeated_stacked_xformer_layer(self, mask_self_attention,
packed_input, cross_attention):
model_dims = 16
p1 = transformers.StackedTransformer.Params().Set(
name='jax_stacked_transformer_layer',
model_dims=model_dims,
hidden_dims=64,
num_heads=8,
mask_self_attention=mask_self_attention,
num_layers=4,
packed_input=packed_input,
cross_attention=cross_attention)
p1_one_layer = p1.Copy()
p1_one_layer.num_layers = 1
p2 = transformers.StackedTransformerRepeated.Params().Set(
name='jax_stacked_transformer_layer_repeated',
block=p1_one_layer,
x_times=p1.num_layers)
seq_len = np.random.randint(10, 32)
batch_size = 10
stacked_transformer_layer = p1.Instantiate()
repeated_transformer_layer = p2.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = stacked_transformer_layer.instantiate_variables(
prng_key)
repeated_transformer_layer.instantiate_variable_configs()
def _stack_vars(*args):
args = [x[jnp.newaxis, :] for x in args]
return jnp.vstack(args)
stacked_vars = tf.nest.map_structure(_stack_vars,
*initial_vars.x_layers)
repeated_vars = py_utils.NestedMap(repeat=py_utils.NestedMap(
sub=py_utils.NestedMap(x_layers=[stacked_vars])))
tf.nest.assert_same_structure(
repeated_vars,
repeated_transformer_layer.instantiate_variables(prng_key))
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, model_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
segment_mask = None
if packed_input:
segment_ids = np.random.random_integers(0, 2,
[batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids,
dtype=np.float32)
cross_inputs = None
cross_paddings = None
cross_segment_mask = None
if cross_attention:
cross_seq_len = np.random.randint(10, 64)
npy_cross_inputs = np.random.normal(
1.0, 0.5,
[batch_size, cross_seq_len, model_dims]).astype('float32')
cross_inputs = jnp.asarray(npy_cross_inputs)
npy_cross_paddings = np.random.randint(
0, 1, [batch_size, cross_seq_len]).astype('float32')
cross_paddings = jnp.asarray(npy_cross_paddings)
if packed_input:
source_segment_ids = np.random.random_integers(
0, 2, [batch_size, cross_seq_len])
cross_segment_mask = attentions.segment_mask(
segment_ids, source_segment_ids, dtype=np.float32)
outputs = test_utils.apply(stacked_transformer_layer,
initial_vars,
stacked_transformer_layer.fprop,
inputs,
paddings,
context_p=None,
segment_mask=segment_mask,
cross_inputs=cross_inputs,
cross_paddings=cross_paddings,
cross_segment_mask=cross_segment_mask)
outputs_repeated = test_utils.apply(
repeated_transformer_layer,
repeated_vars,
repeated_transformer_layer.fprop,
inputs,
paddings,
context_p=None,
segment_mask=segment_mask,
cross_inputs=cross_inputs,
cross_paddings=cross_paddings,
cross_segment_mask=cross_segment_mask)
self.assertAllClose(outputs, outputs_repeated, atol=1e-5)
def test_transformer_moe_dense_layer(self, mask_self_attention,
packed_input, cross_attention):
# Comparing scan over blocks of layers and regular loop
block_p = transformers.StackedTransformer.Params().Set(
name='transformer_block',
num_layers=2,
model_dims=3,
hidden_dims=6,
num_heads=1,
mask_self_attention=mask_self_attention,
packed_input=packed_input,
cross_attention=cross_attention,
num_experts=4,
num_groups=1,
moe_layers=[0])
block_p_repeated = transformers.StackedTransformerRepeated.Params(
).Set(name='stacked_transformer_layer_repeated',
block=block_p.Copy(),
x_times=1)
stack_p = transformers.StackedTransformer.Params().Set(
name='transformer_stack',
num_layers=2, # moe + dense
model_dims=block_p.model_dims,
hidden_dims=block_p.hidden_dims,
num_heads=block_p.num_heads,
mask_self_attention=block_p.mask_self_attention,
packed_input=block_p.packed_input,
cross_attention=block_p.cross_attention,
num_experts=block_p.num_experts,
num_groups=block_p.num_groups,
moe_layers=[0])
moe_p = stack_p.moe_layer_tpl
moe_p.expert_capacity_dim = 2
moe_p.expert_capacity_factor = 0
moe_p = block_p.moe_layer_tpl
moe_p.expert_capacity_dim = 2
moe_p.expert_capacity_factor = 0
transformer_block = block_p_repeated.Instantiate()
transformer_stack = stack_p.Instantiate()
seq_len = 4
batch_size = 3
prng_key = jax.random.PRNGKey(seed=123)
block_initial_vars = transformer_block.instantiate_variables(prng_key)
stack_initial_vars = transformer_stack.instantiate_variables(prng_key)
npy_inputs = np.random.normal(
1.0, 0.5,
[batch_size, seq_len, block_p.model_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
segment_mask = None
if packed_input:
segment_ids = np.random.random_integers(0, 2,
[batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids,
dtype=np.float32)
cross_inputs = None
cross_paddings = None
cross_segment_mask = None
if cross_attention:
cross_seq_len = np.random.randint(10, 64)
npy_cross_inputs = np.random.normal(
1.0, 0.5, [batch_size, cross_seq_len, block_p.model_dims
]).astype('float32')
cross_inputs = jnp.asarray(npy_cross_inputs)
npy_cross_paddings = np.random.randint(
0, 1, [batch_size, cross_seq_len]).astype('float32')
cross_paddings = jnp.asarray(npy_cross_paddings)
if packed_input:
source_segment_ids = np.random.random_integers(
0, 2, [batch_size, cross_seq_len])
cross_segment_mask = attentions.segment_mask(
segment_ids, source_segment_ids, dtype=np.float32)
block_outputs = test_utils.apply(transformer_block,
block_initial_vars,
transformer_block.fprop,
inputs,
paddings,
segment_mask=segment_mask,
cross_inputs=cross_inputs,
cross_paddings=cross_paddings,
cross_segment_mask=cross_segment_mask)
stack_outputs = test_utils.apply(transformer_stack,
stack_initial_vars,
transformer_stack.fprop,
inputs,
paddings,
segment_mask=segment_mask,
cross_inputs=cross_inputs,
cross_paddings=cross_paddings,
cross_segment_mask=cross_segment_mask)
_ = test_utils.to_np(block_outputs)
_ = test_utils.to_np(stack_outputs)
def test_transformer_layer(self, mask_self_attention, packed_input,
cross_attention):
p = transformers.Transformer.Params().Set(
name='jax_transformer_layer',
input_dims=32,
hidden_dims=128,
num_heads=8,
mask_self_attention=mask_self_attention,
packed_input=packed_input,
cross_attention=cross_attention)
seq_len = np.random.randint(10, 32)
batch_size = 10
transformer_layer = p.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = transformer_layer.instantiate_variables(prng_key)
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, p.input_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
causal_mask = None
segment_mask = None
tf_segment_mask = None
attention_mask = attentions.convert_paddings_to_mask(paddings)
if mask_self_attention:
causal_mask = attentions.causal_mask(inputs)
attention_mask = jnp.minimum(attention_mask, causal_mask)
if packed_input:
segment_ids = np.random.random_integers(0, 2,
[batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids,
dtype=np.float32)
attention_mask = jnp.minimum(attention_mask, segment_mask)
if mask_self_attention:
tf_segment_mask = batch_major_attention.CausalSegmentMask(
segment_ids, tf.float32)
else:
tf_segment_mask = batch_major_attention.SegmentMask(
segment_ids, segment_ids)
cross_inputs = None
cross_attention_mask = None
tf_cross_inputs = None
tf_cross_paddings = None
tf_cross_segment_mask = None
if cross_attention:
cross_seq_len = np.random.randint(10, 128)
npy_cross_inputs = np.random.normal(
1.0, 0.5,
[batch_size, cross_seq_len, p.input_dims]).astype('float32')
cross_inputs = jnp.asarray(npy_cross_inputs)
tf_cross_inputs = tf.constant(npy_cross_inputs, dtype=tf.float32)
npy_cross_paddings = np.random.randint(
0, 1, [batch_size, cross_seq_len]).astype('float32')
cross_paddings = jnp.asarray(npy_cross_paddings)
cross_attention_mask = attentions.convert_paddings_to_mask(
cross_paddings)
tf_cross_paddings = tf.constant(npy_cross_paddings,
dtype=tf.float32)
if packed_input:
source_segment_ids = np.random.random_integers(
0, 2, [batch_size, cross_seq_len])
cross_segment_mask = attentions.segment_mask(
segment_ids, source_segment_ids, dtype=np.float32)
cross_attention_mask = jnp.minimum(cross_attention_mask,
cross_segment_mask)
tf_cross_segment_mask = batch_major_attention.SegmentMask(
segment_ids, source_segment_ids)
outputs, _ = test_utils.apply(
transformer_layer,
initial_vars,
transformer_layer.fprop,
inputs,
paddings,
context_p=None,
attention_mask=attention_mask,
cross_inputs=cross_inputs,
cross_attention_mask=cross_attention_mask)
logging.info('initial_vars in transformer layer = %s', initial_vars)
# Test whether tf Transformer layer returns same output
# Modify initial_vars to use TF compatible params
tf_initial_vars = test_utils.replace_jax_attention_vars_to_tf(
initial_vars, cross_attention)
tf_initial_vars = test_utils.to_tf_nmap(tf_initial_vars)
logging.info('tf_initial_vars in transformer layer = %s', initial_vars)
tf_p = batch_major_attention.TransformerLayer.Params().Set(
name='tf_transformer_layer',
input_dim=p.input_dims,
num_heads=p.num_heads,
mask_self_atten=mask_self_attention,
packed_input=packed_input,
has_aux_atten=cross_attention)
tf_p.tr_fflayer_tpl.hidden_dim = p.hidden_dims
tf_p.tr_fflayer_tpl.fflayer_tpl.batch_norm = False
tf_p.tr_fflayer_tpl.fflayer_tpl.has_bias = True
tf_transformer_layer = tf_p.Instantiate()
tf_output, _ = tf_transformer_layer.FProp(
tf_initial_vars,
tf.constant(npy_inputs, dtype=tf.float32),
paddings=test_utils.to_tf_nmap(npy_paddings),
segment_mask=tf_segment_mask,
aux_vec=tf_cross_inputs,
aux_paddings=tf_cross_paddings,
aux_segment_mask=test_utils.to_tf_nmap(tf_cross_segment_mask))
np_outputs = test_utils.to_np(outputs)
tf_np_outputs = test_utils.to_np(tf_output)
self.assertAllClose(tf_np_outputs, np_outputs, atol=1e-5)
def test_transformer_layer_cross_attention_ln(self, packed_input):
p = transformers.Transformer.Params().Set(name='jax_transformer_layer',
input_dims=8,
hidden_dims=32,
num_heads=4,
mask_self_attention=True,
packed_input=packed_input,
cross_attention=True)
seq_len = 5
batch_size = 4
transformer_layer = p.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = transformer_layer.instantiate_variables(prng_key)
# Change the self attention initial vars.
initial_vars.layer_norm.scale = 0.5
initial_vars.layer_norm.bias = 5.0
# Change the cross attention initial vars.
initial_vars.cross_layer_norm.scale = 15
initial_vars.cross_layer_norm.bias = 1.5
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, p.input_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
paddings = jnp.asarray(npy_paddings)
attention_mask = attentions.convert_paddings_to_mask(paddings)
causal_mask = attentions.causal_mask(inputs)
attention_mask = jnp.minimum(causal_mask, attention_mask)
if packed_input:
segment_ids = np.random.random_integers(0, 2,
[batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids,
dtype=np.float32)
attention_mask = jnp.minimum(attention_mask, segment_mask)
with base_layer.JaxContext.new_context(
prng_key=prng_key,
global_step=jnp.array(0, dtype=jnp.uint32)) as jax_context:
jax_context.bind(transformer_layer,
transformer_layer.vars_to_flax_vars(initial_vars))
inputs_normalized = transformer_layer.layer_norm.fprop(inputs)
# Compute self-attention, key/value vectors are the input itself
atten_output, _ = transformer_layer.self_attention.fprop(
inputs_normalized,
inputs_normalized,
inputs_normalized,
atten_mask=attention_mask)
# Residual dropout and connection.
atten_output = transformer_layer.residual_dropout.fprop(
atten_output)
atten_output += inputs
# Normalize atten outputs using cross attention.
atten_output_normalized = transformer_layer.cross_layer_norm.fprop(
atten_output)
inputs_normalized = test_utils.to_np(inputs_normalized)
atten_output_normalized = test_utils.to_np(atten_output_normalized)
self.assertAllClose(initial_vars.layer_norm.bias,
inputs_normalized.mean(),
atol=1e-3)
self.assertAllClose((1.0 + initial_vars.layer_norm.scale)**2,
np.var(inputs_normalized),
atol=5e-3)
self.assertAllClose(initial_vars.cross_layer_norm.bias,
atten_output_normalized.mean(),
atol=1e-3)
self.assertAllClose((1.0 + initial_vars.cross_layer_norm.scale)**2,
np.var(atten_output_normalized),
atol=5e-3)
def test_transformer_layer_extendstep(self, packed_input, cross_attention,
dconv_qkv, use_rotary_position_emb):
p = transformers.Transformer.Params().Set(
name='jax_transformer_layer',
input_dims=8,
hidden_dims=32,
num_heads=4,
mask_self_attention=True,
packed_input=packed_input,
cross_attention=cross_attention)
p.tr_atten_tpl.dconv_qkv = dconv_qkv
p.tr_atten_tpl.use_rotary_position_emb = use_rotary_position_emb
if cross_attention:
p.cross_atten_tpl = p.tr_atten_tpl.Copy()
# Cross attention should not have depth-wise convolution.
p.cross_atten_tpl.dconv_qkv = False
# Cross attention should not have rotary position embedding.
p.cross_atten_tpl.use_rotary_position_emb = False
p.tr_atten_tpl.dconv_kernel_size = 2
seq_len = 4
batch_size = 4
transformer_layer = p.Instantiate()
prng_key = jax.random.PRNGKey(seed=123)
initial_vars = transformer_layer.instantiate_variables(prng_key)
initial_states = transformer_layer.init_states(batch_size, seq_len)
npy_inputs = np.random.normal(
1.0, 0.5, [batch_size, seq_len, p.input_dims]).astype('float32')
inputs = jnp.asarray(npy_inputs)
npy_paddings = np.random.randint(
0, 1, [batch_size, seq_len]).astype('float32')
# npy_paddings = np.zeros([batch_size, seq_len])
paddings = jnp.asarray(npy_paddings)
attention_mask = attentions.convert_paddings_to_mask(paddings)
segment_mask = None
causal_mask = attentions.causal_mask(inputs)
attention_mask = jnp.minimum(causal_mask, attention_mask)
if packed_input:
segment_ids = np.random.random_integers(0, 2,
[batch_size, seq_len])
segment_mask = attentions.segment_mask(segment_ids,
dtype=np.float32)
attention_mask = jnp.minimum(attention_mask, segment_mask)
cross_inputs = None
cross_paddings = None
cross_attention_mask = None
if cross_attention:
cross_seq_len = np.random.randint(10, 32)
npy_cross_inputs = np.random.normal(
1.0, 0.5,
[batch_size, cross_seq_len, p.input_dims]).astype('float32')
cross_inputs = jnp.asarray(npy_cross_inputs)
npy_cross_paddings = np.random.randint(
0, 1, [batch_size, cross_seq_len]).astype('float32')
cross_paddings = jnp.asarray(npy_cross_paddings)
cross_attention_mask = attentions.convert_paddings_to_mask(
cross_paddings)
if packed_input:
source_segment_ids = np.random.random_integers(
0, 2, [batch_size, cross_seq_len])
cross_segment_mask = attentions.segment_mask(
segment_ids, source_segment_ids, dtype=np.float32)
cross_attention_mask = jnp.minimum(cross_attention_mask,
cross_segment_mask)
with base_layer.JaxContext.new_context(
prng_key=prng_key,
global_step=jnp.array(0, dtype=jnp.uint32)) as jax_context:
jax_context.bind(transformer_layer,
transformer_layer.vars_to_flax_vars(initial_vars))
fprop_outputs, _ = transformer_layer.fprop(
inputs,
paddings,
attention_mask=attention_mask,
cross_inputs=cross_inputs,
cross_attention_mask=cross_attention_mask)
decoder_outputs = jnp.zeros(
shape=[seq_len, batch_size, p.input_dims])
atten_states = initial_states
for t in range(seq_len):
attention_mask_t = attention_mask[:, :, t, :]
cross_attention_mask_t = cross_attention_mask
if cross_attention:
cross_attention_mask_t = cross_attention_mask[:, :, t, :]
cross_attention_mask_t = np.expand_dims(
cross_attention_mask_t, axis=2)
atten_states, encoded = transformer_layer.extend_step(
atten_states,
inputs=inputs[:, t, :],
time_step=t,
attention_mask=attention_mask_t,
cross_inputs=cross_inputs,
cross_attention_mask=cross_attention_mask_t)
decoder_outputs = decoder_outputs.at[t].set(encoded)
decoder_out_transposed = jnp.transpose(decoder_outputs, [1, 0, 2])
logging.info('initial_vars in transformer layer = %s', initial_vars)
np_fprop_outputs = test_utils.to_np(fprop_outputs)
np_decoder_outputs = test_utils.to_np(decoder_out_transposed)
self.assertAllClose(np_fprop_outputs, np_decoder_outputs, atol=1e-5)
def fprop(
self,
inputs: JTensor,
input_paddings: JTensor,
targets: JTensor,
target_paddings: JTensor,
labels: Optional[NestedMap] = None,
input_segment_ids: Optional[JTensor] = None,
input_segment_pos: Optional[JTensor] = None,
target_segment_ids: Optional[JTensor] = None,
target_segment_pos: Optional[JTensor] = None,
) -> NestedMap:
"""Computes xent loss given the sequence model inputs.
Args:
inputs: Input ids. An int32 JTensor of shape [B, S].
input_paddings: A 0/1 JTensor of shape [B, S] with 1 denoting padding
correspdonding to the input sequence.
targets: Target ids. An int32 JTensor of shape [B, T].
target_paddings: A 0/1 JTensor of shape [B, T] with 1 denoting padding
corresponding to the target sequence.
labels: A `.NestedMap` containing the following fields: class_weights, a
JTensor with shape [batch, seqlen] containing weights for each target
word. class_ids, a JTensor with shape [B, T] of int32 dtype containing
the target class labels. class_probabilities, a JTensor with shape [B,
T, V] of float values indicating class-membership probabilities.
input_segment_ids: A JTensor of shape [B,S]. The segment that each input
token belongs to.
input_segment_pos: A JTensor of shape [B, S]. The position of each input
token within a segment.
target_segment_ids: A JTensor of shape [B,T]. The segment that each target
token belongs to.
target_segment_pos: A JTensor of shape [B, T]. The position of each target
token within a segment.
Returns:
Returns xent_output, where
`xent_output` is a `.NestedMap` as defined by `SoftmaxLayer`'s return. In
addition, per_sequence_xent is added which equal to the sum of xent loss
for tokens in a sequence.
"""
# Get the input embeddings.
p = self.params
batch, seq_length = inputs.shape
_, target_seq_length = targets.shape
encoder_output = self._encode(inputs, input_paddings,
input_segment_ids, input_segment_pos)
if p.decoder_embedding_tpl is not None:
# Targets have separate embedding params.
target_emb = self.decoder_embedding_lookup.fprop(targets)
else:
# Embedding parameters are shared with targets and softmax.
target_emb = self.softmax.emb_lookup(targets)
if p.decoder_ngrammer_tpl is not None:
target_emb = self.decoder_ngrammer.fprop(
input_ids=targets,
input_embs=target_emb,
paddings=target_paddings,
segment_pos=target_segment_pos)
if p.position_emb_tpl is not None:
targets_position_emb = self.position_emb.fprop(
seq_length=target_seq_length, position=target_segment_pos)
target_emb += targets_position_emb
if input_segment_ids is None:
assert input_segment_pos is None
# Fold the paddings with the segment mask.
input_segment_ids = jnp.asarray(1 - input_paddings, jnp.int32)
input_segment_pos = jnp.tile(
jnp.arange(seq_length, dtype=jnp.int32)[None, :], [batch, 1])
if target_segment_ids is None:
assert target_segment_pos is None
# Fold the paddings with the segment mask.
target_segment_ids = jnp.asarray(1 - target_paddings, jnp.int32)
target_segment_pos = jnp.tile(
jnp.arange(target_seq_length, dtype=jnp.int32)[None, :],
[batch, 1])
# Cross attention.
cross_segment_mask = attentions.segment_mask(target_segment_ids,
input_segment_ids,
target_emb.dtype)
target_segment_mask = attentions.causal_segment_mask(
target_segment_ids, target_emb.dtype)
output = self.decoder.fprop(target_emb,
target_paddings,
target_segment_mask,
cross_inputs=encoder_output,
cross_paddings=input_paddings,
cross_segment_mask=cross_segment_mask)
# Final layer norm for decoder.
output = self.decoder_ln.fprop(output)
return self.compute_loss(output, labels)
def _encode(self,
inputs: JTensor,
input_paddings: JTensor,
input_segment_ids: Optional[JTensor] = None,
input_segment_pos: Optional[JTensor] = None) -> JTensor:
"""Apply the Transformer encoder to the source sequence.
Args:
inputs: Input ids. An int32 JTensor of shape [B, S].
input_paddings: A 0/1 JTensor of shape [B, S] with 1 denoting padding
correspdonding to the input sequence.
input_segment_ids: A JTensor of shape [B,S]. The segment that each input
token belongs to.
input_segment_pos: A JTensor of shape [B, S]. The position of each input
token within a segment.
Returns:
The encoded sequence after applying the Transformer encoder.
"""
p = self.params
batch, seq_length = inputs.shape
if p.encoder_embedding_tpl is not None:
# Encoder has its own embedding lookup table for source ids.
input_emb = self.encoder_embedding_lookup.fprop(inputs)
elif p.decoder_embedding_tpl is not None:
# Encoder shares the same embedding as the target ids.
# The embedding lookup for target ids is separate from the softmax.
input_emb = self.decoder_embedding_lookup.fprop(inputs)
else:
# Encoder and decoder share the softmax and embedding params.
input_emb = self.softmax.emb_lookup(inputs)
if input_segment_ids is None:
assert input_segment_pos is None
# Fold the paddings with the segment mask.
input_segment_ids = jnp.asarray(1 - input_paddings, jnp.int32)
input_segment_pos = jnp.tile(
jnp.arange(seq_length, dtype=jnp.int32)[None, :], [batch, 1])
assert input_segment_ids is not None
assert input_segment_pos is not None
# Add NGrammer to the source embeddings.
if p.encoder_ngrammer_tpl is not None:
input_emb = self.encoder_ngrammer.fprop(
input_ids=inputs,
input_embs=input_emb,
paddings=input_paddings,
segment_pos=input_segment_pos)
if p.position_emb_tpl is not None:
position_emb = self.position_emb.fprop(seq_length=seq_length,
position=input_segment_pos)
input_emb += position_emb
inputs_segment_mask = attentions.segment_mask(input_segment_ids,
dtype=input_emb.dtype)
encoder_output = self.encoder.fprop(input_emb,
input_paddings,
segment_mask=inputs_segment_mask)
# Final layer norm for encoder output.
encoder_output = self.encoder_ln.fprop(encoder_output)
return encoder_output
def fprop(self,
inputs: JTensor,
paddings: JTensor,
labels: Optional[NestedMap] = None,
segment_ids: Optional[JTensor] = None,
segment_pos: Optional[JTensor] = None) -> NestedMap:
"""Computes xent loss given the language model inputs.
Args:
inputs: Input ids. An int32 JTensor of shape [B, T].
paddings: A 0/1 JTensor of shape [B, T] with 1 denoting padding.
labels: A `.NestedMap` containing the following fields: class_weights, a
JTensor with shape [batch, seqlen] containing weights for each target
word. class_ids, a JTensor with shape [B, T] of int32 dtype containing
the target class labels. class_probabilities, a JTensor with shape [B,
T, V] of float values indicating class-membership probabilities.
segment_ids: A JTensor of shape [B, T]. The segment that each token
belongs to.
segment_pos: A JTensor of shape [B, T]. The position of each token in a
segment.
Returns:
Returns xent_output, where
`xent_output` is a `.NestedMap` as defined by `SoftmaxLayer`'s return. In
addition, per_sequence_xent is added which equal to the sum of xent loss
for tokens in a sequence.
"""
p = self.params
# reentrant=True, to enable scan-local context override.
with py_utils.AuxLossContext(reentrant=True) as aux_loss_ctx:
assert aux_loss_ctx is not None
# Get the input embeddings.
if self.params.separate_embedding_tpl is not None:
input_emb = self.embedding_lookup.fprop(inputs)
else:
input_emb = self.softmax.emb_lookup(inputs)
batch, seq_length = inputs.shape
if segment_ids is None:
assert segment_pos is None
# Fold the paddings with the segment mask
segment_ids = jnp.asarray(1 - paddings, jnp.int32)
segment_pos = jnp.tile(
jnp.arange(seq_length, dtype=jnp.int32)[None, :],
[batch, 1])
# Add NGrammer to the source embeddings.
if p.ngrammer_tpl is not None:
input_emb = self.ngrammer.fprop(input_ids=inputs,
input_embs=input_emb,
paddings=paddings,
segment_pos=segment_pos)
if p.position_emb_tpl is not None:
position_emb = self.position_emb.fprop(seq_length=seq_length,
position=segment_pos)
inputs = input_emb + position_emb
else:
inputs = input_emb
if p.masked_lm:
segment_mask = attentions.segment_mask(segment_ids,
segment_ids,
inputs.dtype)
else:
segment_mask = attentions.causal_segment_mask(
segment_ids, inputs.dtype)
output = self.transformer.fprop(inputs,
paddings,
segment_mask=segment_mask,
segment_pos=segment_pos)
# Final layer norm
if p.final_ln_tpl is not None:
output = self.final_ln.fprop(output)
return self.compute_loss(output, labels)