LLM Developer DocsLLM Docs

Loading the search index ...

↑ ↓ move, Enter open, Esc close

GitHub

How the Model Learns

One full learning step from text to updated parameters, with each stage linked to its page and to its call site in train.main.

transformer.py

Training repeats one step many times. This page follows that step from the text to the changed parameters. It adds no new mathematics: each stage links to the page that explains it.

One learning step

Every stage below names the line of train.main it runs from, and the path from there to its own line. The numbers come from one real step on a hand-picked text with a tiny model, so they are labelled Illustrative Example and say nothing about the trained model. The batch comes from one stored shuffle.

One learning step, in the order of train.main. Select a stage.

1. Input text

load_corpus reads one column of the CSV file and joins its rows into a single string, once, before the loop starts.

Page: Data Pipeline

Where it runs, starting from train.main:

  1. text = load_corpus(config.corpus_path, config.corpus_column, config.max_rows) train.py:40

Illustrative Example The text: "the cat sat on the mat\nthe dog sat on the lo..."

data.pyload_corpusGitHub
def load_corpus(path, column, max_rows) -> str:    df = pd.read_csv(path, nrows=max_rows)     if column not in df.columns:        raise ValueError(f"no column {column!r}, the file has {list(df.columns)}")     texts = df[column].dropna().astype(str).str.strip()    return "\n".join(texts) + "\n"

Stages 1 to 3 run once, before the loop. Stages 4 to 12 run for every batch. Several of them share one line of train.main, because batch_loss_and_grads does all of stages 4 to 8 inside one call.

The order in train.main

These are the lines of the loop, as they are in the file. The comment above the loop states the same order.

train.pymainGitHub
    # 3. the loop: gradient -> clip -> learning rate -> Adam    step = 0    for epoch in range(config.epochs):        for xb, yb in iterate_batches(tr_x, tr_y, config.batch_size, rng):            loss, grads = batch_loss_and_grads(params, xb, yb, cfg)            norm = clip_grad_norm(grads, config.grad_clip)            step += 1            lr = lr_schedule(step, config.warmup_steps, total_steps,                             config.lr_max, config.lr_min)            opt.step(params, grads, lr)            if step == 1 or step % 50 == 0:                print(f"step {step:5d}  loss {loss:.3f}  |g| {norm:.2f}  lr {lr:.5f}")

The order is the gradient, then clipping, then the learning rate, then the optimizer step:

  1. batch_loss_and_grads computes the mean loss and the mean gradient of the batch (train.py:61).
  2. clip_grad_norm may scale those gradients in place (train.py:62).
  3. step goes up by one, and lr_schedule gives the learning rate for that step (train.py:63-65).
  4. opt.step changes the parameters (train.py:66).

AdamW / Learning Rate / Clipping covers steps 2 to 4, and Training Loop covers the batches, validation and printed lines.

One weight over two steps

This follows a single weight through the part of the loop after the gradient. The table is computed by this site's copy of clip_grad_norm, lr_schedule and Adam.step, which tests compare with the real functions.

One weight over two stepsIllustrative ExampleThe weight and its gradients are hand-picked, and total_steps is 200, because the real value depends on the corpus. The settings are the values in config.py. Both steps fall inside warmup, so the weight moves very little. The first gradient is above grad_clip, and the second is not.
ValueCodeStep 1Step 2
Weight beforeparams[k]0.50.4999799
Mean gradientgrads[k]30.5
Length before clippingnorm30.5
Gradient after clippinggrads[k]10.5
Stepstep += 112
Learning ratelr0.000020.00004
Moving averageself.m[k]0.10.14
Moving average of the squareself.v[k]0.0010.001249
Bias-correctedmhat10.73684211
Bias-correctedvhat10.62481241
Update before decaymhat / (np.sqrt(vhat) + self.eps)0.999999990.93217963
Update with decayupdate + self.wd * params[k]1.004999990.93717943
Weight afterparams[k] -= lr * update0.49997990.49994241

Two things to read from the table:

  • Clipping changes the first step's gradient but not the second's. The second gradient is smaller than grad_clip, so it passes through as it is.
  • The move depends on more than the current gradient. It is lr times update. lr doubles from step 1 to step 2, and update comes from the history in self.m[k] and self.v[k].

What each number means

train.py prints one line every 50 steps, and also at step 1, with loss, |g| and lr. It prints one line per epoch with the validation loss and ppl.

The two lines these numbers appear onIllustrative ExamplePrinted by evaluating the real format strings of train.main with the values of one real step on a hand-picked text and a tiny model. train.main itself did not run, and these are not numbers from a training run.
step     1  loss 2.930  |g| 1.94  lr 0.00002
epoch 1: val loss 3.007  ppl 20.2
Printed asIn train.mainMeaning
losslossThe mean cross-entropy of this batch, computed before this step's update. It is one batch, so it moves up and down from step to step. A lower value means the model gave the target tokens higher probabilities.
pplperplexity(v)exp of the validation loss, printed once per epoch and read by nothing else. In general, it can be read as the number of equally likely tokens that would give the same loss. A model that gives every token the same probability has ppl equal to V.
|g|normThe length of the mean gradient, measured over all parameters at once, before clipping. So it can be above grad_clip. When it is, the update used the clipped gradient, not the printed length.
lrlrThe learning rate used for this step, from lr_schedule. It is printed with five decimals, so very small values lose digits.

Training Loop lists every line train.py prints.