pretrain_cont

continue_pretraining_base.gin

@@ -1,15 +1,15 @@
 include 't5x/examples/t5/mt5/base.gin'
 include 't5x/configs/runs/pretrain.gin'
+include 't5x/configs/runs/finetune.gin'
 
 # Register necessary SeqIO Tasks/Mixtures.
 import t5.data.mixtures
 import tasks
 
-MIXTURE_OR_TASK_NAME = "span_corruption"
-TASK_FEATURE_LENGTHS = {"inputs": 256, "targets": 256}
-TRAIN_STEPS = 1_001_000
-DROPOUT_RATE = 0.0
+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"
+PjitPartitioner.num_partitions = 2
 
-#Batch size should be the default for pretraining
-#BATCH_SIZE = 256

pretrain_cont.gin

@@ -0,0 +1,111 @@
+# Defaults for pretraining with train.py.
+#
+#
+# You must also include a binding for MODEL.
+#
+# Required to be set
+#
+# - MIXTURE_OR_TASK_NAME
+# - TASK_FEATURE_LENGTHS
+# - TRAIN_STEPS - include pretrain steps
+# - MODEL_DIR: # automatically set when using xm_launch
+#
+# Commonly overridden options:
+#
+# - train/DatasetConfig.batch_size
+# - train_eval/DatasetConfig.batch_size
+# - PjitPartitioner.num_partitions
+# - Trainer.num_microbatches
+# - DROPOUT_RATE
+from __gin__ import dynamic_registration
+
+import __main__ as train_script
+from t5x import gin_utils
+from t5x import partitioning
+from t5x import utils
+from t5x import trainer
+
+MIXTURE_OR_TASK_NAME = %gin.REQUIRED
+TASK_FEATURE_LENGTHS = %gin.REQUIRED
+TRAIN_STEPS = %gin.REQUIRED
+MODEL_DIR = %gin.REQUIRED
+BATCH_SIZE = 128
+USE_CACHED_TASKS = True
+INITIAL_CHECKPOINT_PATH = %gin.REQUIRED
+
+# DEPRECATED: Import the this module in your gin file.
+MIXTURE_OR_TASK_MODULE = None
+SHUFFLE_TRAIN_EXAMPLES = True
+
+# HW RNG is faster than SW, but has limited determinism.
+# Most notably it is not deterministic across different
+# submeshes.
+USE_HARDWARE_RNG = False
+# None always uses faster, hardware RNG
+RANDOM_SEED = None
+
+# Can be overridden with `train.*`.`
+train_script.train:
+  model = %MODEL  # imported from separate gin file
+  model_dir = %MODEL_DIR
+  train_dataset_cfg = @train/utils.DatasetConfig()
+  train_eval_dataset_cfg = @train_eval/utils.DatasetConfig()
+  infer_eval_dataset_cfg = None
+  checkpoint_cfg = @utils.CheckpointConfig()
+  partitioner = @partitioning.PjitPartitioner()
+  trainer_cls = @trainer.Trainer
+  total_steps = %TRAIN_STEPS
+  eval_steps = 20
+  eval_period = 1000
+  random_seed = %RANDOM_SEED
+  use_hardware_rng = %USE_HARDWARE_RNG
+  summarize_config_fn = @gin_utils.summarize_gin_config
+
+partitioning.PjitPartitioner:
+  num_partitions = 1
+  model_parallel_submesh = None
+  logical_axis_rules = @partitioning.standard_logical_axis_rules()
+
+train/utils.DatasetConfig:
+  mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+  task_feature_lengths = %TASK_FEATURE_LENGTHS
+  split = 'train'
+  batch_size = %BATCH_SIZE
+  shuffle = %SHUFFLE_TRAIN_EXAMPLES
+  seed = None  # use a new seed each run/restart
+  use_cached = %USE_CACHED_TASKS
+  pack = True
+  module = %MIXTURE_OR_TASK_MODULE
+
+train_eval/utils.DatasetConfig:
+  mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+  task_feature_lengths = %TASK_FEATURE_LENGTHS
+  split = 'validation'
+  batch_size = %BATCH_SIZE
+  shuffle = False
+  seed = 42
+  use_cached = %USE_CACHED_TASKS
+  pack = True
+  module = %MIXTURE_OR_TASK_MODULE
+
+utils.CheckpointConfig:
+  restore = @utils.RestoreCheckpointConfig()
+  save = @utils.SaveCheckpointConfig()
+utils.RestoreCheckpointConfig:
+  path = %INITIAL_CHECKPOINT_PATH
+  mode = 'specific'
+  dtype = 'float32'
+utils.SaveCheckpointConfig:
+  period = 1000
+  dtype = 'float32'
+  keep = None  # keep all checkpoints
+  save_dataset = False  # don't checkpoint dataset state
+
+trainer.Trainer:
+  num_microbatches = None
+  learning_rate_fn = @utils.create_learning_rate_scheduler()
+
+utils.create_learning_rate_scheduler:
+  factors = 'constant * rsqrt_decay'
+  base_learning_rate = 0.5 #This is set to half of the original since it is continued training
+  warmup_steps = 10000  # 10k to keep consistent with T5/MTF defaults.

