Step 9 — Run a real training job

Step 9 — Run a real training job 🚂¶

Goal: train the biggest Sanskrit model your data and budget actually justify — and survive the run.

Why this step matters¶

Everything before this was on toy scale. Now you meet the real problems: memory limits, speed, instability, and cost.

Based on Step 6, that is probably a model somewhere between 100 and 500 million parameters. Be honest about the size your data supports. A larger model trained on too little data is worse than a smaller one trained properly, and it costs more.

The four lines at the centre of this chapter are the same four lines you wrote in Step 1. Everything else in this chapter exists to stop those four lines from silently wasting a week of GPU time.

The unbreakable rule: smoke first 🚦¶

Before anything else in this chapter, internalise one habit.

The shape of this chapter 🗺️¶

Six additions. The first three buy you memory and speed. The last three buy you the ability to recover when a run goes wrong — which it will.

1 · bf16, never fp16 🔢¶

Storing most numbers in 16 bits instead of 32 roughly halves memory and speeds everything up. But which 16-bit format matters enormously:

def pick_dtype(device):
    if device == "cuda" and torch.cuda.is_bf16_supported():
        return torch.bfloat16
    return torch.float32

Both are 16 bits. The difference is what they spend the bits on:

Language model training cares far more about range than precision, because gradients span many orders of magnitude. fp16 overflows, produces inf, and needs a loss scaler babysitting it. bf16 just works.

If your card predates Ampere it has no bf16 and you fall back to fp32. That is fine — slower, not wrong.

And FlashAttention, which you already turned on in Step 8 by using F.scaled_dot_product_attention, computes exactly the same attention without ever materialising the T × T matrix. Faster, far less memory, identical results. There is no downside.

2 · Gradient accumulation — a big batch on a small card ➕¶

Run several small batches, add up their gradients, and only update the weights once at the end:

opt.zero_grad(set_to_none=True)
for micro in range(cfg["accum"]):
    xb, yb = get_batch(splits["train"], model.cfg, device)
    with torch.autocast(device_type=device, dtype=dtype, enabled=use_amp):
        _, loss = model(xb, yb)
    (loss / cfg["accum"]).backward()      # ← the division is not optional

Four micro-batches of 8 behaves almost exactly like one batch of 32, but only ever holds 8 sequences of activations in memory.

3 · Gradient checkpointing — trade time for memory 💾¶

Normally the model saves every intermediate value during the forward pass so it can use them in the backward pass. Gradient checkpointing throws most of them away and recomputes them when needed.

for blk in model.blocks:
    blk.forward = _checkpointed(blk.forward)     # ~30% slower, much less memory

You trade about thirty percent extra time for a large memory saving. Worth it whenever memory — not speed — is your limit. Turn it on only when you actually hit an out-of-memory error; it is not free.

4 · The learning rate schedule 📈¶

Two schedules, one function:

def lr_at(step, cfg):
    warm, total = cfg["warmup"], cfg["steps"]
    lr, lr_min = cfg["lr"], cfg["lr"] * 0.1

    if step < warm:                               # linear warmup, both schedules
        return lr * (step + 1) / warm

    if cfg["schedule"] == "wsd":
        decay_from = int(total * 0.8)             # last 20% decays
        if step < decay_from:
            return lr                             # the long flat "stable" phase
        frac = (step - decay_from) / max(total - decay_from, 1)
        return lr_min + (lr - lr_min) * (1 - frac)

    frac = (step - warm) / max(total - warm, 1)   # cosine
    return lr_min + 0.5 * (lr - lr_min) * (1 + math.cos(math.pi * frac))

Cosine decay is the standard choice and the right default. Its one drawback is that the shape depends on the total step count, so you must commit to a run length before you start.

Warmup-Stable-Decay (WSD) warms up, holds steady for most of the run, then decays only at the end. Because the middle is flat, you can branch a new experiment off any point in it without redoing warmup — and you can decide to train longer halfway through. Very useful while you are still exploring.

5 · Watch the gradient norm — your earliest warning 🌡️¶

This is the most useful number on your dashboard, and most people never plot it.

gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["clip"])
opt.step()

clip_grad_norm_ returns the norm before clipping, which is exactly the number you want. It usually starts climbing a few hundred steps before the loss does anything visible on the chart. If the gradient norm spikes, you have time to react — lower the learning rate, or roll back — before the loss follows.

The file logs it on every evaluation, and also watches for the two failures you already caused on purpose in Step 2:

if not math.isfinite(losses["train"]):
    print("\n  NaN loss. Learning rate too high, or fp16 overflow.")
    print("  Roll back to best.pt, halve the LR, continue. This is routine.")
    break
