# Determine the number of steps for each training run

There are two types of training workloads:

- compute-bound
- not compute bound

**Compute-bound** training is limited by how much time you can spend
on training, not by how much training data you have or some other factor.
In other words, the "optimal" training time is always "as long as you
can afford." If you somehow can train longer or more efficiently,
training loss should drop. (With proper tuning, validation loss should also
drop.)

Speeding up compute-bound training is equivalent to improving training.
That said, just because a workload is compute-limited doesn't mean
training longer or faster is the only way to improve results.

When training is not compute-bound, you can afford to train as long as you
would like to. However, training a model longer might not help much or might
even cause [overfitting](https://developers.google.com/machine-learning/glossary#overfitting). When
training is not compute-bound:

- You can train to very low training loss, to the point where additional training might slightly reduce the training loss but does not meaningfully reduce the validation loss.
- You can tune more easily, especially when tuning learning rate decay schedules, since they have a particularly strong interaction with the training budget. In contrast, getting a low training loss on compute-bound training might require a learning rate decay schedule tuned to perfection.

Regardless of whether a given workload is compute-bound or not, methods that
increase the variance of the gradients (across batches) typically slow
training progress, and thus may increase the number of training steps
required to reach a particular validation loss. Any of the following
can cause a high gradient variance:

- Using a smaller batch size.
- Adding data augmentation.
- Adding some types of regularization (for example, dropout regularization).

## Decide how long to train when training is *not* compute-bound

> Your Goal: Train long enough for the model to reach the best possible result
> without wasting training steps.

Your main goal is to ensure that you train long enough for the model to reach
the best possible result without wasting unnecessary training steps.
When in doubt, err on the side of training longer. Your evaluation metrics (for
example, precision, recall, AUC, or F~1~) should never
degrade when training longer, assuming you properly use retrospective
checkpoint selection and you checkpoint frequently enough.

Never tune the `max_train_steps` number in a study. Instead, pick a value and
use that same value for all trials. From these trials, plot the training step
that retrospective checkpoint selection finds in order to refine the choice of
`max_train_steps`.

For example, if the best step is always during the first
10% of training, then the maximum number of steps is way too high.
Alternatively, if the best step is consistently in the last 25% of training,
you might benefit from training longer and retuning the decay schedule.
The ideal number of training steps can change when the architecture or data
changes (for example, adding data augmentation).
The next section describes how to pick an initial candidate value for
`max_train_steps` based on the number of steps necessary to "perfectly fit"
the training set using a constant learning rate.

> [!NOTE]
> **Note:** We are not using the phrase "perfectly fit the training set" in a precise or mathematically well-defined way. The phrase is merely meant as an informal descriptor to indicate a very low training loss. For example, when training with the log loss, absent regularization terms, you might see the training loss slowly improve until training reaches floating point precision limits. That's because the network weights grow without bound and the model's predictions on the training set become increasingly confident. In this case, the model "perfectly fit" the training set around the time the misclassification error reached zero on the training set. You might need to increase the starting value for `max_train_steps` if the amount of gradient noise in the training procedure increases; for example, when data augmentation or regularizers like dropout are introduced to the model.

It may be possible to decrease `max_train_steps` if the training process
improves somehow; for example, with a better-tuned optimizer or a
better-tuned learning rate schedule.

### Algorithm for picking an initial candidate for max_train_steps using a learning rate sweep

You can pick an initial candidate for `max_train_steps` with
a learning rate sweep algorithm. The following algorithm assumes it is
possible to not only "perfectly" fit the training set, but also
to do so using a constant learning rate schedule.

1. If it is possible to perfectly fit the entire training set, then there must exist a configuration (with some value of `max_train_steps`) that perfectly fits the training set. Find any such configuration and use its value of `max_train_steps` as a starting point `N`.
2. Run a constant learning rate sweep (that is, grid search the learning rate) without data augmentation and without regularization where each trial trains for `N` steps. The number of steps required for the fastest trial in the learning rate sweep to reach perfect training performance should be your initial guess for `max_train_steps`.

NOTE: Bad search spaces can lead to self-deception.
For example, if all the learning rates in a study are too small, you might
incorrectly conclude that a very large value of `max_train_steps` is necessary.
At a minimum, check that the optimal learning rate in the study is not at
the boundary of the search space.

## Decide how long to train when training is compute-bound

In some cases, training loss keeps improving indefinitely, so your patience and
your computational resources become the limiting factors. But should you train
as long as you can afford? Not necessarily. Consider the following:

- You might be able to tune more effectively by running a larger number of shorter experiments, reserving the longest "production length" runs for the models you hope to launch.
- As the training time for trials approaches your patience limit, tuning experiments become more relevant for your potential launch candidates, but you can complete fewer of them.
- You can probably answer many questions while only training for \~10% of the production length. However, your conclusions at this time limit might not apply to experiments at 20% of the production length, let alone 100%.

Tuning during multiple rounds with increasing, per-trial training step limits
is a sensible approach. You can run as many rounds as you want, but usually
1-3 rounds are the most practical. Essentially, try to obtain as much
understanding of the problem as possible using trials with a very quick
turnaround time, trading off the following:

- Tuning thoroughness.
- Relevance to the final, longest runs.

Once a given per-trial time limit has generated useful insights, increase
the training time and continue tuning, double-checking your conclusions
from the shorter runs as needed. As a starting point, we recommend two
rounds of tuning:

- Round 1: Shorter duration runs to find good model and optimizer hyperparameters.
- Round 2: Very few long duration runs on good hyperparameter points to get the final model.

The biggest question going from Round 1 to Round 2 is:
> How do you adjust learning rate decay schedules.

One common pitfall when adjusting learning rate schedules between rounds is
using all the extra training steps with a learning rate that is too small.

### Round 1: a lot of short training runs

Unfortunately, there is no guarantee that good hyperparameters found in
short, incomplete training are still good choices when you significantly
increase training length. However, for some hyperparameters, the good choices
are often correlated *enough* for Round 1 to be useful.
What hyperparameter values found in shorter runs transfer successfully
to longer training runs? We don't know; we need more research.
But based on what we know so far, here are our suspicions in decreasing
probability of transferring:

- **Very likely to transfer.** Early training instability can be resolved in the first round of tuning using a smaller number of training steps. The following hyperparameters are the likeliest to transfer:
  - Warmup length
  - Initialization
- **Likely to transfer.** A dramatic win in the model architecture usually transfers, but many counterexamples are probable.
- **Might transfer.** The following hyperparameters might transfer:
  - Optimization algorithm and hyperparameters would "loosely" transfer.
  - Data augmentation.
  - Regularization. If it isn't possible to perfectly fit the training set, the model might be in a regime where regularization is unlikely to help very much.
- **Unlikely to transfer.** Learning rate schedule is unlikely to transfer perfectly. [Training Compute-Optimal Large Language
  Models](https://arxiv.org/abs/2203.15556) suggests that even decay schedule transfers, but we don't believe this is true in general. For example, tuning sqrt decay on a small number of training steps and then extending to a large number causes the majority of training to occur at overly small steps. You can likely do "good enough" with most schedules in the limit of extreme training budget, but you will likely see noticeable performance improvements if it is tuned. [Understanding Short-Horizon Bias in Stochastic
  Meta-Optimization](https://arxiv.org/abs/1803.02021) describes the dangers of trying to pick learning rates myopically.

### Round 2: fewer runs, but of longer duration

Run the best hyperparameter configuration from Round 1.

**Speculation** : 🤖 Use the extra steps to extend the period of training
at a high learning rate. For example, if you are using a linear schedule,
then keep the length of the decay fixed from Round 1 and extend the period
of constant `lr` in the beginning. For cosine decay, keep the base `lr`
from Round 1 and extend `max_train_steps` as described in
[Training Compute-Optimal Large Language
Models](https://arxiv.org/abs/2203.15556).

Additional training rounds might make sense for teams with all of the following:

- Very mature modeling
- Tuning pipelines
- Very long and expensive production training runs

However, additional training runs are often unproductive.

We've already described how to transfer from Round 1 to Round 2.
If you don't care about analysis time and if making efficient use
of computing resources is your overriding concern, then we recommend
exponentially increasing the length of training runs (and thus the
end-to-end time to complete a study) over many different rounds of tuning:

- At each round, systematically ensure that your choices continue to provide good results.
- Put new ideas through a pipeline that progressively derisks them using increasingly long-running experiments from Step i to Step i+1.