Running a tuning job
Lesson 4 of 6
4Running a tuning job
The big pictureA first tuning run usually ends with "out of memory". This lesson is about the road from there to a loss curve you can read: where memory goes, which setting to change first, and how to tell from the curve that training is going well or needs stopping.
What the picture shows: Training loss keeps falling, but validation rises after the dashed line: memorisation started here. The best model is at that line, and load_best_model_at_end keeps it automatically even if training carries on.
Map of this part
The concepts in order, click any one to jump to it:
Two ways into the same idea
Running a tuning job is like cooking a recipe for the first time in an unfamiliar oven: it may burn at first because the heat is too high, or stay raw because the time is short. You watch the dish (the loss curve) and change one setting at a time so you know which one caused what.
Out-of-memory errors are handled with smaller batches plus gradient accumulation, shorter sequences and gradient checkpointing. You start with a small learning rate and a warmup, log training and validation loss, and save checkpoints to return to the best one.
Before You Start
🪡LoRA
You know LoRA freezes the weights and trains small adapters, so the frozen weights have no gradients.
🩺Debugging training
You have read loss curves before and know what overfitting looks like.
The Core Idea
A first tuning run usually ends out of memory. In this lesson you see where memory goes, how to shrink the batch without losing anything, which settings to change first, and how to read the curve.
1Where memory goes
Memory in training goes to four things: the weights, their gradients, the optimiser state (Adam keeps two numbers per parameter), and the activations saved for the backward pass.
Look at the picture: for a 7B model, full tuning needs 84 GB before activations, LoRA about 14, and QLoRA about 4.
def gb(n_bytes):
return n_bytes / 1e9
P = 7e9 # a 7B model
lora = 40e6 # trainable LoRA parameters, about 0.6%
full = P * 2 + P * 2 + P * 8 # bf16 weights + bf16 grads + Adam's two fp32 states
lora_bf16 = P * 2 + lora * (2 + 8)
qlora = P * 0.5 + lora * (2 + 8) # base weights in 4 bits
print(f"full fine-tuning: {gb(full):6.1f} GB (+ activations)")
print(f"LoRA on bf16 base: {gb(lora_bf16):6.1f} GB (+ activations)")
print(f"QLoRA on 4-bit base:{gb(qlora):6.1f} GB (+ activations)")
Full tuning of a 7B model needs 84 gigabytes before activations, more than the largest common card. LoRA removes gradients and optimiser state for all frozen weights, bringing it to 14 gigabytes. QLoRA compresses the base model itself to 4 bits, making it about 4 gigabytes. That is what made tuning possible on a single ordinary card.
2The batch and gradient accumulation
Activations scale with the batch size, so it is the first thing you lower. A batch of 1 or 2 per card is normal when tuning large models.
Look at the picture: four small batches of 8 add up their gradients, then one optimiser step. That is an effective batch of 32 with the memory of a batch of 8.
The gradient of one batch of 32 examples, against summing the gradients of four batches of 8.
import numpy as np
rng = np.random.default_rng(0)
X = rng.normal(size=(32, 3))
y = X @ np.array([1.0, -2.0, 0.5]) + rng.normal(scale=0.1, size=32)
w = np.zeros(3)
def grad(Xb, yb, w):
return 2 * Xb.T @ (Xb @ w - yb) / len(yb) # mean squared error gradient
big = grad(X, y, w) # one batch of 32
accum = np.zeros(3)
for i in range(4): # four micro-batches of 8
accum += grad(X[i * 8:(i + 1) * 8], y[i * 8:(i + 1) * 8], w) / 4 # divide by the number of steps
print("batch of 32: ", big.round(4))
print("4 x 8, accumulated: ", accum.round(4))
print("identical:", bool(np.allclose(big, accum)))
Exactly identical, provided you divide by the number of steps (/ 4), since each small gradient is a mean over 8, not a sum. Forgetting this division makes the gradient four times larger, as if you quadrupled the learning rate without knowing. Libraries handle it for you, which is why you see the loss divided by gradient_accumulation_steps in hand-written training code.
3The settings to change first
Three settings first: the learning rate (1e-4 to 3e-4 for LoRA), the number of epochs (1 to 3), and the effective batch (16 to 64 through accumulation). Leave the rest at defaults until these work.
Then bf16, which halves weight and activation memory without fp16's overflow issues. And if a batch of 1 is still too much, turn on gradient checkpointing.
from trl import SFTConfig, SFTTrainer
config = SFTConfig(
output_dir="runs/lora-r16-lr2e-4",
per_device_train_batch_size=2,
gradient_accumulation_steps=16, # effective batch = 2 x 16 = 32
learning_rate=2e-4,
num_train_epochs=2,
bf16=True,
gradient_checkpointing=True,
max_length=1024, # chosen after printing the length distribution
eval_strategy="steps", eval_steps=50,
save_strategy="steps", save_steps=50,
load_best_model_at_end=True, # keep the checkpoint with the lowest validation loss
logging_steps=10,
)
trainer = SFTTrainer(model=model, args=config, train_dataset=train, eval_dataset=val)
trainer.train()
The output folder's name carries the settings you vary, so the name alone tells you how two runs differ. And load_best_model_at_end with eval_steps automatically keeps the checkpoint with the lowest validation loss, so it does no harm if training continues past the start of memorisation.
4Reading the curve
Training loss falls smoothly and validation loss follows it. If validation rises while training falls, that is memorisation: stop and take the best checkpoint.
Look at the picture: every 50 steps you look at both losses together, and at a sample of the model's outputs with your own eyes.
How It Works, Step by Step: how to get from out-of-memory to a working run
Print the distribution of example lengths in tokens, and choose the maximum length so no answer gets cut.
Lower the per-card batch and raise gradient accumulation, so the effective batch stays the same.
Turn on bf16, then gradient checkpointing if a batch of 1 is still too much.
Watch training and validation loss together with a sample of outputs, and take the checkpoint at the lowest validation loss.
Going deeper
Another technique saves a lot of memory: gradient checkpointing. Instead of keeping every activation for the backward pass, it keeps some and recomputes the rest when needed. Training slows by about 30%, and memory falls a lot. Turn it on when a batch of 1 is still too much.
Worked example
Memory in training goes to four things: the weights, their gradients, the optimiser state (Adam keeps two numbers per parameter), and the activations saved for the backward pass. The first three are a multiplication; the fourth grows with batch size and sequence length.
Try it without code
Myth vs Reality
| Myth | Reality |
|---|---|
| Training loss near zero means success. | In the third epoch it usually means memorisation. Look at validation loss and at a sample of outputs with your own eyes. |
| Gradient accumulation approximates a large batch. | The result is mathematically identical, provided you divide by the number of steps. Forgetting the division scales the gradient as if you raised the learning rate. |
| Cutting the maximum length is a safe fix for out-of-memory. | If examples are longer than it, their answers get truncated and the model learns cut-off answers. |
Try it yourself
You run out of memory on the first step. What do you try first?
Activations scale with the batch, and accumulation keeps the same maths. After that, try bf16 and gradient checkpointing.
Training loss is falling and validation loss has been rising since step 400. What do you do?
The model started memorising training after step 400. Your best model is at that point, and training past it does harm.
Quick review cards
bf16 halves weight and activation memory and speeds up compute, without fp16's overflow issues. Use it if your card supports it.Check Your Understanding
Your batch is 2 and accumulation 16, and you raise the batch to 8 now that memory allows. What do you change in accumulation?
Lower it to 4, so the effective batch stays 32. If you keep 16 the effective batch becomes 128, training behaves differently and needs a different learning rate, and you can no longer compare the two runs fairly. A larger per-card batch should only speed training up, not change its result.
Why does a run folder name such as runs/lora-r16-lr2e-4 carry the settings you vary?
So the name alone tells you how two runs differ, letting you compare curves without opening the settings, and not mixing two experiments up weeks later.
Resources for this lesson
Grouped by type. Start with the official docs for precision, or an article or video for the explanation.
Done with this lesson?
Sign in to save your progress.