# Auxiliary Job


## Introduction

This page outlines the steps needed to use the auxiliary job capabilities
available in T5X.

## Overview

There are a variety of situations in which running a single job is insufficient
or suboptimal. For example, consider the following scenarios:

+   You want to keep track of evaluation (`infer_eval` or `train_eval`) metrics
    per checkpoint, but evaluation takes a very long time due to having a large
    eval dataset, slow decoding, or multiple tasks to evaluate.

+   You want to finetune every checkpoint on a downstream task as you train.

+   You have customized evaluation code that you want to run on every checkpoint
    as you train, but that does not naturally fit within a `seqio.Evaluator`
    framework.

In cases like these, users can make use of the auxiliary job functionality. At a
high-level, the auxiliary job will launch a new job every time a new checkpoint
is saved. This new job can either re-use the `train.py` binary (e.g. for
continuous finetuning) or a different one. For example, this allows users to
perform continuous evaluation (using `eval.py`) without slowing down the
training job. We will provide detailed examples showing how to use the auxiliary
job for these use-cases.

When this new job is launched, the controller will replace four gin macros:
`MODEL_DIR`, `MIXTURE_OR_TASK_NAME`,`INITIAL_CHECKPOINT_PATH`, `TRAIN_STEPS`.
The second of these is set by the user-controlled flag (more on this below), and
the third one is equal to the last checkpoint seen. Aside from this, users are
free to modify the configuration as needed. Beyond gin macros, the auxiliary job
can also have different resource requirements, priority, and even cell placement
from the train job.

## Example 1: Separate evaluation job.

### Step 1: Choose a model architecture.

Similar to pretraining, we will need some gin configuration. For this example,
we will use the T5-1.1-Base model.

### Step 2: Choose a SeqIO Task/Mixture for training and evaluation.

In this example, we will use the classic task of English-French translation from
WMT14, which is conveniently available as a SeqIO task in the tasks file from
the T5 tasks under the name `'wmt_enfr14_v003'`.

### Step 3: Write a Gin config.

Unlike pretraining or finetuning, we will need two gin files for this setup: one
for the training job, and one for the auxiliary job. The train gin file will
have the same requirements as the gin file for pretraining or finetuning. The
auxiliary job gin file can leverage these gin files or be its own independent
gin file, depending on the user’s choice. For this example, we will make a new
gin which is mostly a wrapper around `pretrain.gin` with some additional
hardcoded features. We will use this gin file for the train job and `eval.gin`
for the auxiliary job.

### Step 4: Launch your experiment.

Our sample script will be quite similar to the one used in pretraining and
finetuning, but with a few additional flags which we describe below.

+   `auxiliary_job_mixtures`: This is a comma-separated list of mixtures. A
    separate auxiliary job will be run for each mixture and will replace the gin
    macro `MIXTURE_OR_TASK_NAME`. Note that you need this flag even if you are
    using a custom binary, which does not need a mixture since otherwise no
    auxiliary job will run.

+   `auxiliary_job_gin_file`: This is identical to `gin_file`, except it is used
    for the auxiliary job instead of the train job.

+   `replace_gin_file`: If True, this auxiliary launcher will not use any of the
    gin files from train job. This is necessary when using a binary different
    from `train.py`, since the top-level functions will not match.

+   `auxiliary_job_cell`: The cell in which to run your job. Note that this can
    be different from the training cell.

+   `auxiliary_job_platform`: The platform to use for the auxiliary. Note that
    this can be different from the one use for the train job, allowing users to
    use smaller configurations for evaluation than needed for training.

+   `auxiliary_job_build_target`: The binary to use for auxiliary job.

+   `final_auxiliary_job_steps`: This flag controls how many additional steps to
    take when using the auxiliary job for finetuning. Setting to 0 enables
    continuous evaluation.

We provide the sample script below.

```sh
declare -a ARGS=(
--cell=iz
--platform=jd=2x2
--final_auxiliary_job_steps=0
--replace_gin_file=True
--auxiliary_job_mixtures=wmt14_enfr_v003
--auxiliary_job_gin_file=t5x/examples/t5/t5_1_1/examples/base_wmt14enfr_eval.gin
--auxiliary_job_cell=iz
--auxiliary_job_platform=jd=2x2
--auxiliary_job_build_target_path=//t5x:eval
--gin_file=t5x/examples/t5/t5_1_1/examples/base_wmt14enfr_train.gin
)

gxm t5x/google/xm_launch.py "${ARGS[@]}"
```

