| # 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. | |