Upload 3 files

aglib/meliad/transformer/configs/base_htrans.gin

@@ -0,0 +1,96 @@
+# Base configuration for the Hierarchical Transformer.
+
+include "trainer_configuration.gin"
+
+# Imports
+from transformer import attention
+from transformer import decoder_stack
+from transformer import models
+from transformer import nn_components
+from transformer import transformer_base
+from transformer import transformer_layer
+
+
+NUM_LAYERS = 12
+NUM_HEADS = 8
+HEAD_DIM = 128
+EMBED_DIM = 512        # Size of embedding vector for each token
+MLP_DIM = 2048         # Number of hidden units in transformer FFN
+NUM_EMBEDDINGS = 256   # Number of tokens in vocabulary.
+DROPOUT_RATE = 0.05
+ATTN_DROPOUT_RATE = 0.05
+
+# For training on TPU.
+DTYPE="bfloat16"
+
+# Configure the input task.
+decoder_stack.TransformerTaskConfig:
+  dataset_name = "synthetic"
+  train_split = "train"
+  test_split = "test"
+  sequence_length = 512
+  batch_size = 8
+  vocab_size = %NUM_EMBEDDINGS
+
+transformer_layer.TransformerLayer:
+  num_heads = %NUM_HEADS
+  head_size = %HEAD_DIM
+  window_length = 512
+  use_long_xl_architecture = True
+  max_unrolled_windows = -1           # Always unroll.
+  relative_position_type = "t5"       # Can be "fourier", "t5", or None.
+  use_causal_mask = True
+  attn_dropout_rate = %ATTN_DROPOUT_RATE   # Attention matrix dropout.
+
+  memory_num_neighbors = 0
+  compute_importance = False
+  dtype = %DTYPE
+
+transformer_base.TransformerBase:
+  attn_mlp_factory = @transformer_attn/nn_components.MLP
+  ffn_factory = @transformer_ffn/nn_components.MLP
+  normalize_keys = True               # More stable with Transformer XL.
+  dropout_rate = %DROPOUT_RATE
+  pre_attn_dropout = True
+  post_attn_dropout = False
+  pre_ffn_dropout = False
+  post_ffn_dropout = True
+
+transformer_attn/nn_components.MLP:
+  num_layers = 1                      # Just a single dense matmul.
+  num_hidden_units = 0
+  hidden_activation = None
+  use_bias = False
+
+transformer_ffn/nn_components.MLP:
+  num_layers = 2
+  num_hidden_units = %MLP_DIM
+  hidden_activation = "relu"
+  use_bias = False
+
+decoder_stack.DecoderStack:
+  # task_config will be passed in from DecoderOnlyLanguageModel.
+  num_layers = %NUM_LAYERS
+  embedding_size = %EMBED_DIM
+  embedding_stddev = 1.0
+  layer_factory = @transformer_layer.TransformerLayer
+  dstack_window_length = 0
+  use_absolute_positions = False
+  use_final_layernorm = True          # Final layernorm before token lookup.
+  final_dropout_rate = %DROPOUT_RATE  # Dropout before token lookup.
+  final_mlp_factory = None            # Final MLP to predict target tokens.
+  recurrent_layer_indices = ()
+  memory_factory = None     # e.g. @memory_factory.memory_on_tpu_factory
+  memory_layer_indices = ()
+  dtype = %DTYPE
+
+models.DecoderOnlyLanguageModel:
+  task_config = @decoder_stack.TransformerTaskConfig()
+  decoder_factory = @decoder_stack.DecoderStack
+
+nn_components.LayerNorm:
+  use_scale = True
+  use_bias = False
+  use_mean = False   # Calculate and adjust for the mean as well as the scale.
+  dtype = %DTYPE
+

aglib/meliad/transformer/configs/memory_configuration.gin

@@ -0,0 +1,22 @@
+# Configure external memory.
+# This file should be included after base_htrans.gin.
+
+import  training_loop
+from transformer import memory_factory
+
+MEMORY_HEAD_DIM = %HEAD_DIM
+NUM_MEMORY_HEADS = %gin.REQUIRED
+
+memory_factory.memory_on_tpu_factory:
+  num_heads = %NUM_MEMORY_HEADS
+  key_size = %MEMORY_HEAD_DIM
+  value_size = %MEMORY_HEAD_DIM
+  database_size = 8192
+  dtype = %DTYPE  # defined in base_htrans.gin
+
+
+training_loop.Trainer:
+  log_every_steps = 100    # memory can slow down training, need responsive stats
+  checkpoint_every_steps = 1000
+  generate_every_steps = 0  # disable generate mode when using external memory.
+

aglib/meliad/transformer/configs/trainer_configuration.gin

@@ -0,0 +1,56 @@
+
+import  optimizer_config
+import  training_loop
+from transformer import models
+from transformer import text_dataset
+
+# Training setup.
+training_loop.Trainer:
+  model_definition = @models.DecoderOnlyLanguageModel
+
+  num_steps = 250_000
+  status_every_steps = 10
+  log_every_steps = 1000
+  test_every_steps = 1000
+  num_test_steps = 400
+  generate_every_steps = 5000
+  print_input_every_steps = 5000
+  checkpoint_every_steps = 5000
+  save_checkpoints = True
+  restore_checkpoints = True
+  use_separate_metric_directories = False
+
+  optimizer_factory = @optimizer_config.FlaxAdafactorConfig()
+  learning_rate_schedule = @optimizer_config.lr_cosine_decay
+  max_scheduled_steps = 0   # Use num_steps as max_scheduled_steps.
+  warmup_steps = 1000
+  learning_rate_multiplier = 1.0
+
+  rng_key_names = ("dropout", "sample")
+
+text_dataset.load_text_dataset:
+  verbose = False  # if true, prints the start of every book/repo read from disk
+
+# Use cosine decay to max_scheduled_steps, as described in Chinchilla:
+# https://arxiv.org/abs/2203.15556
+optimizer_config.lr_cosine_decay:
+    max_lr = 0.01
+    min_lr = 0.001
+    decay_after = True
+    spike_steps = 0
+    spike_lr = 0.0
+
+
+# Adam optimizer configuration.
+# optimizer_config.AdamConfig:
+#   learning_rate = 0.05  # Will be multiplied by the LR schedule.
+#   beta1 = 0.9
+#   beta2 = 0.98
+#   weight_decay_rate = 0.0
+
+# Adafactor optimizer configuration.
+optimizer_config.FlaxAdafactorConfig:
+  learning_rate = 1.0   # Will be multiplied by the LR schedule.
+  beta1 = 0.9           # Can be "None".
+
+