Here is TF code I used for training:
PROJECT=...
ZONE=...
BUCKET=...
TPU=...
MODEL_DIR=...

TASK=natural_questions_open_light_mixture
PRETRAINED_DIR=gs://t5-data/pretrained_models/small
PRETRAINED_STEPS=1000000
FINETUNE_STEPS=15000


# Run fine-tuning
python -m t5.models.mesh_transformer_main \
  --module_import="nqopen_tasks" \
  --tpu="${TPU}" \
  --gcp_project="${PROJECT}" \
  --tpu_zone="${ZONE}" \
  --model_dir="${MODEL_DIR}" \
  --gin_file="dataset.gin" \
  --gin_file="${PRETRAINED_DIR}/operative_config.gin" \
  --gin_param="utils.tpu_mesh_shape.tpu_topology = '8x8'" \
  --gin_param="MIXTURE_NAME = '${TASK}'" \
  --gin_param="utils.run.save_checkpoints_steps=1000" \
  --gin_param="utils.run.batch_size=('tokens_per_batch', 393216)" \
  --gin_param="utils.run.train_steps=$((PRETRAINED_STEPS + FINETUNE_STEPS))" \
  --gin_param="utils.run.init_checkpoint='${PRETRAINED_DIR}/model.ckpt-${PRETRAINED_STEPS}'" \
  --gin_param="utils.run.learning_rate_schedule=@learning_rate_schedules.constant_learning_rate" \
  --gin_param="constant_learning_rate.learning_rate=1e-3" \
  --t5_tfds_data_dir="${BUCKET}/t5-tfds"

where natural_questions_open_light_mixture is the name of the dataset I was training on.