test.sh

# Model dir to save logs, ckpts, etc. in "gs://model_dir" format.
MODEL_DIR="gs://bigscience/experiment_d/multilingual_t0/test-run/model"

# Data dir to save the processed dataset in "gs://data_dir" format.
TFDS_DATA_DIR="gs://bigscience/experiment_d/multilingual_t0/test-run/data"
T5X_DIR=${HOME}"/t5x"  # directory where the T5X repo is cloned.

python3 ${T5X_DIR}/t5x/train.py \
  --gin_file=${T5X_DIR}"/t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_finetune.gin" \
  --gin.MODEL_DIR="'${MODEL_DIR}'" \
  --tfds_data_dir=${TFDS_DATA_DIR} \
  --gin.partitioning.ModelBasedPjitPartitioner.model_parallel_submesh="(4,1,1,1)"