Decide when to stop a training trial

At epoch four, one configuration has the lowest validation loss. Another is still improving slowly. Stopping the second run saves compute, but it also removes the evidence that would tell you whether the ranking changes after warmup or a learning-rate transition.

Early stopping makes a prediction about a run's unfinished future. The useful question is how often that prediction discards configurations you would have wanted to keep at the full training budget.

This guide develops a calibration procedure around Polyaxon experiments: collect full learning curves, replay proposed stopping rules, inspect their mistakes, and record decisions before allowing them to stop new runs. It extends the PyTorch tuning walkthrough, whose example deliberately gives every trial the same training budget.

Constructed example: stopping B at epoch four hides its later improvement. The chart shows quality increasing upward; the code below uses loss, where lower is better.

Choose what progress means

A useful intermediate metric needs a common training coordinate. Wall-clock time mixes model quality with startup costs, hardware speed, and input-pipeline throughput. “Epoch four” can also represent different numbers of optimizer updates when batch size or the dataset changes.

For a study with fixed data, batch size, and evaluation frequency, epoch can be a reasonable coordinate. If those vary, consider optimizer updates, examples, or tokens processed. State what is held equal and what remains different. Matching token count, for example, does not make two model sizes equally expensive.

Use the same validation split and metric implementation at every comparison. Keep the final test split outside both threshold calibration and model selection. Record the learning-rate schedule and warmup alongside each run; a low early rank during a deliberately slower schedule needs a different interpretation from a flat curve after training has settled.

Resource allocation is the central idea behind methods such as Hyperband. The procedure here is a smaller, explicit experiment for assessing a stopping rule on your own workload. It does not assume a particular search algorithm will produce the same decisions.

Start with complete pilot runs

Choose a modest pilot that covers the training behaviors you intend to search. Include the schedules or batch sizes that might learn slowly, rather than calibrating only on the fastest starters. Run those configurations to the same final budget with quality-based pruning disabled.

Polyaxon mapping studies make a small, declared comparison easy to inspect. This is a matrix fragment for a training component that accepts learning_rate and seed inputs; retain that component and its fixed dataset, batch size, schedule, and epoch budget:

matrix:
  kind: mapping
  concurrency: 2
  values:
    - learning_rate: 0.0003
      seed: 19
    - learning_rate: 0.0003
      seed: 23
    - learning_rate: 0.001
      seed: 19
    - learning_rate: 0.001
      seed: 23

These four recipes illustrate the mapping; they are not enough to establish that a pruning policy is reliable across a larger search space. Expand the pilot to represent that space. A timeout and failure policy can still bound broken jobs; keep those operational limits separate from quality decisions.

Log validation loss at every evaluation point using the actual progress coordinate. In an epoch-based trainer:

from polyaxon import tracking


def report_validation(epoch: int, validation_loss: float) -> None:
    tracking.log_metrics(validation_loss=validation_loss, step=epoch)

Call this inside your initialized tracking run after validation. Keep a complete per-epoch history in the run's artifacts as well, so a later policy review can use the original values. Synchronized charts help inspect curves at a shared step, but plotting the same step does not establish that the underlying training budgets are equivalent.

Replay a rule against the futures you observed

Start with a small rule you can explain. For example: after a chosen minimum epoch, stop a trial if its current validation loss exceeds a precomputed ceiling at one of the declared comparison epochs. Lower loss is better in this example.

Build those ceilings from a calibration subset of the full pilot histories. A ceiling is a policy choice, not a universal accuracy requirement. Choose its comparison epochs and tolerance, then replay the frozen rule on separate pilot configurations or later full-budget runs. Keep repeated seeds from the same configuration together when separating calibration from review.

For each replayed trial, find the first point where the rule would stop it. Record two consequences:

• Work potentially avoided: the training progress between that point and the full budget. This is a replay estimate; actual GPU time also depends on runtime and stopping delay.
• Useful configurations lost: full-budget finalists that the rule would have stopped. Define “finalist” before inspecting the replay, using the selection criterion your study actually needs.

Inspect the individual losses, not just an aggregate percentage. A rule that repeatedly rejects a whole schedule family is changing the effective search space. If the pilot provides too little evidence to assess that risk, use a later comparison point or keep those trials at full budget.

Record decisions before enforcing them

The following helper implements the declared epoch-based rule. Supply a frozen mapping of comparison epochs to finite loss ceilings from your calibration record. An unlisted epoch produces no pruning decision; a non-finite observed loss is an execution problem that needs diagnosis.

