GPU Hyperparameter Tuning, in depth: systems knobs vs optimisation knobs, batch and learning-rate coupling, precision traps and budget-fair search

Two families of knobs, one identity

The single identity to keep in your head is:

global_batch = micro_batch_per_gpu * grad_accum_steps * data_parallel_size
tokens_per_step = global_batch * sequence_length

The optimiser only sees the averaged gradient over the global batch. So if you hold global batch, sequence length, precision and data order fixed, you can trade micro-batch against accumulation against replica count freely, and the training curve per step stays the same up to floating-point noise. That gives a clean split into two classes:

The practical rule follows: tune systems knobs with a stopwatch, not with a hyperparameter search. Each trial of a search should be a training run that differs in the math; spending search budget on micro-batch size wastes it.

Step one: benchmark the systems knobs

Step one, before any search, is to find the largest micro-batch that fits and the throughput it delivers. Activation memory grows roughly linearly with micro-batch and sequence length, while weights, gradients and optimizer state do not depend on batch at all, so the first OOM tells you most of what you need. The harness below runs a few real optimizer steps per candidate, records peak memory and tokens per second, and treats an out-of-memory error as a result, not a crash.

import time, torch

def bench(make_model, make_batch, micro_batches, seq_len, steps=8, warmup=3):
    results = []
    for mb in micro_batches:
        model = make_model().cuda()
        opt = torch.optim.AdamW(model.parameters(), lr=1e-4, fused=True)
        torch.cuda.reset_peak_memory_stats()
        try:
            for i in range(warmup + steps):
                if i == warmup:
                    torch.cuda.synchronize(); t0 = time.perf_counter()
                x, y = make_batch(mb, seq_len)
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    loss = torch.nn.functional.cross_entropy(
                        model(x).flatten(0, 1), y.flatten())
                loss.backward()
                opt.step(); opt.zero_grad(set_to_none=True)
            torch.cuda.synchronize()
            dt = (time.perf_counter() - t0) / steps
            results.append((mb, mb * seq_len / dt,
                            torch.cuda.max_memory_allocated() / 2**30))
        except torch.cuda.OutOfMemoryError:
            results.append((mb, None, None))
            break
        finally:
            del model, opt
            torch.cuda.empty_cache()
    return results   # (micro_batch, tokens_per_sec, peak_GiB)

Read the output as a curve, not a maximum. Throughput usually rises steeply at small micro-batches, where the GPU is underfed, and then flattens. Choose the smallest micro-batch on the flat part and leave 10 to 15 percent memory headroom, because fragmentation and longer real sequences push peak memory above the benchmark. Keep hidden sizes and micro-batch times sequence length at multiples of 8 or larger powers of two so matrix multiplies map onto tensor-core tiles cleanly.

If the micro-batch you want does not fit, activation checkpointing recomputes activations in the backward pass instead of storing them. Full recomputation costs roughly one extra forward pass, so about a third more compute per step; selective checkpointing of only the most memory-hungry layers costs less. It never changes the loss, so it is a systems knob: enable it only when it lets you reach a better point on the throughput curve.

Global batch and the learning rate

Global batch is where systems and optimisation meet. Larger batches reduce gradient noise and let you take bigger steps, but only up to a point. McCandlish and colleagues (2018) described a critical batch size estimated from the gradient noise scale: below it, doubling the batch roughly halves the steps needed; above it, extra batch buys little and you are burning GPU hours. The noise scale typically grows during training, which is one motivation for batch-size warmup schedules.

When you do change the global batch, the learning rate must move with it. Two heuristics are widely used, and both are starting points for a short re-tune, not laws:

• Linear scaling for SGD. Goyal et al. (2017) scaled ImageNet training to large batches by multiplying the learning rate by the batch ratio, with a gradual warmup to avoid early instability.
• Square-root scaling for Adam-family optimisers. Analyses of adaptive optimisers, such as Malladi et al. (2022), motivate scaling the learning rate by the square root of the batch ratio, often adjusting other optimiser constants too.

Worked example. You tuned a model on 8 GPUs with micro-batch 16, accumulation 4 and sequence length 2,048: a global batch of 16 x 4 x 8 = 512 sequences, or 1,048,576 tokens per step, with AdamW at a peak learning rate of 3e-4. You now get 32 GPUs.

The first plan is almost always the right first move: it reuses every tuned value and turns new hardware into wall-clock speed. Only take the second if you have evidence you are below the critical batch size, for example because the first plan's GPUs are underutilised at accumulation 1. Note that the token budget is unchanged in both: the second plan takes a quarter as many optimizer steps, so step-indexed values such as warmup length and schedule decay points must be rescaled too.

When systems knobs secretly change the math

The promise that systems knobs do not change the math has three common exceptions. Check each one explicitly.

Loss normalisation under accumulation. Averaging each micro-batch's mean loss and dividing by the number of accumulation steps is correct only when every micro-batch has the same number of loss tokens. With padding or packed variable-length sequences it over-weights tokens in short micro-batches, so changing the accumulation split changes the effective objective. Normalise by the total number of loss tokens across the whole accumulation window:

