OptimisationGENERALFRAMEWORK-SPECIFIC

Epoch, Batch, Step

An epoch is one pass over the data. A batch is the subset used for one gradient estimate. A step is one parameter update. "We trained for ten epochs" says nothing until you know the batch size.

Target & dataWhat to measureWhat must stay true

The problem, the obvious approach, and why it breaks

Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.

The question

Two runs both "trained for ten epochs" and one is far better than the other. What is the unit of training actually being counted, and why is an epoch not it?

The problem

A team reproduces a colleague's training run from the notes "10 epochs, learning rate 3e-4, cosine schedule" on a bigger GPU and gets a worse model. The notes were followed exactly. They want to know what was left out of the notes.

The obvious approach

An epoch is a natural unit — the model has "seen all the data once" — so training length is expressed in epochs, and the schedule and checkpoints are laid out in epochs too. Batch size is a memory setting chosen to fill the GPU.

Why it breaks

Quadrupling the batch size quartered the number of updates. The second run made a quarter as many steps with the same learning rate, and the cosine schedule — defined in epochs — decayed the rate over a quarter as many updates. The model is undertrained and the notes contain no hint of it.

How it breaks — usually after the offline metric looked fine
  • Quadrupling the batch size quartered the number of updates. The second run made a quarter as many steps with the same learning rate, and the cosine schedule — defined in epochs — decayed the rate over a quarter as many updates. The model is undertrained and the notes contain no hint of it.
  • Checkpoints and evaluation were tied to epoch boundaries. On a dataset that grew between runs, "epoch 5" now means a different number of updates than it did, and the learning-curve comparison between the runs is between different points in training.
  • A warm-up "of one epoch" became a warm-up of a quarter as many steps on the big GPU; the early large steps that warm-up exists to soften were taken at nearly full rate, and the run was unstable in a way the first never was.
ProblemTargetDataRepresentationSplitModelTrainingEvaluationValidationDeploymentInferenceMonitoringDriftRetraining

What is being predicted, and from what data

This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.

Target
  • The optimiser's target is the loss after a number of *parameter updates*; the notes express training length in passes over the data, which is a different unit and only convertible with the batch size.
  • The surrounding system, a text classifier, is unchanged; the question is whether the second run performed the same optimisation as the first.
Data
  • A training set of n examples. Each epoch visits every example once, in shuffled order, grouped into batches; each batch produces one gradient estimate and one update.
  • The bigger GPU held four times the batch. The dataset, the epoch count and the learning rate were the same. Steps per epoch were therefore a quarter, and so was the total number of updates.

How it actually works

Precisely enough to predict its behaviour — not a framework API.

  • Three units, three things. An epoch is one traversal of the training set, a data-loader fact. A batch is the set of examples whose gradients are averaged into one estimate, a memory-and-noise fact. A step is one application of the update rule, an optimiser fact. Steps per epoch equals n / batch_size; total steps equals epochs times that.
  • The optimiser only ever counts steps. A learning-rate schedule, a warm-up, a momentum estimate and an Adam moment are all functions of the step index; the optimiser has no notion of an epoch. Describing a schedule in epochs is a convenience that silently multiplies by n / batch_size.
  • Two runs with the same epoch count and different batch sizes are therefore different optimisations: different numbers of updates, different gradient noise per update, and — if the rate was not adjusted — different effective progress per pass (Batch Size and Learning Rate).

Three units, one conversion

The optimiser counts updates. The data loader counts passes. The batch size is the exchange rate between them, and any statement of training length that leaves it out is incomplete in the way a price without a currency is incomplete.

The arithmetic is trivial and the consequences are not: a schedule, a warm-up, a checkpoint interval and an early-stopping patience defined in epochs all change their meaning when the batch size changes, without any line of the config that mentions them being edited.

Express the run in the optimiser's unit
1import math
2
3n = 1_200_000 # training examples
4batch = 256 # examples per gradient estimate
5epochs = 10
6
7steps_per_epoch = math.ceil(n / batch) # 4688
8total_steps = epochs * steps_per_epoch # 46_880 parameter updates
9
10# the schedule is a function of the step index, never of the epoch
11def lr_at(step, peak=3e-4, warmup=2000):
12 if step < warmup:
13 return peak * step / warmup
14 progress = (step - warmup) / max(1, total_steps - warmup)
15 return 0.5 * peak * (1 + math.cos(math.pi * progress))
16
17# same n, same epochs, batch 1024 -> 11_720 updates, a quarter of the training

The last line is the reproduction failure from the problem. Nothing in the notes — epochs, rate, schedule — changed; the number of updates did, and the schedule compressed with it.

The same notes, a different optimisation

The reproduction run followed the notes and performed a quarter of the optimisation. The offline evaluation reported the result honestly — the model was worse — and the team looked for the cause in the data and the code, because the config was "the same".

Two ways to write down a training run
In epochs
"10 epochs, lr 3e-4, cosine, warm-up 1 epoch." Reproducible only on the same batch size and the same dataset size, neither of which is written down.
In steps, with the exchange rate
"46,880 steps at batch 256 on 1.2M examples (10 passes); peak lr 3e-4; cosine over steps; warm-up 2,000 steps." Reproducible on any hardware; the derived epoch count is there for intuition.

