finetuning-anlg.sh

declare -a sizes=("11B" )

TASK=anlg
PRETRAINED_STEPS=1000000
FINETUNE_STEPS=20000
export BUCKET=gs://danielk-files/t5-models
export TPU_NAME=...
for SIZE in "${sizes[@]}"; do
  PRETRAINED_DIR="gs://t5-data/pretrained_models/${SIZE}"
  MODEL_DIR="${BUCKET}/${TASK}/${SIZE}"

  # Run fine-tuning
 python -m t5.models.mesh_transformer_main \
   --module_import="genie" \
   --tpu="${TPU_NAME}" \
   --gcp_project="${PROJECT}" \
   --tpu_zone="${ZONE}" \
   --model_dir="${MODEL_DIR}" \
   --gin_file="dataset.gin" \
   --gin_file="${PRETRAINED_DIR}/operative_config.gin" \
   --gin_param="utils.run.save_checkpoints_steps=1000" \
   --gin_param="utils.tpu_mesh_shape.tpu_topology = 'v3-128'" \
   --gin_param="MIXTURE_NAME = '${TASK}'" \
   --gin_param="utils.run.batch_size=('tokens_per_batch', 24576)" \
   --gin_param="utils.run.train_steps=$((PRETRAINED_STEPS + FINETUNE_STEPS))" \
   --gin_param="utils.run.init_checkpoint='${PRETRAINED_DIR}/model.ckpt-${PRETRAINED_STEPS}'" \
   --t5_tfds_data_dir="${BUCKET}/t5-tfds"

  # Run eval
  python -m t5.models.mesh_transformer_main \
    --module_import="genie" \
    --tpu="${TPU_NAME}" \
    --gcp_project="${PROJECT}" \
    --tpu_zone="${ZONE}" \
    --model_dir="${MODEL_DIR}" \
    --gin_file="dataset.gin" \
    --gin_file="${MODEL_DIR}/operative_config.gin" \
    --gin_file="eval.gin" \
    --gin_param="utils.tpu_mesh_shape.tpu_topology = 'v3-128'" \
    --gin_param="MIXTURE_NAME = '${TASK}'" \
    --gin_param="utils.run.dataset_split = 'dev'" \
    --gin_param="utils.run.batch_size=('tokens_per_batch', 24576)" \
    --gin_param="utils.run.eval_checkpoint_step='all'" \
    --t5_tfds_data_dir="${BUCKET}/t5-tfds"

  python -m t5.models.mesh_transformer_main \
    --module_import="genie" \
    --tpu="${TPU_NAME}" \
    --gcp_project="${PROJECT}" \
    --tpu_zone="${ZONE}" \
    --model_dir="${MODEL_DIR}" \
    --gin_file="dataset.gin" \
    --gin_file="${MODEL_DIR}/operative_config.gin" \
    --gin_file="eval.gin" \
    --gin_param="utils.tpu_mesh_shape.tpu_topology = 'v3-128'" \
    --gin_param="MIXTURE_NAME = '${TASK}'" \
    --gin_param="utils.run.dataset_split = 'test'" \
    --gin_param="utils.run.batch_size=('tokens_per_batch', 24576)" \
    --gin_param="utils.run.eval_checkpoint_step='all'" \
    --t5_tfds_data_dir="${BUCKET}/t5-tfds"
done