tasks.py

@@ -11,15 +11,18 @@ from seqio import FunctionDataSource, utils
 
 TaskRegistry = seqio.TaskRegistry
 
-vocabulary = seqio.SentencePieceVocabulary('gs://t5-data/vocabs/mc4.250000.100extra/sentencepiece.model', extra_ids=0)
+
+
 
 DEFAULT_OUTPUT_FEATURES = {
     "inputs": seqio.Feature(
-        vocabulary=vocabulary, add_eos=True,
+        vocabulary=t5.data.get_default_vocabulary(), add_eos=True,
         required=False),
     "targets": seqio.Feature(
-        vocabulary=vocabulary, add_eos=True)
+        vocabulary=t5.data.get_default_vocabulary(), add_eos=True)
 }
+# Custom vocabs can also be defined and loaded
+# vocabulary = seqio.SentencePieceVocabulary("gs://t5-data/vocabs/mc4.250000.100extra/sentencepiece.model")
 
 
 def gen_dataset(split, shuffle=False, seed=None, column="text", dataset_params=None):
@@ -48,15 +51,15 @@ def target_to_key(x, key_map, target_key):
 
 
 # Final pretraining task used in Raffel et al., 2019 adaptated to NCC
-dataset_name = 'NbAiLab/NCC_small'
-dataset_params = {"path": dataset_name, "use_auth_token": True}
+dataset_name = 'NbAiLab/NCC'
+dataset_params = {"path": dataset_name}
 dataset_shapes = {'train': 20830348, 'validation': 473079}
 TaskRegistry.add(
-    'span_corruption',
+    "ncc_span_corruption",
     source=seqio.FunctionDataSource(
         dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
         splits=("train", "validation"),
-        #caching_permitted=True,
+        caching_permitted=False,
         num_input_examples=dataset_shapes,
     ),
     preprocessors=[
@@ -66,24 +69,52 @@ TaskRegistry.add(
                 "targets": None,
             }, target_key="targets"),
         seqio.preprocessors.tokenize,
-        #seqio.CacheDatasetPlaceholder(),
-        preprocessors.span_corruption,
+        # seqio.CacheDatasetPlaceholder(),
+        preprocessors.span_corruption, 
         seqio.preprocessors.append_eos_after_trim,
     ],
-    output_features={"targets": DEFAULT_OUTPUT_FEATURES["targets"]},
+    output_features={"targets": seqio.Feature(vocabulary=vocabulary, add_eos=True)},
     metric_fns=[]
 )
 
-# Final pretraining task used in Raffel et al., 2019 adaptated to nbailab_extended
-dataset_name = 'NbAiLab/nbailab_extended'
+# Final pretraining task used in Raffel et al., 2019 adaptated to NCC
+dataset_name = 'NbAiLab/NCC_amall'
+dataset_params = {"path": dataset_name}
+dataset_shapes = {'train': 20830348, 'validation': 473079}
+TaskRegistry.add(
+    "ncc_small_span_corruption",
+    source=seqio.FunctionDataSource(
+        dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
+        splits=("train", "validation"),
+        caching_permitted=False,
+        num_input_examples=dataset_shapes,
+    ),
+    preprocessors=[
+        functools.partial(
+            target_to_key, key_map={
+                "inputs": None,
+                "targets": None,
+            }, target_key="targets"),
+        seqio.preprocessors.tokenize,
+        # seqio.CacheDatasetPlaceholder(),
+        preprocessors.span_corruption, 
+        seqio.preprocessors.append_eos_after_trim,
+    ],
+    output_features={"targets": seqio.Feature(vocabulary=vocabulary, add_eos=True)},
+    metric_fns=[]
+)
+
+
+# Final pretraining task used in Raffel et al., 2019 adaptated to NCC
+dataset_name = 'NbAiLab/NCC'
 dataset_params = {"path": dataset_name, "use_auth_token": True, "streaming": True}
 dataset_shapes = None
 TaskRegistry.add(
-    'span_corrpution_stream',
+    "NCC_span_corruption_stream",
     source=seqio.FunctionDataSource(
         dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
         splits=("train", "validation"),
-        caching_permitted=True,
+        caching_permitted=False,
         num_input_examples=dataset_shapes,
     ),
     preprocessors=[
@@ -93,10 +124,10 @@ TaskRegistry.add(
                 "targets": None,
             }, target_key="targets"),
         seqio.preprocessors.tokenize,
-        seqio.CacheDatasetPlaceholder(),
+        # seqio.CacheDatasetPlaceholder(),
         preprocessors.span_corruption,
         seqio.preprocessors.append_eos_after_trim,
     ],