n_tokens = sum((y != PAD).sum() for _, y in window)        # whole window
for x, y in window:
    with torch.autocast("cuda", dtype=torch.bfloat16):
        logits = model(x)
    loss_sum = torch.nn.functional.cross_entropy(
        logits.flatten(0, 1).float(), y.flatten(),
        ignore_index=PAD, reduction="sum")
    (loss_sum / n_tokens).backward()                       # not / accum_steps
# in data parallel, all-reduce n_tokens across ranks first
opt.step(); opt.zero_grad(set_to_none=True)

Batch normalisation. Batch-norm statistics are computed over the per-GPU micro-batch, not the global batch. Changing micro-batch changes the model's behaviour even at fixed global batch, unless you use synchronised batch norm.

Precision. BF16 has the same exponent range as FP32 and normally needs no loss scaling. FP16 has a narrow range, so small gradients underflow without dynamic loss scaling, which PyTorch provides as torch.amp.GradScaler. FP8 training changes numerics further and depends on per-tensor or finer-grained scaling recipes. Treat a precision change as a coupled change: re-validate the learning rate and keep optimizer state and master weights in FP32 unless your framework's recipe says otherwise.

Spending the search budget on GPUs

With systems knobs fixed, spend the search budget on the optimisation knobs, in roughly this order of sensitivity: peak learning rate first, then warmup and schedule, then weight decay, then Adam betas and epsilon. Three GPU-specific habits make the budget go further.

Fix the budget in tokens or GPU-hours, not steps. Compare trials at equal cost, otherwise the search rewards whichever configuration happens to run longest. Prune early with successive halving so most trials die after a small fraction of the budget. Pack small trials: a model that uses a fifth of a GPU's memory leaves the rest idle, so run several trials per device, or partition supported data-centre GPUs with MIG so each trial gets isolated memory and compute. The objective below wires these together with Optuna.

import optuna

TOKEN_BUDGET = 2_000_000_000          # identical for every trial
EVAL_EVERY = 100_000_000

def objective(trial):
    cfg = dict(
        lr=trial.suggest_float("lr", 1e-5, 3e-3, log=True),
        warmup_frac=trial.suggest_float("warmup_frac", 0.005, 0.05, log=True),
        weight_decay=trial.suggest_float("weight_decay", 1e-3, 0.3, log=True),
        beta2=trial.suggest_categorical("beta2", [0.95, 0.98, 0.999]),
        # systems knobs come from the benchmark, never from the search
        micro_batch=BENCH_MICRO_BATCH, grad_accum=BENCH_ACCUM,
    )
    run = start_training(cfg, seed=trial.number)
    seen = 0
    while seen < TOKEN_BUDGET:
        seen = run.train_tokens(EVAL_EVERY)
        val = run.eval_loss()
        if val != val:                                # NaN: diverged
            raise optuna.TrialPruned()
        trial.report(val, step=seen // EVAL_EVERY)
        if trial.should_prune():
            raise optuna.TrialPruned()
    return val

study = optuna.create_study(direction="minimize",
                            pruner=optuna.pruners.HyperbandPruner(),
                            storage="sqlite:///hpo.db", study_name="lm-v3",
                            load_if_exists=True)
study.optimize(objective, n_trials=64)

Finally, re-run the top three configurations with two or three different seeds before choosing. Run-to-run variance from data order and non-deterministic GPU kernels is often as large as the gap between the best few trials, and a winner that does not survive a reseed was noise. For transferring results from small proxy models to a large target, standard parametrisation does not preserve the optimal learning rate across width; maximal update parametrisation (muP) was designed to, and is covered on its own page.

Failure modes

• Tuning at one global batch, training at another. The tuned learning rate is valid only for the batch it was tuned at. Log global batch with every trial.
• Treating OOM as a failed trial. In a search, an out-of-memory crash looks like a bad configuration and biases the sampler. Fix the micro-batch from the benchmark so memory never depends on sampled values.
• Step-indexed schedules after a batch change. Warmup of 2,000 steps means four times as many tokens after a 4x batch increase. Express schedules as fractions of the token budget.
• Divergence mistaken for a bad learning rate. Loss spikes in FP16 can be scaler underflow or overflow; check the scaler's scale value before shrinking the learning rate.
• Throughput benchmarks on synthetic data. Random tokens with fixed length hide data-loader stalls and padding waste. Benchmark with the real loader at least once.
• Shared-GPU interference. Packed trials without isolation contend for memory bandwidth, so their throughput numbers are not comparable to dedicated runs.

Trade-offs

What to do next

1. Write down your current global batch, tokens per step and precision, and log them with every run from now on.
2. Run the micro-batch benchmark on your real data loader and pick the smallest micro-batch on the flat part of the throughput curve, with memory headroom.
3. Fix the accumulation-window loss normalisation to divide by total loss tokens, and verify two different accumulation splits give matching losses for a few hundred steps.
4. When hardware changes, keep global batch fixed first; only grow it with evidence and a learning-rate re-tune.
5. Run the optimisation search at a fixed token budget with Hyperband pruning and persistent storage.
6. Reseed the top three configurations before committing to a winner.

Keep learning: Grid, random, Bayesian, ASHA and PBT in practice, Hyperparameter optimisation architecture, Batch size math, muP and learning-rate transfer and Mixed precision training.