ready to train

finetune_mt5_base.gin

@@ -1,17 +0,0 @@
-from __gin__ import dynamic_registration
-
-include 't5x/examples/t5/mt5/base.gin'
-include 't5x/configs/runs/finetune.gin'
-
-# Register necessary SeqIO Tasks/Mixtures.
-import t5.data.mixtures
-
-import tasks
-
-INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/mt5_base/checkpoint_1000000"
-MIXTURE_OR_TASK_NAME = "ncc_small_span_corruption"
-TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
-TRAIN_STEPS = 1100000
-DROPOUT_RATE = 0.0
-# BATCH_SIZE = 256
-

norwegian_base.gin

@@ -1,7 +1,7 @@
 include 't5x/examples/t5/mt5/base.gin'
-#include 'pretrain_cont.gin'
+include 'pretrain_cont.gin'
 #include 't5x/configs/runs/pretrain.gin'
-include 't5x/configs/runs/finetune.gin'
+#iinclude 't5x/configs/runs/finetune.gin'
 
 
 # Register necessary SeqIO Tasks/Mixtures.
@@ -12,6 +12,6 @@ MIXTURE_OR_TASK_NAME = "ncc_span_corruption_stream"
 TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
 TRAIN_STEPS = 1_100_000
 DROPOUT_RATE = 0.0 # Changed from the default since T5-1.1 recomments this.
-INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/mt5_xxl/checkpoint_1000000"
+INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/mt5_base/checkpoint_1000000"
 #PjitPartitioner.num_partitions = 2
 

pretrain_cont.gin

@@ -93,6 +93,7 @@ utils.CheckpointConfig:
   save = @utils.SaveCheckpointConfig()
 utils.RestoreCheckpointConfig:
   path = %INITIAL_CHECKPOINT_PATH
+  mode = 'specific'
   dtype = 'float32'
 utils.SaveCheckpointConfig:
   period = 1000

train_base.sh

@@ -1,9 +1,9 @@
 PROJECT_DIR=${HOME}"/models/pk-nb-t5x"
 T5X_DIR="../../t5x"  # directory where the t5x is cloned.
-MODEL_DIR="gs://nb-t5x-us-central2/pk_nb_t5x_base_test2"
+MODEL_DIR="gs://nb-t5x-us-central2/pk_nb_t5x_base_run1"
 export PYTHONPATH=${PROJECT_DIR}
 
 python3 ${T5X_DIR}/t5x/train.py \
   --gin_search_paths=${PROJECT_DIR} \
-  --gin_file="finetune_mt5_base.gin" \
+  --gin_file="norwegian_base.gin" \
   --gin.MODEL_DIR="'${MODEL_DIR}'" \