-    output_features={"targets": DEFAULT_OUTPUT_FEATURES["targets"]},
+    output_features={"targets": seqio.Feature(vocabulary=vocabulary, add_eos=True)},
     metric_fns=[]
 )

tasks_old.py

@@ -0,0 +1,102 @@
+import functools
+
+import seqio
+import tensorflow as tf
+import t5.data
+from datasets import load_dataset
+from t5.data import postprocessors
+from t5.data import preprocessors
+from t5.evaluation import metrics
+from seqio import FunctionDataSource, utils
+
+TaskRegistry = seqio.TaskRegistry
+
+vocabulary = seqio.SentencePieceVocabulary('gs://t5-data/vocabs/mc4.250000.100extra/sentencepiece.model', extra_ids=0)
+
+DEFAULT_OUTPUT_FEATURES = {
+    "inputs": seqio.Feature(
+        vocabulary=vocabulary, add_eos=True,
+        required=False),
+    "targets": seqio.Feature(
+        vocabulary=vocabulary, add_eos=True)
+}
+
+
+def gen_dataset(split, shuffle=False, seed=None, column="text", dataset_params=None):
+    dataset = load_dataset(**dataset_params)
+    if shuffle:
+        if seed:
+            dataset = dataset.shuffle(seed=seed)
+        else:
+            dataset = dataset.shuffle()
+    while True:
+        for item in dataset[str(split)]:
+            yield item[column]
+
+
+def dataset_fn(split, shuffle_files, seed=None, dataset_params=None):
+    return tf.data.Dataset.from_generator(
+        functools.partial(gen_dataset, split, shuffle_files, seed, dataset_params=dataset_params),
+        output_signature=tf.TensorSpec(shape=(), dtype=tf.string, name=dataset_name)
+    )
+
+
+@utils.map_over_dataset
+def target_to_key(x, key_map, target_key):
+    """Assign the value from the dataset to target_key in key_map"""
+    return {**key_map, target_key: x}
+
+
+# Final pretraining task used in Raffel et al., 2019 adaptated to NCC
+dataset_name = 'NbAiLab/NCC_small'
+dataset_params = {"path": dataset_name, "use_auth_token": True}
+dataset_shapes = {'train': 20830348, 'validation': 473079}
+TaskRegistry.add(
+    'span_corruption',
+    source=seqio.FunctionDataSource(
+        dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
+        splits=("train", "validation"),
+        #caching_permitted=True,
+        num_input_examples=dataset_shapes,
+    ),
+    preprocessors=[
+        functools.partial(
+            target_to_key, key_map={
+                "inputs": None,
+                "targets": None,
+            }, target_key="targets"),
+        seqio.preprocessors.tokenize,
+        #seqio.CacheDatasetPlaceholder(),
+        preprocessors.span_corruption,
+        seqio.preprocessors.append_eos_after_trim,
+    ],
+    output_features={"targets": DEFAULT_OUTPUT_FEATURES["targets"]},
+    metric_fns=[]
+)
+
+# Final pretraining task used in Raffel et al., 2019 adaptated to nbailab_extended
+dataset_name = 'NbAiLab/nbailab_extended'
+dataset_params = {"path": dataset_name, "use_auth_token": True, "streaming": True}
+dataset_shapes = None
+TaskRegistry.add(
+    'span_corrpution_stream',
+    source=seqio.FunctionDataSource(
+        dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
+        splits=("train", "validation"),
+        caching_permitted=True,
+        num_input_examples=dataset_shapes,
+    ),
+    preprocessors=[
+        functools.partial(
+            target_to_key, key_map={
+                "inputs": None,
+                "targets": None,
+            }, target_key="targets"),
+        seqio.preprocessors.tokenize,
+        seqio.CacheDatasetPlaceholder(),
+        preprocessors.span_corruption,
+        seqio.preprocessors.append_eos_after_trim,
+    ],
+    output_features={"targets": DEFAULT_OUTPUT_FEATURES["targets"]},
+    metric_fns=[]
+)

train_base.sh

@@ -1,6 +1,6 @@
 PROJECT_DIR=${HOME}"/models/pk-nb-t5x"
 T5X_DIR="../../t5x"  # directory where the t5x is cloned.
-MODEL_DIR="gs://nb-t5x/pk_nb_t5x_base"
+MODEL_DIR="gs://nb-t5x-us-central2/pk_nb_t5x_base_test1"
 export PYTHONPATH=${PROJECT_DIR}
 
 python3 ${T5X_DIR}/t5x/train.py \