import math
from collections.abc import Mapping


def exceeds_loss_ceiling(
    epoch: int,
    validation_loss: float,
    loss_ceilings: Mapping[int, float],
    min_epoch: int,
) -> bool:
    if epoch < 1 or min_epoch < 1:
        raise ValueError("Epochs must be positive")
    if not math.isfinite(validation_loss):
        raise FloatingPointError("Validation loss is not finite")
    if epoch < min_epoch or epoch not in loss_ceilings:
        return False
    ceiling = loss_ceilings[epoch]
    if not math.isfinite(ceiling):
        raise ValueError("Loss ceilings must be finite")
    return validation_loss > ceiling

Evaluate the rule after validation, and save a recoverable checkpoint before acting on it. This excerpt belongs inside your existing training loop; train_one_epoch, validate, and save_checkpoint are the trainer's own functions. loss_ceilings, min_epoch, and policy_revision come from the saved policy, and pruning_mode must be either observe or enforce.

from polyaxon import tracking

if pruning_mode not in {"observe", "enforce"}:
    raise ValueError("Unknown pruning mode")

tracking.log_inputs(
    pruning_policy=policy_revision,
    pruning_mode=pruning_mode,
    min_epoch=min_epoch,
    total_epochs=total_epochs,
)
training_complete = True
for epoch in range(1, total_epochs + 1):
    train_one_epoch()
    validation_loss = validate()

    proposed_stop = exceeds_loss_ceiling(
        epoch, validation_loss, loss_ceilings, min_epoch
    )
    save_checkpoint(epoch)
    tracking.log_metrics(
        validation_loss=validation_loss,
        proposed_stop=int(proposed_stop),
        step=epoch,
    )

    if pruning_mode == "enforce" and proposed_stop and epoch < total_epochs:
        training_complete = False
        break

tracking.log_outputs(
    training_complete=training_complete,
    stopped_epoch=epoch,
    pruning_policy=policy_revision,
    pruning_mode=pruning_mode,
)
if training_complete:
    tracking.log_outputs(final_validation_loss=validation_loss)

The helper must be in scope, the trainer must initialize tracking, and total_epochs must be a positive integer. save_checkpoint must persist enough state for the recovery boundary you need; checkpoint recovery explains that part. Handle broader training exceptions through your normal failure path rather than publishing completion outputs from a finally block.

In observe mode, the trial records every proposed stop and continues to the full budget. That gives you fresh completed histories against which to review the frozen policy. In enforce mode, the trainer finishes normally after a deliberate stop, so the run can have a Succeeded status. Its training_complete=false output identifies the shorter experiment. Rank full-budget candidates using final_validation_loss, which the excerpt publishes only for complete training.

Save the policy mapping, its revision, the pilot run IDs, and the comparison protocol with the study. Validate the complete policy file before training, including the epoch coordinate and all thresholds. A missing evaluation must not be silently replaced with an older metric value.

Choose who owns the stop

This example lets the trainer make an exact decision at its own evaluation boundary, using a fixed reference from earlier runs. Polyaxon records the metrics, outputs, checkpoints, and resulting run state. The frozen reference does not adapt as the new wave finishes.

Polyaxon also offers native median, difference, and truncation policies through metric early stopping. Use those when their comparison behavior fits the study. Do not assume that setting step=epoch turns a native policy into the fixed, matched-epoch rule implemented above. Avoid enabling both owners for the same pruning experiment, because a second stop path makes its decisions harder to audit.

A metric threshold without a pruning policy can finish the whole study and stop its running and pending children. Reaching an acceptable quality target is a different objective from deciding which unfinished trial deserves more training.

Keep evidence after the first pruned wave

Review completed, deliberately stopped, and failed trials separately. A failed data loader has not established that its hyperparameters are poor. An epoch-six loss does not become a full-budget result because no later value exists.

Reserve part of the budget for full training of selected would-be-stopped configurations. Choose that audit sample in advance, including some cases beyond those that look almost acceptable. Their complete histories help reveal whether the policy is discarding useful regions. Revisit the rule when the dataset, model family, schedule, or training coordinate changes.

Confirm the remaining finalists across matched training seeds at the full budget before the final held-out assessment. Keep the pruning policy fixed during that confirmation. The evidence you need is both the work the policy avoids and the useful configurations it excludes.