def _enc_dec_concat():
"""Layers to merge encoder and decoder."""
if reversible_encoder:
return [
tl.ReversibleSelect([0, 1, 4, 2,
3]), # v_e v_d mask_e tok_e tok_d
t2.ConcatWithPadding2(mode=mode), # v_ed v_ed tok_e tok_d
]
else:
return [
tl.ReversibleSelect([0, 3, 1,
2]), # v_e v_d mask_e tok_e tok_d
t2.ConcatWithPadding(mode=mode), # v_ed tok_e tok_d
tl.ReversibleSelect([0, 0]), # v_ed v_ed tok_e tok_d
]
def _Encoder(): # vec_e mask_e tok_e tok_d tok_d
layers = [
tl.ReversibleSelect([0, 0]),
_ReversibleSerialForget(
[_EncoderBlock() for _ in range(n_encoder_layers)], d_model1,
n_layers_forget, forget_dense)
]
if not reversible_encoder:
layers += [
_XYAvg(),
tl.Dense(d_model1, use_bfloat16=use_bfloat16),
tl.LayerNorm(),
]
if mode == 'predict':
return tl.Cache(tl.Serial(layers))
else:
return tl.Serial(layers)
def Reformer2(input_vocab_size,
output_vocab_size=None,
d_model=512,
d_ff=2048,
d_attention_key=None,
d_attention_value=None,
n_encoder_layers=6,
n_decoder_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
encoder_attention_type=tl.SelfAttention,
encoder_decoder_attention_type=tl.SelfAttention,
axial_pos_shape='fixed-base',
d_axial_pos_embs=None,
ff_activation=tl.Relu,
ff_use_sru=0,
ff_chunk_size=0,
ff_dropout=None,
ff_sparsity=0,
loss_sparsity_type='mult',
loss_sparsity=0,
loss_d_lowrank=0,
loss_sparsity_prob=None,
attention_chunk_size=0,
n_layers_forget=0,
n_decoder_attention_layers=2,
use_bfloat16=False,
reversible_encoder=False,
mode='train'):
"""Reversible transformer encoder-decoder model.
This model expects an input pair: source, target.
At the moment, this model supports dot-product attention only. For the
attention types in the Reformer paper, see ReformerLM.
Args:
input_vocab_size: int: vocab size of the source.
output_vocab_size: int (optional): vocab size of the target. If None, the
source and target are assumed to have the same vocab.
d_model: int: depth of embedding
d_ff: int: depth of feed-forward layer
d_attention_key: int: depth of key vector for each attention head
d_attention_value: int: depth of value vector for each attention head
n_encoder_layers: int: number of encoder layers
n_decoder_layers: int: number of decoder layers
n_heads: int: number of attention heads
dropout: float: dropout rate (how much to drop out)
max_len: int: maximum symbol length for positional encoding
encoder_attention_type: class: attention class to use, such as SelfAttention
encoder_decoder_attention_type: class: attention class to use, such as
SelfAttention
axial_pos_shape: tuple of ints: input shape to use for the axial position
encoding. If unset, axial position encoding is disabled.
d_axial_pos_embs: tuple of ints: depth of position embedding for each axis.
Tuple length must match axial_pos_shape, and values must sum to d_model.
ff_activation: the non-linearity in feed-forward layer
ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
ff_dropout: float: (optional) separate dropout rate at feed-forward
nonlinearity. This is called relu_dropout in T2T.
ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
loss_sparsity_type: str, type of sparsity to used in loss layer. See
SparseDenseWithOptions for options. None if no sparsity should be used.
loss_sparsity: int, the sparsity for loss layer (if used)
loss_d_lowrank: int, the dimensions for intermediate layer (if used)
loss_sparsity_prob: float, the probability for sparse version of loss to be
used. If None, only sparse version is used.
attention_chunk_size: int, if > 0 run attention chunked at this size
n_layers_forget: how often to have a forgetting block between layers
n_decoder_attention_layers: how many attention layers in a decoder block
use_bfloat16: whether to use bfloat16 for weights (default: False)
reversible_encoder: whether to be reversible through the encoder
mode: str: 'train' or 'eval'
Returns:
A Reformer model as a layer that maps from a target, source pair to
activations over a vocab set.
"""
# Set default dimensions for attention head key and value sizes.
if d_attention_key is None:
if d_model % n_heads != 0:
raise ValueError(
f'n_heads ({n_heads}) must divide d_model ({d_model})')
d_attention_key = d_model // n_heads
if d_attention_value is None:
if d_model % n_heads != 0:
raise ValueError(
f'n_heads ({n_heads}) must divide d_model ({d_model})')
d_attention_value = d_model // n_heads
# Vector embeddings.
in_encoder, out_encoder, output_vocab_size = (
ct.EmbeddingAndPositionalEncodings(
input_vocab_size,
d_model,
mode,
dropout,
[-2], # dropout_shared_axes
max_len,
output_vocab_size=output_vocab_size,
axial_pos_shape=axial_pos_shape,
d_axial_pos_embs=d_axial_pos_embs,
use_bfloat16=use_bfloat16))
# pylint: disable=g-complex-comprehension
encoder_blocks = [
EncoderBlock(d_model,
d_ff,
n_heads,
encoder_attention_type,
dropout=dropout,
ff_activation=ff_activation,
ff_dropout=ff_dropout,
ff_use_sru=ff_use_sru,
ff_chunk_size=ff_chunk_size,
ff_sparsity=ff_sparsity,
attention_chunk_size=attention_chunk_size,
use_bfloat16=use_bfloat16,
mode=mode) for _ in range(n_encoder_layers)
]
# pylint: enable=g-complex-comprehension
encoder = [ # vec_e mask_e tok_e tok_d tok_d
tl.ReversibleSelect([0, 0]), # vec_e1 vec_e2 mask_e tok_e tok_d tok_d
_ReversibleSerialForget(encoder_blocks, d_model, n_layers_forget)
]
if not reversible_encoder:
encoder += [
tl.Fn('XYAvg', lambda x, y: (x + y) / 2.0),
tl.Dense(d_model, use_bfloat16=use_bfloat16),
tl.LayerNorm(),
]
encoder = tl.Serial(encoder)
if mode == 'predict':
encoder = tl.Cache(encoder)
decoder_blocks = []
if isinstance(encoder_decoder_attention_type, (tuple, list)):
assert n_decoder_layers % len(encoder_decoder_attention_type) == 0
else:
encoder_decoder_attention_type = [encoder_decoder_attention_type]
for layer_idx in range(n_decoder_layers):
layer_attention_type = encoder_decoder_attention_type[
layer_idx % len(encoder_decoder_attention_type)]
decoder_block = DecoderBlock(
d_model,
d_ff,
d_attention_key,
d_attention_value,
n_heads,
attention_type=layer_attention_type,
dropout=dropout,
ff_activation=ff_activation,
ff_dropout=ff_dropout,
ff_use_sru=ff_use_sru,
ff_chunk_size=ff_chunk_size,
ff_sparsity=ff_sparsity,
attention_chunk_size=attention_chunk_size,
n_attention_layers=n_decoder_attention_layers,
use_bfloat16=use_bfloat16,
mode=mode)
decoder_blocks.append(decoder_block)
dense_loss_layer = tl.SparseDenseWithOptions(
output_vocab_size,
d_input=d_model,
sparsity_type=loss_sparsity_type,
sparsity=loss_sparsity,
d_lowrank=loss_d_lowrank,
prob_sparse=loss_sparsity_prob,
use_bfloat16=use_bfloat16,
mode=mode)
# Layers to merge encoder and decoder, see below for details.
if reversible_encoder:
encdec_layers = [
tl.ReversibleSelect([0, 1, 4, 2,
3]), # vec_e vec_d mask_e tok_e tok_d
t2.ConcatWithPadding2(mode=mode), # vec_ed vec_ed tok_e tok_d
]
else:
encdec_layers = [
tl.ReversibleSelect([0, 3, 1,
2]), # vec_e vec_d mask_e tok_e tok_d
t2.ConcatWithPadding(mode=mode), # vec_ed tok_e tok_d
tl.ReversibleSelect([0, 0]), # vec_ed vec_ed tok_e tok_d
]
# Assemble and return the model.
return tl.Serial(
# Input: encoder_side_tokens, decoder_side_tokens
# Copy decoder tokens for use in loss.
tl.Select([0, 0, 0, 1, 1]), # tok_e tok_e tok_e tok_d tok_d
# Embed in and out tokens; done together as weights may be shared.
tl.Parallel(
in_encoder,
[],
[], # vec_e tok_e tok_e vec_d tok_d
[tl.ShiftRight(mode=mode), out_encoder]),
tl.Parallel([], [
tl.PaddingMask(),
tl.Fn('Squeeze', lambda x: jnp.squeeze(x, (1, 2)), n_out=1)
]),
# # vec_e mask_e tok_e vec_d tok_d
# Encode.
encoder, # vec_e mask_e tok_e vec_d tok_d
# Concat encoder and decoder, given encoder mask.
encdec_layers,
# Run decoder blocks.
_ReversibleSerialForget(
decoder_blocks, d_model,
n_layers_forget), # vec_ed1 vec_ed2 tok_e tok_d
tl.Fn('XYAvg', lambda x, y: (x + y) / 2.0), # vec_ed tok_e tok_d
tl.LayerNorm(), # vec_ed tok_e tok_d
# Separate out the encoder part from the concatenated vector.
tl.Select([0, 1, 2, 2]), # vec_ed tok_e tok_d tok_d
t2.StripFromConcatenateWithPadding(mode=mode), # vec_d tok_d
# Map to output vocab.
dense_loss_layer, # vec_d tok_d
)
def Reformer2(input_vocab_size,
output_vocab_size=None,
d_model=512,
d_ff=2048,
d_attention_key=None,
d_attention_value=None,
n_encoder_layers=6,
n_decoder_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
encoder_attention_type=tl.SelfAttention,
encoder_decoder_attention_type=tl.SelfAttention,
pos_type='fixed-base',
pos_axial_shape=(),
pos_d_axial_embs=None,
pos_start_from_zero_prob=1.0,
pos_max_offset_to_add=0,
ff_activation=tl.Relu,
ff_use_sru=0,
ff_chunk_size=0,
ff_dropout=None,
ff_sparsity=0,
loss_sparsity_type='mult',
loss_sparsity=0,
loss_d_lowrank=0,
loss_sparsity_prob=None,
attention_chunk_size=0,
n_layers_forget=0,
forget_dense=True,
n_decoder_attention_layers=2,
use_bfloat16=False,
reversible_encoder=False,
use_two_swaps_per_encoder_block=True,
center_layernorm=True,
half_before_layer=None,
double_after_layer=None,
mode='train'):
"""Reversible transformer encoder-decoder model.
This model expects an input pair: source, target.
At the moment, this model supports dot-product attention only. For the
attention types in the Reformer paper, see ReformerLM.
Args:
input_vocab_size: int: vocab size of the source.
output_vocab_size: int (optional): vocab size of the target. If None, the
source and target are assumed to have the same vocab.
d_model: int: depth of embedding
d_ff: int: depth of feed-forward layer
d_attention_key: int: depth of key vector for each attention head
d_attention_value: int: depth of value vector for each attention head
n_encoder_layers: int: number of encoder layers
n_decoder_layers: int: number of decoder layers
n_heads: int: number of attention heads
dropout: float: dropout rate (how much to drop out)
max_len: int: maximum symbol length for positional encoding
encoder_attention_type: class: attention class to use, such as SelfAttention
encoder_decoder_attention_type: class: attention class to use, such as
SelfAttention
pos_type: string, the type of positional embeddings to use.
pos_axial_shape: tuple of ints: input shape to use for the axial position
encoding. If unset, axial position encoding is disabled.
pos_d_axial_embs: tuple of ints: depth of position embedding for each axis.
Tuple length must match pos_axial_shape, and values must sum to d_model.
pos_start_from_zero_prob: how often to start from 0 during training,
(if 1.0, we always start from position 0, if less, we randomize).
pos_max_offset_to_add: maximum offset to add to positions during training
when randomizing; this offset plus input length must still be less than
max_len for all training examples.
ff_activation: the non-linearity in feed-forward layer
ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
ff_dropout: float: (optional) separate dropout rate at feed-forward
nonlinearity. This is called relu_dropout in T2T.
ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
loss_sparsity_type: str, type of sparsity to used in loss layer. See
SparseDenseWithOptions for options. None if no sparsity should be used.
loss_sparsity: int, the sparsity for loss layer (if used)
loss_d_lowrank: int, the dimensions for intermediate layer (if used)
loss_sparsity_prob: float, the probability for sparse version of loss to be
used. If None, only sparse version is used.
attention_chunk_size: int, if > 0 run attention chunked at this size
n_layers_forget: how often to have a forgetting block between layers
forget_dense: whether to use Dense or no-op (Serial) as a forget layer.
n_decoder_attention_layers: how many attention layers in a decoder block
use_bfloat16: whether to use bfloat16 for weights (default: False)
reversible_encoder: whether to be reversible through the encoder
use_two_swaps_per_encoder_block: whether to allow even number of swaps in
the encoder
center_layernorm: whether to use centering in LayerNorm (default) or if
to skip it, which is known as RMS normalization.
half_before_layer: int, half d_model and d_ff before that layer
double_after_layer: int, double d_model and d_ff after that layer
mode: str: 'train' or 'eval'
Returns:
A Reformer model as a layer that maps from a target, source pair to
activations over a vocab set.
"""
# Set default dimensions for attention head key and value sizes.
if (d_model / 2) % n_heads != 0:
raise ValueError(
f'n_heads ({n_heads}) must divide d_model/2 ({d_model/2})')
if d_attention_key is None:
d_attention_key = d_model // n_heads
if d_attention_value is None:
d_attention_value = d_model // n_heads
# Set values of d_model, d_ff and d_qkv for the first stage.
d_model1, d_ff1 = d_model, d_ff
d_attention_key1, d_attention_value1 = d_attention_key, d_attention_value
if half_before_layer:
d_model1, d_ff1 = d_model / 2, d_ff / 2
d_attention_key1 = d_attention_key / 2
d_attention_value1 = d_attention_value / 2
# Set values of d_model, d_ff and d_qkv for the final stage.
d_model2, d_ff2 = d_model, d_ff
d_attention_key2, d_attention_value2 = d_attention_key, d_attention_value
if double_after_layer:
d_model2, d_ff2 = d_model * 2, d_ff * 2
d_attention_key2 = d_attention_key * 2
d_attention_value2 = d_attention_value * 2
# Vector embeddings.
in_encoder, out_encoder, output_vocab_size = (
ct.EmbeddingAndPositionalEncodings(
input_vocab_size,
d_model1,
mode,
dropout,
[-2], # dropout_shared_axes
max_len,
output_vocab_size=output_vocab_size,
pos_type=pos_type,
pos_axial_shape=pos_axial_shape,
pos_d_axial_embs=pos_d_axial_embs,
pos_start_from_zero_prob=pos_start_from_zero_prob,
pos_max_offset_to_add=pos_max_offset_to_add,
use_bfloat16=use_bfloat16))
# pylint: disable=g-complex-comprehension
encoder_blocks = [
EncoderBlock(d_model1,
d_ff1,
n_heads,
encoder_attention_type,
dropout=dropout,
ff_activation=ff_activation,
ff_dropout=ff_dropout,
ff_use_sru=ff_use_sru,
ff_chunk_size=ff_chunk_size,
ff_sparsity=ff_sparsity,
attention_chunk_size=attention_chunk_size,
center_layernorm=center_layernorm,
use_bfloat16=use_bfloat16,
use_two_swaps_per_block=use_two_swaps_per_encoder_block,
mode=mode) for _ in range(n_encoder_layers)
]
# pylint: enable=g-complex-comprehension
encoder = [ # vec_e mask_e tok_e tok_d tok_d
tl.ReversibleSelect([0, 0]), # vec_e1 vec_e2 mask_e tok_e tok_d tok_d
_ReversibleSerialForget(encoder_blocks, d_model1, n_layers_forget,
forget_dense)
]
if not reversible_encoder:
encoder += [
tl.Fn('XYAvg', lambda x, y: (x + y) / 2.0),
tl.Dense(d_model1, use_bfloat16=use_bfloat16),
tl.LayerNorm(),
]
encoder = tl.Serial(encoder)
if mode == 'predict':
encoder = tl.Cache(encoder)
decoder_blocks = []
if isinstance(encoder_decoder_attention_type, (tuple, list)):
assert n_decoder_layers % len(encoder_decoder_attention_type) == 0
else:
encoder_decoder_attention_type = [encoder_decoder_attention_type]
for layer_idx in range(n_decoder_layers):
layer_attention_type = encoder_decoder_attention_type[
layer_idx % len(encoder_decoder_attention_type)]
# Grow d_model, d_ff, and d_qkv if requested.
d_m, d_f, d_k, d_v = d_model1, d_ff1, d_attention_key1, d_attention_value1
if half_before_layer and layer_idx >= half_before_layer:
d_m, d_f, d_k, d_v = d_model, d_ff, d_attention_key, d_attention_value
if double_after_layer and layer_idx > double_after_layer:
d_m, d_f, d_k, d_v = d_model2, d_ff2, d_attention_key2, d_attention_value2
decoder_block = DecoderBlock(
d_m,
d_f,
d_k,
d_v,
n_heads,
attention_type=layer_attention_type,
dropout=dropout,
ff_activation=ff_activation,
ff_dropout=ff_dropout,
ff_use_sru=ff_use_sru,
ff_chunk_size=ff_chunk_size,
ff_sparsity=ff_sparsity,
attention_chunk_size=attention_chunk_size,
n_attention_layers=n_decoder_attention_layers,
center_layernorm=center_layernorm,
use_bfloat16=use_bfloat16,
mode=mode)
decoder_blocks.append(decoder_block)
if half_before_layer and layer_idx == half_before_layer - 1:
decoder_blocks.append(tl.ReversibleConcatenatePair())
if double_after_layer and layer_idx == double_after_layer:
decoder_blocks.append(tl.ReversibleConcatenatePair())
dense_loss_layer = tl.SparseDenseWithOptions(
output_vocab_size,
d_input=d_model2,
sparsity_type=loss_sparsity_type,
sparsity=loss_sparsity,
d_lowrank=loss_d_lowrank,
prob_sparse=loss_sparsity_prob,
use_bfloat16=use_bfloat16,
mode=mode)
# Layers to merge encoder and decoder, see below for details.
if reversible_encoder:
encdec_layers = [
tl.ReversibleSelect([0, 1, 4, 2,
3]), # vec_e vec_d mask_e tok_e tok_d
t2.ConcatWithPadding2(mode=mode), # vec_ed vec_ed tok_e tok_d
]
else:
encdec_layers = [
tl.ReversibleSelect([0, 3, 1,
2]), # vec_e vec_d mask_e tok_e tok_d
t2.ConcatWithPadding(mode=mode), # vec_ed tok_e tok_d
tl.ReversibleSelect([0, 0]), # vec_ed vec_ed tok_e tok_d
]
# Assemble and return the model.
return tl.Serial(
# Input: encoder_side_tokens, decoder_side_tokens
# Copy decoder tokens for use in loss.
tl.Select([0, 0, 0, 1, 1]), # tok_e tok_e tok_e tok_d tok_d
# Embed in and out tokens; done together as weights may be shared.
tl.Parallel(
in_encoder,
[],
[], # vec_e tok_e tok_e vec_d tok_d
[tl.ShiftRight(mode=mode), out_encoder]),
# Predict mode doesn't work with padding in encoder. Raising an exception
# in jitted function isn't possible, so the second next best thing is
# to convert every embedding to NaNs, so the user will not get subtly
# wrong results, but clearly wrong results.
(_ConvertToNaNsOnAnyZero() if mode == 'predict' else []),
tl.Parallel([], [
tl.PaddingMask(),
tl.Fn('Squeeze', lambda x: jnp.squeeze(x, (1, 2)), n_out=1)
]),
# # vec_e mask_e tok_e vec_d tok_d
# Encode.
encoder, # vec_e mask_e tok_e vec_d tok_d
# Concat encoder and decoder, given encoder mask.
encdec_layers,
# Run decoder blocks.
_ReversibleSerialForget(decoder_blocks, d_model2, n_layers_forget,
forget_dense), # vec_ed1 vec_ed2 tok_e tok_d
tl.Fn('XYAvg', lambda x, y: (x + y) / 2.0), # vec_ed tok_e tok_d
tl.LayerNorm(), # vec_ed tok_e tok_d
# Separate out the encoder part from the concatenated vector.
tl.Select([0, 1, 2, 2]), # vec_ed tok_e tok_d tok_d
t2.StripFromConcatenateWithPadding(mode=mode), # vec_d tok_d
# Map to output vocab.
dense_loss_layer, # vec_d tok_d
)