if losses["val"] > best_val * 1.15 and step > cfg["warmup"] * 3:
    print(f"  ^ val loss {losses['val']/best_val:.0%} of best "
          f"— overfitting (Step 2, experiment 5)")

You have seen both before. You will recognise them.

6 · Checkpoint often, and test the resume 💾¶

if losses["val"] < best_val:
    best_val = losses["val"]
    save_ckpt(out_dir / "best.pt", model, opt, step, cfg, tok, best_val)
save_ckpt(out_dir / "latest.pt", model, opt, step, cfg, tok, best_val)

Two files, doing different jobs. best.pt is what you ship. latest.pt is what you resume from. Save the optimizer state as well as the weights — AdamW carries running moment estimates, and resuming without them restarts the optimizer cold and puts a visible kink in your loss curve.

When a run goes bad, you roll back to the last good checkpoint, lower the learning rate, and continue. Everyone does this. It is normal. It is not a sign of failure.

7 · Scale up in stages 🪜¶

1. One GPU. Get it working.

2. DDP — copy the model to each GPU, split the data. Simple.

3. FSDP — split the model itself across GPUs. Use when the model no longer fits on one.

4. Tensor and pipeline parallel — most readers never need these.

Do not skip a rung. Each one adds a new class of bug, and debugging two new classes at once is how people lose weeks.

What you should see ▶️¶

python train.py --smoke

device=cpu dtype=torch.float32 params=0.10M
schedule=cosine steps=60 micro_batch=8 accum=1 -> effective batch 8
   step        lr    train      val   |grad|     tok/s
--------------------------------------------------------
      0  6.00e-05   4.4031   4.4083    1.846     6,680
      6  3.00e-04   4.2436   4.2421    1.467    13,846
     12  2.89e-04   4.1221   4.1241    1.328    15,345
     24  2.28e-04   3.9727   3.9788    1.435    16,803
     36  1.38e-04   3.8455   3.8322    1.381    17,298
     48  6.05e-05   3.7541   3.7649    1.423    17,586
     59  3.02e-05   3.7307   3.7324    1.520    17,595

best val loss 3.7324  ->  run/best.pt

Read every column, because each one is telling you something:

• lr rises then falls. Warmup then cosine decay, exactly as configured. If it does not, your schedule is wrong and nothing downstream will make sense.

• train and val fall together. No gap has opened yet, so no overfitting yet. The moment they separate, see Step 2, experiment 5.

• |grad| sits flat around 1.4. That is a healthy run. A climb here is your first warning; a spike is your last one.

• tok/s climbs then flattens. It is a running average, so the early numbers include startup cost. The flat value is your real throughput, and it is what you multiply out to estimate the cost of the full run before launching it.

Switch schedules and the flat phase is obvious:

python train.py --smoke --schedule wsd --out run2

      6  3.00e-04   4.2300   4.2277    1.618    14,215
     18  3.00e-04   3.9680   3.9934    1.297    15,672
     30  3.00e-04   3.7613   3.7728    1.422    16,691
     42  3.00e-04   3.5298   3.5243    1.670    17,170
     54  1.65e-04   3.3531   3.3761    1.673    17,476
     59  5.25e-05   3.3341   3.3264    1.617    17,443

Constant 3.00e-04 for most of the run, then a short decay — and in this particular sixty-step run WSD lands at 3.33 where cosine landed at 3.73, because cosine spent most of its budget decaying rather than learning. Do not over-read that: on a run this short, the schedule that stays at a high learning rate longer is bound to look better, and the ranking can reverse at real length. It is a good illustration of the shape, not a recommendation.

Then confirm the resume works before you trust it:

$ python train.py --smoke --resume run/latest.pt
resumed from run/latest.pt at step 59

Where people usually get stuck¶

Launching one giant expensive run without a small test run first. Covered at the top, and still the single most expensive mistake in this chapter.

Forgetting to divide by the accumulation count. Silent, and it looks exactly like an unstable learning rate.

Resuming without the optimizer state. Your loss curve gets a visible kink and you spend an afternoon looking for a data bug.

Not logging the gradient norm. You lose your only early warning and find out about instability from the loss, hours later.

Reading tok/s from the first few steps. It is a running average with startup cost baked in. Use the flattened value for cost estimates.

Treating a rollback as a failure. Roll back, halve the learning rate, continue. That is the job.

You are ready to move on when¶

You have a finished training run, a saved checkpoint, a log.jsonl you can plot, and a loss curve you can explain to someone else — including any spikes, and what you did about them.

You should also be able to say, from tok/s and your token count, roughly what the run cost.