exp

finetune_large_mt5_sentencefix_v4_16.gin

@@ -25,35 +25,17 @@ RANDOM_SEED = 0
 INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/mt5_large/checkpoint_1000000"
 
 train_script.train:
-  eval_period = 100
-  partitioner = @partitioning.PjitPartitioner()
+  eval_period = 500
+  partitioner = @partitioning.ModelBasedPjitPartitioner()
 
-utils.RestoreCheckpointConfig:
-  path = %INITIAL_CHECKPOINT_PATH
-  mode = 'specific'
+# `num_decodes` is equivalent to a beam size in a beam search decoding.
+models.EncoderDecoderModel.predict_batch_with_aux.num_decodes = 4
 
-#train_script.train:
-#  train_dataset_cfg = @train/utils.DatasetConfig()
-#  train_eval_dataset_cfg = @train_eval/utils.DatasetConfig()
-#  infer_eval_dataset_cfg = @infer_eval/utils.DatasetConfig()
+partitioning.ModelBasedPjitPartitioner.num_partitions = 2
 
-models.EncoderDecoderModel.predict_batch_with_aux.num_decodes = 4
 
-infer_eval/utils.DatasetConfig:
-  mixture_or_task_name = %MIXTURE_OR_TASK_NAME
-  task_feature_lengths = %TASK_FEATURE_LENGTHS
-  split = 'validation'
-  batch_size = 64
-  shuffle = False
-  seed = 42
-  use_cached = %USE_CACHED_TASKS
-  pack = False
-  module = %MIXTURE_OR_TASK_MODULE
-
-partitioning.PjitPartitioner:
-  num_partitions = 4
-  model_parallel_submesh = None
-  logical_axis_rules = @partitioning.standard_logical_axis_rules()
-
-#partitioning.PjitPartitioner.num_partitions = 4
+#from t5.models import mesh_transformer
+#import t5.models
+#mesh_transformer.learning_rate_schedules.constant_learning_rate.learning_rate = 0.0005
+#run.learning_rate_schedule = @learning_rate_schedules.constant_learning_rate