Add config

config.json

@@ -3,9 +3,9 @@
   "architectures": [
     "T5ForConditionalGeneration"
   ],
-  "d_ff": 16384,
-  "d_kv": 32,
-  "d_model": 256,
+  "d_ff": 3072,
+  "d_kv": 64,
+  "d_model": 768,
   "decoder_start_token_id": 0,
   "dropout_rate": 0.1,
   "eos_token_id": 1,
@@ -16,7 +16,7 @@
   "model_type": "t5",
   "n_positions": 512,
   "num_decoder_layers": 12,
-  "num_heads": 4,
+  "num_heads": 12,
   "num_layers": 12,
   "pad_token_id": 0,
   "relative_attention_num_buckets": 32,

operative_config.gin

@@ -9,15 +9,15 @@ import t5.models.mesh_transformer
 
 # Macros:
 # ==============================================================================
-d_ff = 16384
-d_kv = 32
-d_model = 256
+d_ff = 3072
+d_kv = 64
+d_model = 768
 dropout_rate = 0.0
 inputs_length = 512
 mean_noise_span_length = 3.0
 MIXTURE_NAME = 'c4_v220_unsupervised'
 noise_density = 0.15
-num_heads = 4
+num_heads = 12
 num_layers = 12
 
 # Parameters for adafactor_decay_rate_pow:
@@ -146,7 +146,6 @@ encoder/make_layer_stack.num_layers = %num_layers
 mesh_train_dataset_fn.mixture_or_task_name = %MIXTURE_NAME
 mesh_train_dataset_fn.pack = True
 mesh_train_dataset_fn.seed = None
-mesh_train_dataset_fn.shuffle = True
 mesh_train_dataset_fn.use_cached = 1
 
 # Parameters for noise_span_to_unique_sentinel:
@@ -195,7 +194,6 @@ rewrite_stack_variables.max_combined_variable_size = 536870912
 # ==============================================================================
 run.autostack = True
 run.batch_size = ('tokens_per_batch', 65536)
-run.checkpoint_input_pipeline = False
 run.dataset_split = 'train'
 run.ensemble_inputs = None
 run.eval_checkpoint_step = None
@@ -217,7 +215,7 @@ run.optimizer = @optimize.AdafactorOptimizer
 run.output_eval_examples = True
 run.perplexity_eval_steps = 100
 run.predict_fn = None
-run.save_checkpoints_steps = 5000
+run.save_checkpoints_steps = 10000
 run.seen_data_init_step = 0
 run.sequence_length = {'inputs': 512, 'targets': 128}
 run.skip_seen_data = False
@@ -312,7 +310,7 @@ tpu_estimator_model_fn.tpu_summaries = False
 # Parameters for tpu_mesh_shape:
 # ==============================================================================
 tpu_mesh_shape.ensemble_parallelism = None
-tpu_mesh_shape.model_parallelism = 2
+tpu_mesh_shape.model_parallelism = 1
 tpu_mesh_shape.tpu_topology = '4x4'
 
 # Parameters for unit_scaling_convention: