Upload 33 files

aglib/meliad/transformer/configs/options/debug_mode.gin

@@ -0,0 +1,5 @@
+
+training_loop.Trainer:
+  replicate_mode = False
+  trace_debug_mode = True
+  print_variables = True

aglib/meliad/transformer/configs/options/enable_scan.gin

@@ -0,0 +1,5 @@
+
+DTYPE="float32"  # Required when using scan for stability.
+
+transformer_layer.TransformerLayer:
+  max_unrolled_windows = 0   # Never unroll.

aglib/meliad/transformer/configs/options/external_memory_32k.gin

@@ -0,0 +1,17 @@
+
+include "memory_configuration.gin"
+
+from transformer import attention
+from transformer import memory_factory
+
+NUM_MEMORY_HEADS = %NUM_HEADS
+
+decoder_stack.DecoderStack:
+  memory_factory = @memory_factory.memory_on_tpu_factory
+  memory_layer_indices = (-4,)
+
+transformer_layer.TransformerLayer:
+  memory_num_neighbors = 128
+
+memory_factory.memory_on_tpu_factory:
+  database_size = 32768

aglib/meliad/transformer/configs/options/external_memory_8k.gin

@@ -0,0 +1,17 @@
+
+include "memory_configuration.gin"
+
+from transformer import attention
+from transformer import memory_factory
+
+NUM_MEMORY_HEADS = %NUM_HEADS
+
+decoder_stack.DecoderStack:
+  memory_factory = @memory_factory.memory_on_tpu_factory
+  memory_layer_indices = (-4,)
+
+transformer_layer.TransformerLayer:
+  memory_num_neighbors = 128
+
+memory_factory.memory_on_tpu_factory:
+  database_size = 8192

aglib/meliad/transformer/configs/options/final_mlp.gin

@@ -0,0 +1,14 @@
+
+from transformer import decoder_stack
+from transformer import nn_components
+
+decoder_stack.DecoderStack:
+  final_mlp_factory = @decoder_final_mlp/nn_components.MLP
+
+# Add a final MLP for token prediction after the last transformer layer.
+decoder_final_mlp/nn_components.MLP:
+  num_hidden_units = %MLP_DIM
+  num_layers = 2
+  activation_function = "relu"
+  use_bias = False
+  dtype = %DTYPE

aglib/meliad/transformer/configs/options/lr_cosine_decay.gin

@@ -0,0 +1,13 @@
+
+import  optimizer_config
+
+training_loop.Trainer:
+  learning_rate_schedule = @optimizer_config.lr_cosine_decay
+
+optimizer_config.lr_cosine_decay:
+    max_lr = 0.01
+    min_lr = 0.001
+    decay_after = True
+    spike_steps = 0
+    spike_lr = 0.0
+

aglib/meliad/transformer/configs/options/lr_cosine_decay_spike.gin

@@ -0,0 +1,13 @@
+
+import  optimizer_config
+
+training_loop.Trainer:
+  learning_rate_schedule = @optimizer_config.lr_cosine_decay
+
+optimizer_config.lr_cosine_decay:
+    max_lr = 0.01
+    min_lr = 0.001
+    decay_after = True
+    spike_steps = 10_000
+    spike_lr = 0.04
+

aglib/meliad/transformer/configs/options/lr_linear_decay.gin

@@ -0,0 +1,10 @@
+
+import  optimizer_config
+
+training_loop.Trainer:
+  learning_rate_schedule = @optimizer_config.lr_linear_decay
+
+optimizer_config.lr_linear_decay:
+    max_lr = 0.01
+    min_lr = 0.001
+    decay_after = True

aglib/meliad/transformer/configs/options/lr_max_steps_125k.gin

@@ -0,0 +1,3 @@
+
+training_loop.Trainer:
+  max_scheduled_steps = 125_000

aglib/meliad/transformer/configs/options/lr_max_steps_250k.gin

@@ -0,0 +1,3 @@
+
+training_loop.Trainer:
+  max_scheduled_steps = 250_000

aglib/meliad/transformer/configs/options/lr_max_steps_500k.gin

@@ -0,0 +1,3 @@
+
+training_loop.Trainer:
+  max_scheduled_steps = 500_000

aglib/meliad/transformer/configs/options/lr_rsqrt_decay.gin

@@ -0,0 +1,9 @@
+
+import  optimizer_config
+
+training_loop.Trainer:
+  learning_rate_schedule = @optimizer_config.lr_rsqrt_decay
+
+optimizer_config.lr_rsqrt_decay:
+    max_lr = 0.05
+    min_lr = 0.001

aglib/meliad/transformer/configs/options/lr_rsqrt_decay_std.gin

@@ -0,0 +1,11 @@
+
+import  optimizer_config
+
+# Implement standard rsqrt decay as used in the memorizing and block-recurrent
+# transformer papers, which does not decay to a specified minimum learning
+# rate over max_steps.
+training_loop.Trainer:
+  learning_rate_schedule = @optimizer_config.lr_rsqrt_decay_std
+
+optimizer_config.lr_rsqrt_decay_std:
+    max_lr = None

aglib/meliad/transformer/configs/options/lr_scale_05.gin