## Example 2: Continuous finetuning job.

In this example, we will be pretraining a model on a span corruption task on the
C4 dataset, and finetuning it on the WMT'14 English-French translation task. As
before, we will launch a new auxiliary job once every checkpoint is saved.
However, instead of using the `eval.py` binary, we will use the `train.py`
binary.

### Step 1: Choose a model architecture.

We will use the T5-1.1-Base model as in the previous example.

### Step 2: Choose a SeqIO Task/Mixture for training and evaluation.

For pretraining, we re-use the span coprruption task `c4_v220_span_corruption`
available in the T5 mixtures `tasks.py` file.

### Step 3: Write a Gin config.

As before, we need our gin files to contain all the desired macros in them. We
thus create two new gin files: `base_c4_pretrain.gin` for the train job and
`base_wmtenfr14_finetune.gin` for the auxiliary job.

### Step 4: Launch your experiment.

Our script is quite similar to the first example, with the same flags as before
but with the appropiate changes. The main distinction is that we must change the
flag `final_auxiliary_job_steps` to be non-zero to start finetuning. We will
settle for a modest 200 steps for the sake of demonstration (and evaluate every
100 steps), but users should use larger steps in realistic scenarios. We also
use `train.py` binary instead of `eval.py`.

We provide the sample script below.

```sh
declare -a ARGS=(
--cell=iz
--platform=jd=2x2
--final_auxiliary_job_steps=200
--replace_gin_file=True
--auxiliary_job_mixtures=wmt14_enfr_v003
--auxiliary_job_gin_file=t5x/examples/t5/t5_1_1/examples/base_wmt14enfr_finetune.gin
--auxiliary_job_cell=iz
--auxiliary_job_platform=jd=2x2
--auxiliary_job_build_target_path=//t5x:train
--gin_file=t5x/examples/t5/t5_1_1/examples/base_c4_pretrain.gin
)

gxm t5x/google/xm_launch.py "${ARGS[@]}"
```

## Common Gotchas.

We outline a few common error patterns that we have encountered.

+   **Not passing a value for the `auxiliary_mixtures` flag.** Even if you have
    the desired task in your gin file, or you use a differently named macro, you
    should still pass a value for this flag, since launch script will launch a
    new job per value of this flag.

+   **Not setting `replace_gin_file=True` when using a different binary from
    train.py.** This will usually yield an error that there is no `train`
    function.

+   **No metrics being logged.** It can be tempting to use gin files usually
    used for evaluation. However, one must ensure that the corresponding SeqIO
    evaluators still log to the tensorboard, otherwise you won’t see the
    metrics.

+   **Slow `train_eval`.** While the approach outlined above separates out the
    infer_eval job, it may be that even train_eval is too slow. In these
    situations, we suggest adding the metrics from train_eval into the
    `metrics_fn` argument of the SeqIO task and have them be computed in the
    auxiliary job as well. To do this with teacher forcing, you will have to
    use `train.py` instead of `eval.py`.

+   **Using `CHECKPOINT_PATH` rather `INITIAL_CHECKPOINT_PATH`.** For legacy
    reasons, the auxiliary job uses the macro `INITIAL_CHECKPOINT_PATH` rather
    than `CHECKPOINT_PATH` as found in `eval.gin`. Make sure to use the latter
    macro building your gin scripts.

+   **Gin macros being ignored when passed through the format
    `gin.{MACRO}={VAL}`.** In the current setup, you must include all gin macros
    in the gin script. Attempting to pass them as additional flags will usually
    not work.

+   **Not setting `final_auxiliary_job_steps=0` when performing continuous
    evaluation.** The current parameter controller uses this as a check. When
    this is true, it will replace the `EVAL_OUTPUT_DIR` folder with the current
    `MODEL_DIR`, so that the evaluation metrics are saved in the right place and
    the metrics are showed correctly on the tensorboard.