The optimiser, the schedule and the moment estimates are all functions of the step index. Writing the run in that unit means the config describes the optimisation that happened, not a data-loader statistic that happens to correlate with it.

Pipelines that count the wrong thing

A retraining pipeline configured in epochs has a hidden dependency on the dataset size. Data accumulates, so each retraining run is a little longer than the last; a deduplication or a retention policy makes one run abruptly shorter. The step count drifts and nobody changed the config.

The assumption is that training length is what it was when the recipe was validated. It is checkable at start-up in one line, and almost never checked.

must stay trueTraining length means what it meant

Each retraining run performs the number of parameter updates the recipe was validated with, on a schedule that spans those updates.

holds when Length and schedule are configured in steps, or the pipeline recomputes them from the current dataset size and batch size and logs the result.

breaks when The dataset grows or shrinks under a fixed epoch count; the batch size is changed to fit new hardware; gradient accumulation is added and the step counter counts micro-batches.

how you would know A start-up log of n, batch size, steps per epoch and total steps compared against the previous run; a schedule plot checked at the first checkpoint; an alert when total steps moves without a config diff (Experiment Tracking).

respond Fix the unit in the config, not the model. Then decide deliberately what to hold constant when the data grows, and record the decision.

How to build it

Most important first.

  • Specify training length in steps, and schedules, warm-up and checkpoints in steps. Report epochs as a derived quantity if it is useful to know how many times the data was seen.
  • Record batch size, dataset size and total steps together in every experiment record, so any two runs can be compared in the unit the optimiser used (Experiment Tracking).
  • When batch size changes, decide explicitly what to hold constant — total steps, total examples seen, or examples-per-update times rate — and say which, because they cannot all be held constant at once.
  • Evaluate on a step schedule, not at epoch boundaries, so learning curves from runs with different batch sizes line up on the same axis (Learning Curves).

What to measure

Which number actually maps to the decision — and which numbers look relevant and are not.

  • Held-out loss against steps, and against examples seen, on the same plot. The two runs in the problem overlap on the examples axis and diverge on the steps axis, which is the whole explanation.
  • Steps per second and examples per second as separate throughput numbers; a bigger batch raises the second and lowers the first, and which one matters depends on what is holding training back.
  • Do not compare runs by "epochs trained". Without batch size it is not a quantity.

What must stay true after deployment

The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.

Assumptions
  • The batch size, the dataset size and the schedule unit recorded with a run are the ones it actually used, so its step count can be reconstructed and compared with the next run.
  • The retraining pipeline holds training length constant in steps — or deliberately in examples seen — when the dataset grows, rather than letting a fixed epoch count silently change the optimisation.
  • Every place that counts steps — the schedule, the logging, the checkpointing — counts parameter updates, not micro-batches or data-loader iterations.
How to verify — offline, online, and over time
  • Offline: assert in the training script that the schedule's total length in steps equals epochs times ceil(n / batch_size), and log all four numbers at start-up.
  • During training: log the learning rate against the step index and check the curve against the intended schedule at the first checkpoint; a warm-up that ended too early is visible immediately.
  • Over time: compare consecutive retraining runs on total steps and examples seen, and alert when either moves without an explicit config change.

What can go wrong

Failure modes in production
  • A dataset that grows week over week under a fixed epoch count silently lengthens the run and shifts the schedule; a dataset that is deduplicated shortens it. Both runs report the same epoch count.
  • The last batch of an epoch is smaller than the others and, if the loss is summed rather than averaged, produces a smaller gradient on every epoch boundary — a periodic artefact in the loss curve that sends people looking for a data problem.
  • Gradient accumulation is used to simulate a large batch on a small GPU, and the logging counts each micro-batch as a step; the recorded step count is several times the number of updates and every schedule is off by the same factor.
What the recommended approach costs
  • Step-based schedules require knowing the dataset size before training starts, which is awkward for streamed or sharded data and needs a size estimate or a count pass.
  • Reporting in steps loses the intuition that an epoch carries — "the model has seen everything once" is a useful sanity check for overfitting and for whether the data loader is exhausted.
  • Holding total steps constant under a larger batch costs proportionally more compute; holding examples seen constant undertrains unless the rate is rescaled. There is no free conversion.
Misreads
  • "We trained for the same number of epochs, so the runs are comparable." They saw the same data the same number of times. They made different numbers of updates and followed different schedules, and the optimiser only knows about updates.
  • "A bigger batch is faster, so we train more in the same wall-clock." A bigger batch processes more examples per second and fewer steps per second. Whether that is more training depends on what the rate and schedule do with it.
  • "The loss dips at every epoch boundary — the data must be sorted somehow." Check the last-batch size and the loss reduction first; a short final batch under a summed loss is the common cause.

Where this applies

ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.

  • GENERALThe three units are properties of mini-batch training itself and mean the same thing for any model trained by stochastic gradient descent; the definitions do not change with the model family or the data.
  • FRAMEWORK-SPECIFICWhether a framework counts steps as optimiser updates or as data-loader iterations differs, and gradient-accumulation wrappers differ in which one they expose to the scheduler and the logger; check the one in use rather than assuming.

Where the depth lives

This domain teaches the model and hands the rest off by name.

Computer Architecturegpu-parallelism