@@ -0,0 +1,3 @@
+
+training_loop.Trainer:
+  learning_rate_multiplier = 0.5

aglib/meliad/transformer/configs/options/lr_scale_2.gin

@@ -0,0 +1,3 @@
+
+training_loop.Trainer:
+  learning_rate_multiplier = 2.0

aglib/meliad/transformer/configs/options/no_norm.gin

@@ -0,0 +1,3 @@
+
+transformer_base.TransformerBase:
+  normalize_keys = False

aglib/meliad/transformer/configs/options/positions_absolute.gin

@@ -0,0 +1,7 @@
+
+transformer_layer.TransformerLayer:
+  relative_position_type = None
+
+decoder_stack.DecoderStack:
+  use_absolute_positions = True
+

aglib/meliad/transformer/configs/options/positions_fourier.gin

@@ -0,0 +1,7 @@
+
+transformer_layer.TransformerLayer:
+  relative_position_type = "fourier"
+
+decoder_stack.DecoderStack:
+  use_absolute_positions = False
+

aglib/meliad/transformer/configs/options/positions_rotary.gin

@@ -0,0 +1,7 @@
+
+transformer_layer.TransformerLayer:
+  relative_position_type = "rotary"
+
+decoder_stack.DecoderStack:
+  use_absolute_positions = False
+

aglib/meliad/transformer/configs/options/positions_t5.gin

@@ -0,0 +1,7 @@
+
+transformer_layer.TransformerLayer:
+  relative_position_type = "t5"
+
+decoder_stack.DecoderStack:
+  use_absolute_positions = False
+

aglib/meliad/transformer/configs/options/seq_1024.gin

@@ -0,0 +1,7 @@
+
+decoder_stack.TransformerTaskConfig:
+  sequence_length = 1024
+  batch_size = 4
+
+transformer_layer.TransformerLayer:
+  use_long_xl_architecture = True

aglib/meliad/transformer/configs/options/seq_1024_nocache.gin

@@ -0,0 +1,8 @@
+
+decoder_stack.TransformerTaskConfig:
+  sequence_length = 1024
+  batch_size = 4
+
+transformer_layer.TransformerLayer:
+  window_length = 1024
+  use_long_xl_architecture = False

aglib/meliad/transformer/configs/options/seq_2048.gin

@@ -0,0 +1,7 @@
+
+decoder_stack.TransformerTaskConfig:
+  sequence_length = 2048
+  batch_size = 2
+
+transformer_layer.TransformerLayer:
+  use_long_xl_architecture = True

aglib/meliad/transformer/configs/options/seq_2048_nocache.gin

@@ -0,0 +1,8 @@
+
+decoder_stack.TransformerTaskConfig:
+  sequence_length = 2048
+  batch_size = 2
+
+transformer_layer.TransformerLayer:
+  window_length = 2048
+  use_long_xl_architecture = False

aglib/meliad/transformer/configs/options/seq_4096.gin

@@ -0,0 +1,7 @@
+
+decoder_stack.TransformerTaskConfig:
+  sequence_length = 4096
+  batch_size = 1
+
+transformer_layer.TransformerLayer:
+  use_long_xl_architecture = True

aglib/meliad/transformer/configs/options/seq_512.gin

@@ -0,0 +1,7 @@
+
+decoder_stack.TransformerTaskConfig:
+  sequence_length = 512
+  batch_size = 8
+
+transformer_layer.TransformerLayer:
+  use_long_xl_architecture = True

aglib/meliad/transformer/configs/options/seq_512_nocache.gin

@@ -0,0 +1,8 @@
+
+decoder_stack.TransformerTaskConfig:
+  sequence_length = 512
+  batch_size = 8
+
+transformer_layer.TransformerLayer:
+  window_length = 512
+  use_long_xl_architecture = False

aglib/meliad/transformer/configs/options/stack_window_512.gin

@@ -0,0 +1,6 @@
+
+transformer_layer.TransformerLayer:
+  window_length = 512
+
+decoder_stack.DecoderStack:
+  dstack_window_length = 512

aglib/meliad/transformer/configs/options/window_1024.gin

@@ -0,0 +1,4 @@
+
+# Sequence length must be larger than window_length.
+transformer_layer.TransformerLayer:
+  window_length = 1024

aglib/meliad/transformer/configs/options/window_128.gin

@@ -0,0 +1,4 @@
+
+# Sequence length must be larger than window_length.
+transformer_layer.TransformerLayer:
+  window_length = 128

aglib/meliad/transformer/configs/options/window_2048.gin

@@ -0,0 +1,4 @@
+
+# Sequence length must be larger than window_length.
+transformer_layer.TransformerLayer:
+  window_length = 2048

aglib/meliad/transformer/configs/options/window_256.gin

@@ -0,0 +1,4 @@
+
+# Sequence length must be larger than window_length.
+transformer_layer.TransformerLayer:
+  window_length = 256

aglib/meliad/transformer/configs/options/window_512.gin

@@ -0,0 +1,4 @@
+
+# Sequence length must be larger than window_length.
+transformer_layer.TransformerLayer:
+  window_length = 512