Build one training step by hand

Problem

When you need this. A training step should make the expected next token more likely. Calculate that change in a small model so you can see how a prediction error reaches each weight.

A language-model framework hides one training step behind three calls: compute the loss, call backward, step the optimizer. When training goes wrong, you need to know what those calls compute. In this build you write that step yourself, small enough to check every number.

The model predicts the next token. Its vocabulary has 4 tokens: “red”, “green”, “stop”, “go”. The training batch holds 3 examples: “red” is followed by “stop”, “green” by “go”, and “red” by “stop” again. The model starts with small random weights, so it gives every next token roughly the same probability.

The program has a fixed contract. It takes a batch of input token IDs, a batch of target token IDs, the current parameters and a learning rate. It returns the mean loss, a gradient for every parameter and a new set of parameters. It rejects an empty batch, a token ID outside the vocabulary, inputs and targets of different lengths and a learning rate that is not a positive finite number.

Success has three checks, all reproducible with the supplied tests. Every hand-written gradient agrees with a finite-difference estimate to within 10⁻⁶. One update with learning rate 0.5 lowers the loss on the batch from 1.2952 to 1.1345. Every bad input raises a ValueError that names the problem. A lower loss on three training examples does not show that the model has learned anything useful about new text; later builds in this series measure that on held-out data.

Concepts

The model must turn a token into numbers before it can calculate a prediction. A token ID selects a row from the embedding table: a short vector of learned numbers representing that token. A parameter is a number training can change. Here the parameters are the embedding entries, the prediction weights and the biases.

A tensor holds numbers in an array; its shape gives the size of each dimension. This implementation stores tensors as rectangular Python lists. A batch groups examples processed with the same parameters. Each input has a target, the token it should predict. The input IDs have shape (3): one ID per example. The model’s 20 parameters sit in three tensors:

The forward pass computes a prediction from the current parameters. It looks up each input embedding, multiplies it by weights and adds bias. The resulting logits are raw scores, one per possible next token. They are not probabilities: they can be negative and need not sum to one. Softmax exponentiates and normalizes those scores into positive probabilities that sum to one.

A loss turns prediction errors into a number training can minimize. Cross-entropy here is −ln of the probability assigned to the target token. Assigning the target a low probability gives a larger loss. Averaging the examples gives one batch loss to differentiate.

A gradient contains the partial derivative of the loss with respect to each parameter: how the loss changes for a small change in that parameter, with the others fixed. The backward pass computes these derivatives for all 20 parameters. It uses two facts. First, for softmax followed by mean cross-entropy, the gradient with respect to the logits is the probabilities minus a one-hot row for the target (1 at the target’s position, 0 elsewhere), divided by the batch size. Second, the chain rule carries that gradient backward: through the matrix multiplication to weights, bias and the embedding rows, and through the lookup into the embedding table.

A gradient check estimates each derivative independently with a finite difference: raise it by a small ε, lower it by ε and divide the loss difference by 2ε. If a hand-written gradient disagrees with that estimate, the backward pass has a bug. The update is plain stochastic gradient descent (SGD) on one batch: each parameter moves against its gradient by the learning rate times the gradient. The learning rate sets the size of the update; it must be small enough for the local derivative to remain useful.

Architecture

The figure follows one step through the program. Input IDs enter the Forward pass, which looks up embedding rows, multiplies them by the weights and applies softmax. Loss compares those probabilities with the Targets. The Backward pass turns the loss into Gradients, one per parameter. The dashed branch sends the gradients to the finite-difference Check, which changes nothing. SGD update changes the weights, and the return path labeled new weights carries them to the next forward pass. The vocabulary, the batch and the shapes stay fixed.

tiny_step.py is one Python process with no files, network calls or saved state. forward() returns every intermediate value. backward() uses those values to return the gradients. numerical_gradients() reruns forward() twice for each of the 20 parameters. sgd_step() returns new parameters and leaves the old ones unchanged, so the trace can show values before and after.

Working code

Runtime: Python 3.10 or newer. Dependencies: Python standard library only.

Download and extract the complete lab using Download the complete lab above. Open build-one-training-step-by-hand/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.

"""One training step for a tiny next-token model, written with Python lists.

The model maps each input token ID to an embedding row, multiplies that row by
an output weight matrix, adds a bias and turns the scores into probabilities.
Training compares the probabilities with the target token, computes gradients by
hand, checks them against finite differences and applies one SGD update.

Run:  python3 tiny_step.py           print the recorded trace as JSON
      python3 tiny_step.py --write   rewrite step-trace.json
      python3 tiny_step.py --check   exit 1 if step-trace.json is out of date
Python 3.10+, standard library only.
"""
import json
import math
import random
import sys
from pathlib import Path

VOCAB = ["red", "green", "stop", "go"]
INPUTS = ["red", "green", "red"]
TARGETS = ["stop", "go", "stop"]
EMBED_DIM = 2
SEED = 7
LEARNING_RATE = 0.5
EPSILON = 1e-5


# ---- tensors as nested lists -------------------------------------------------

def shape(t):
    """Return the shape of a rectangular nested list, e.g. (3, 2)."""
    dims = []
    while isinstance(t, list):
        dims.append(len(t))
        t = t[0] if t else None
    return tuple(dims)


def matmul(a, b):
    """(n, k) @ (k, m) -> (n, m). Rejects mismatched inner sizes."""
    (n, k), (k2, m) = shape(a), shape(b)
    if k != k2:
        raise ValueError(f"cannot multiply {shape(a)} by {shape(b)}: inner sizes {k} and {k2} differ")
    return [[math.fsum(a[i][p] * b[p][j] for p in range(k)) for j in range(m)] for i in range(n)]


def transpose(a):
    return [list(row) for row in zip(*a)]


# ---- model -------------------------------------------------------------------

def init_params(vocab_size=len(VOCAB), dim=EMBED_DIM, seed=SEED):
    """Small random starting values, rounded to two decimals so a reader can follow them."""
    rng = random.Random(seed)
    value = lambda: round(rng.uniform(-0.5, 0.5), 2)
    return {
        "embedding": [[value() for _ in range(dim)] for _ in range(vocab_size)],  # (V, d)
        "weights": [[value() for _ in range(vocab_size)] for _ in range(dim)],    # (d, V)
        "bias": [0.0] * vocab_size,                                               # (V,)
    }


def encode(tokens, vocab=VOCAB):
    index = {token: i for i, token in enumerate(vocab)}
    unknown = [t for t in tokens if t not in index]
    if unknown:
        raise ValueError(f"tokens not in the vocabulary: {unknown}")
    return [index[t] for t in tokens]


def check_batch(ids, targets, vocab_size):
    if not ids:
        raise ValueError("empty batch: a training step needs at least one example")
    if len(ids) != len(targets):
        raise ValueError(f"batch has {len(ids)} inputs but {len(targets)} targets")
    for i in list(ids) + list(targets):
        if isinstance(i, bool) or not isinstance(i, int) or not 0 <= i < vocab_size:
            raise ValueError(f"token ID {i} is outside the vocabulary of size {vocab_size}")


def softmax(row):
    top = max(row)  # subtracting the largest score keeps exp() from overflowing
    exps = [math.exp(v - top) for v in row]
    total = math.fsum(exps)
    return [e / total for e in exps]


def forward(params, ids, targets):
    """Return every intermediate value of the forward pass."""
    vocab_size = len(params["embedding"])
    check_batch(ids, targets, vocab_size)
    x = [list(params["embedding"][i]) for i in ids]                   # (B, d)
    scores = matmul(x, params["weights"])                             # (B, V)
    logits = [[s + b for s, b in zip(row, params["bias"])] for row in scores]
    probs = [softmax(row) for row in logits]                          # (B, V)
    losses = [-math.log(p[t]) for p, t in zip(probs, targets)]        # (B,)
    return {"x": x, "logits": logits, "probs": probs, "losses": losses,
            "loss": math.fsum(losses) / len(losses)}


def backward(params, ids, targets, cache, accumulate=True):
    """Gradients of the mean loss with respect to every parameter.

    accumulate=False reproduces a common bug: writing each example's embedding
    gradient into its row instead of adding it, so a repeated token ID keeps
    only its last example's contribution.
    """
    batch = len(ids)
    vocab_size = len(params["embedding"])
    d_logits = [[(p - (1.0 if j == t else 0.0)) / batch for j, p in enumerate(row)]
                for row, t in zip(cache["probs"], targets)]             # (B, V)
    d_weights = matmul(transpose(cache["x"]), d_logits)               # (d, V)
    d_bias = [math.fsum(col) for col in zip(*d_logits)]                     # (V,)
    d_x = matmul(d_logits, transpose(params["weights"]))              # (B, d)
    dim = len(d_x[0])
    d_embedding = [[0.0] * dim for _ in range(vocab_size)]            # (V, d)
    for i, row in zip(ids, d_x):
        if accumulate:
            d_embedding[i] = [a + g for a, g in zip(d_embedding[i], row)]
        else:
            d_embedding[i] = list(row)
    return {"embedding": d_embedding, "weights": d_weights, "bias": d_bias}


def loss_only(params, ids, targets):
    return forward(params, ids, targets)["loss"]


def numerical_gradients(params, ids, targets, eps=EPSILON):
    """Central differences: nudge one value up and down, compare the losses."""
    grads = {}
    for name, value in params.items():
        if isinstance(value[0], list):
            grads[name] = [[0.0] * len(value[0]) for _ in value]
            cells = [(r, c) for r in range(len(value)) for c in range(len(value[0]))]
        else:
            grads[name] = [0.0] * len(value)
            cells = [(c,) for c in range(len(value))]
        for cell in cells:
            original = _get(value, cell)
            _set(value, cell, original + eps)
            up = loss_only(params, ids, targets)
            _set(value, cell, original - eps)
            down = loss_only(params, ids, targets)
            _set(value, cell, original)
            _set(grads[name], cell, (up - down) / (2 * eps))
    return grads


def _get(t, cell):
    return t[cell[0]][cell[1]] if len(cell) == 2 else t[cell[0]]


def _set(t, cell, v):
    if len(cell) == 2:
        t[cell[0]][cell[1]] = v
    else:
        t[cell[0]] = v


def flat(t):
    return [v for row in t for v in row] if isinstance(t[0], list) else list(t)


def max_gradient_error(analytic, numeric):
    return max(abs(a - n) for name in analytic for a, n in zip(flat(analytic[name]), flat(numeric[name])))


def sgd_step(params, grads, learning_rate):
    if not (isinstance(learning_rate, (int, float)) and math.isfinite(learning_rate) and learning_rate > 0):
        raise ValueError(f"learning rate must be a positive finite number, got {learning_rate!r}")
    new = {}
    for name, value in params.items():
        g = grads[name]
        if isinstance(value[0], list):
            new[name] = [[v - learning_rate * gv for v, gv in zip(row, grow)] for row, grow in zip(value, g)]
        else:
            new[name] = [v - learning_rate * gv for v, gv in zip(value, g)]
    return new


def train_step(params, ids, targets, learning_rate=LEARNING_RATE):
    cache = forward(params, ids, targets)
    grads = backward(params, ids, targets, cache)
    updated = sgd_step(params, grads, learning_rate)
    return cache["loss"], grads, updated, loss_only(updated, ids, targets)


# ---- recorded trace ------------------------------------------------------------

def r6(t):
    if isinstance(t, dict):
        return {k: r6(v) for k, v in t.items()}
    if isinstance(t, list):
        return [r6(v) for v in t]
    return round(t, 6)


def trace():
    params = init_params()
    ids, targets = encode(INPUTS), encode(TARGETS)
    cache = forward(params, ids, targets)
    grads = backward(params, ids, targets, cache)
    numeric = numerical_gradients(params, ids, targets)
    overwrite = backward(params, ids, targets, cache, accumulate=False)
    loss_before, _, updated, loss_after = train_step(params, ids, targets)
    return {
        "vocab": VOCAB,
        "inputs": INPUTS,
        "targets": TARGETS,
        "input_ids": ids,
        "target_ids": targets,
        "seed": SEED,
        "learning_rate": LEARNING_RATE,
        "shapes": {
            "input_ids": list(shape(ids)),
            "embedding": list(shape(params["embedding"])),
            "x": list(shape(cache["x"])),
            "weights": list(shape(params["weights"])),
            "logits": list(shape(cache["logits"])),
            "bias": list(shape(params["bias"])),
        },
        "parameter_count": sum(len(flat(v)) for v in params.values()),
        "params_before": r6(params),
        "x": r6(cache["x"]),
        "logits": r6(cache["logits"]),
        "probs": r6(cache["probs"]),
        "losses": r6(cache["losses"]),
        "loss_before": round(loss_before, 6),
        "grads": r6(grads),
        "gradient_check": {
            "epsilon": EPSILON,
            # The exact error depends on floating-point rounding; its power of ten is stable.
            "max_error_power_of_ten": math.floor(math.log10(max_gradient_error(grads, numeric))),
            "max_error_overwrite_bug": round(max_gradient_error(overwrite, numeric), 6),
            "loss_after_overwrite_bug": round(loss_only(sgd_step(params, overwrite, LEARNING_RATE), ids, targets), 6),
        },
        "params_after": r6(updated),
        "loss_after": round(loss_after, 6),
        "learning_rate_sweep": [
            {"learning_rate": lr, "loss_after": round(train_step(params, ids, targets, lr)[3], 6)}
            for lr in (0.1, 0.5, 1, 2, 5, 10, 20)
        ],
        "large_learning_rate": {
            "learning_rate": 20,
            "loss_after": round(train_step(params, ids, targets, 20)[3], 6),
        },
    }


if __name__ == "__main__":
    path = Path(__file__).with_name("step-trace.json")
    text = json.dumps(trace(), indent=2) + "\n"
    if "--write" in sys.argv:
        path.write_text(text)
    elif "--check" in sys.argv:
        sys.exit(0 if path.read_text() == text else "step-trace.json is out of date")
    else:
        print(text, end="")

matmul() checks inner sizes before it multiplies, so a shape mistake fails at the line that made it. softmax() subtracts the largest score before math.exp(); the probabilities are the same, and large scores cannot overflow. init_params() draws starting values from a seeded generator and rounds them to two decimals, so you can follow the arithmetic by hand.

Run the program to print the full trace as JSON:

python3 tiny_step.py

The recorded output is step-trace.json. python3 tiny_step.py --check exits with an error if your run differs from it.

Trace the forward pass

The inputs “red”, “green”, “red” become IDs 0, 1, 0. The lookup copies embedding rows 0, 1, 0, so x holds the “red” row twice: [-0.1800, -0.3500], [0.1500, -0.4300], [-0.1800, -0.3500].

Multiplying the first row by weights and adding the zero bias gives the logits [0.1108, -0.1029, 0.2104, 0.1718], one per vocabulary token. Softmax turns them into [0.2516, 0.2032, 0.2779, 0.2674]. The target is “stop”, ID 2, with probability 0.2779, so the loss for this example is −ln 0.2779 = 1.2805.

The three per-example losses are 1.2805, 1.3247, 1.2805, and their mean is 1.2952. That is a little below ln 4 ≈ 1.3863, the loss of a model that gives all four tokens equal probability. The random start has learned nothing yet.

Compute the gradients

backward() starts at the logits. For the first example it subtracts 1 from the probability of the target, “stop”, and divides every entry by the batch size, 3. The two “red” examples produce the same row; the “green” example produces its own.

The other gradients follow from that row:

• weights gradient: x transposed (2, 3), times the logit gradient (3, 4), gives (2, 4), the same shape as weights.
• bias gradient: the sum of the logit gradient over the batch, [0.2484, 0.2072, -0.3892, -0.0665]. The “stop” entry is negative, so the update raises the “stop” bias.
• Embedding gradient: the logit gradient times weights transposed gives one row per example, shape (3, 2). Each row is added to the embedding row it came from.

The last point matters. “red” appears twice, so its embedding row receives the sum of two contributions: [0.0473, 0.1643]. “stop” and “go” never appear as inputs, so their rows get zero gradient and do not change.

Check the gradients

numerical_gradients() nudges each of the 20 parameters by ε = 0.00001 in both directions. The largest difference between a hand-written gradient and its estimate is of the order of 10^-11 on the recorded run. The exact value depends on floating-point rounding, so the trace records only its power of ten and the test requires it to be below 10⁻⁶.

Apply one update

With learning rate 0.5, every parameter moves by −0.5 × its gradient. The “red” embedding row moves from [-0.1800, -0.3500] to [-0.2036, -0.4321]. The bias becomes [-0.1242, -0.1036, 0.1946, 0.0332]. Recomputing the forward pass with the new parameters gives a loss of 1.1345, down from 1.2952.

{
  "vocab": [
    "red",
    "green",
    "stop",
    "go"
  ],
  "inputs": [
    "red",
    "green",
    "red"
  ],
  "targets": [
    "stop",
    "go",
    "stop"
  ],
  "input_ids": [
    0,
    1,
    0
  ],
  "target_ids": [
    2,
    3,
    2
  ],
  "seed": 7,
  "learning_rate": 0.5,
  "shapes": {
    "input_ids": [
      3
    ],
    "embedding": [
      4,
      2
    ],
    "x": [
      3,
      2
    ],
    "weights": [
      2,
      4
    ],
    "logits": [
      3,
      4
    ],
    "bias": [
      4
    ]
  },
  "parameter_count": 20,
  "params_before": {
    "embedding": [
      [
        -0.18,
        -0.35
      ],
      [
        0.15,
        -0.43
      ],
      [
        0.04,
        -0.13
      ],
      [
        -0.44,
        0.01
      ]
    ],
    "weights": [
      [
        -0.46,
        -0.07,
        -0.43,
        -0.41
      ],
      [
        -0.08,
        0.33,
        -0.38,
        -0.28
      ]
    ],
    "bias": [
      0.0,
      0.0,
      0.0,
      0.0
    ]
  },
  "x": [
    [
      -0.18,
      -0.35
    ],
    [
      0.15,
      -0.43
    ],
    [
      -0.18,
      -0.35
    ]
  ],
  "logits": [
    [
      0.1108,
      -0.1029,
      0.2104,
      0.1718
    ],
    [
      -0.0346,
      -0.1524,
      0.0989,
      0.0589
    ],
    [
      0.1108,
      -0.1029,
      0.2104,
      0.1718
    ]
  ],
  "probs": [
    [
      0.251559,
      0.203156,
      0.277904,
      0.267381
    ],
    [
      0.242148,
      0.215239,
      0.276732,
      0.265881
    ],
    [
      0.251559,
      0.203156,
      0.277904,
      0.267381
    ]
  ],
  "losses": [
    1.28048,
    1.324706,
    1.28048
  ],
  "loss_before": 1.295222,
  "grads": {
    "embedding": [
      [
        0.047291,
        0.164298
      ],
      [
        0.018513,
        0.050684
      ],
      [
        0.0,
        0.0
      ],
      [
        0.0,
        0.0
      ]
    ],
    "weights": [
      [
        -0.01808,
        -0.013617,
        0.100488,
        -0.068792
      ],
      [
        -0.093405,
        -0.078254,
        0.128824,
        0.042835
      ]
    ],
    "bias": [
      0.248422,
      0.207184,
      -0.389153,
      -0.066452
    ]
  },
  "gradient_check": {
    "epsilon": 1e-05,
    "max_error_power_of_ten": -11,
    "max_error_overwrite_bug": 0.082149,
    "loss_after_overwrite_bug": 1.142513
  },
  "params_after": {
    "embedding": [
      [
        -0.203646,
        -0.432149
      ],
      [
        0.140743,
        -0.455342
      ],
      [
        0.04,
        -0.13
      ],
      [
        -0.44,
        0.01
      ]
    ],
    "weights": [
      [
        -0.45096,
        -0.063192,
        -0.480244,
        -0.375604
      ],
      [
        -0.033298,
        0.369127,
        -0.444412,
        -0.301417
      ]
    ],
    "bias": [
      -0.124211,
      -0.103592,
      0.194577,
      0.033226
    ]
  },
  "loss_after": 1.134459,
  "learning_rate_sweep": [
    {
      "learning_rate": 0.1,
      "loss_after": 1.261467
    },
    {
      "learning_rate": 0.5,
      "loss_after": 1.134459
    },
    {
      "learning_rate": 1,
      "loss_after": 0.99613
    },
    {
      "learning_rate": 2,
      "loss_after": 0.797525
    },
    {
      "learning_rate": 5,
      "loss_after": 0.751234
    },
    {
      "learning_rate": 10,
      "loss_after": 1.401903
    },
    {
      "learning_rate": 20,
      "loss_after": 3.277132
    }
  ],
  "large_learning_rate": {
    "learning_rate": 20,
    "loss_after": 3.277132
  }
}

Check your understanding

The trace also records the loss after one step at 7 learning rates, all from the same starting weights and batch. Before you open the comparison, decide which rate you would use to train this model, and what this table can and cannot tell you.

Failure cases

Each case below has a named test in test_tiny_step.py.

The learning-rate case is specific to this batch and these starting weights. It shows that a rate can be too large; it does not give a safe range for any other model.

Download test_tiny_step.py and step-trace.json beside tiny_step.py, then run the 15 standard-library tests:

python3 -m unittest -v test_tiny_step.py

"""Regression and failure tests for tiny_step.py. Run: python3 -m unittest -v test_tiny_step.py"""
import json
import math
import unittest
from pathlib import Path

import tiny_step as t

TOLERANCE = 1e-6


def batch():
    return t.init_params(), t.encode(t.INPUTS), t.encode(t.TARGETS)


def close(a, b, tol=TOLERANCE):
    if isinstance(a, dict):
        return a.keys() == b.keys() and all(close(a[k], b[k], tol) for k in a)
    if isinstance(a, list):
        return len(a) == len(b) and all(close(x, y, tol) for x, y in zip(a, b))
    if isinstance(a, (int, float)) and isinstance(b, (int, float)):
        return abs(a - b) <= tol
    return a == b


class RecordedTrace(unittest.TestCase):
    def test_recorded_trace(self):
        recorded = json.loads(Path(__file__).with_name("step-trace.json").read_text())
        actual = t.trace()
        self.assertLess(actual["gradient_check"]["max_error_power_of_ten"], -6)
        self.assertTrue(close(recorded, actual))


class Mechanism(unittest.TestCase):
    def test_softmax_handles_large_scores(self):
        self.assertEqual(t.softmax([1000.0, 0.0]), [1.0, math.exp(-1000.0)])

    def test_shapes(self):
        params, ids, targets = batch()
        cache = t.forward(params, ids, targets)
        self.assertEqual(t.shape(cache["x"]), (3, 2))
        self.assertEqual(t.shape(cache["logits"]), (3, 4))
        self.assertTrue(all(math.isclose(sum(row), 1.0) for row in cache["probs"]))

    def test_gradient_check_passes(self):
        params, ids, targets = batch()
        analytic = t.backward(params, ids, targets, t.forward(params, ids, targets))
        numeric = t.numerical_gradients(params, ids, targets)
        self.assertLess(t.max_gradient_error(analytic, numeric), TOLERANCE)

    def test_unused_embedding_rows_get_zero_gradient(self):
        params, ids, targets = batch()
        grads = t.backward(params, ids, targets, t.forward(params, ids, targets))
        for token in ("stop", "go"):
            self.assertEqual(grads["embedding"][t.VOCAB.index(token)], [0.0, 0.0])

    def test_one_step_lowers_loss(self):
        params, ids, targets = batch()
        before, _, _, after = t.train_step(params, ids, targets, t.LEARNING_RATE)
        self.assertLess(after, before)

    def test_step_does_not_modify_input_params(self):
        params, ids, targets = batch()
        snapshot = json.dumps(params)
        t.train_step(params, ids, targets)
        self.assertEqual(json.dumps(params), snapshot)


class FailureCases(unittest.TestCase):
    def test_repeated_ids_must_accumulate(self):
        params, ids, targets = batch()
        cache = t.forward(params, ids, targets)
        numeric = t.numerical_gradients(params, ids, targets)
        buggy = t.backward(params, ids, targets, cache, accumulate=False)
        self.assertGreater(t.max_gradient_error(buggy, numeric), 0.01)

    def test_large_learning_rate_raises_loss(self):
        params, ids, targets = batch()
        before, _, _, after = t.train_step(params, ids, targets, 20)
        self.assertGreater(after, before)

    def test_unknown_token(self):
        with self.assertRaisesRegex(ValueError, "not in the vocabulary"):
            t.encode(["red", "blue"])

    def test_token_id_out_of_range(self):
        params, _, _ = batch()
        for ids, targets in (([0, 4], [2, 3]), ([-1], [2]), ([0], [4]), ([0], [-1]), ([True], [2])):
            with self.subTest(ids=ids, targets=targets), self.assertRaisesRegex(ValueError, "outside the vocabulary"):
                t.forward(params, ids, targets)

    def test_batch_length_mismatch(self):
        params, _, _ = batch()
        with self.assertRaisesRegex(ValueError, "3 inputs but 2 targets"):
            t.forward(params, [0, 1, 0], [2, 3])

    def test_empty_batch(self):
        params, _, _ = batch()
        with self.assertRaisesRegex(ValueError, "empty batch"):
            t.forward(params, [], [])

    def test_matmul_shape_mismatch(self):
        with self.assertRaisesRegex(ValueError, "inner sizes 2 and 3 differ"):
            t.matmul([[1.0, 2.0]], [[1.0], [2.0], [3.0]])

    def test_invalid_learning_rate(self):
        params, ids, targets = batch()
        grads = t.backward(params, ids, targets, t.forward(params, ids, targets))
        for bad in (0, -0.1, float("nan"), float("inf")):
            with self.subTest(learning_rate=bad), self.assertRaisesRegex(ValueError, "positive finite"):
                t.sgd_step(params, grads, bad)


if __name__ == "__main__":
    unittest.main()

test_recorded_trace compares a fresh run with the recorded trace. test_gradient_check_passes and test_one_step_lowers_loss check the mechanism. The other tests are the failure cases above. If you change the formula for the logit gradient, for example by removing the division by the batch size, the gradient-check test fails.

Deployment

There is nothing to deploy. tiny_step.py is a command-line program that runs one step and exits. It does not save parameters, so nothing persists between runs. Run the tests on any machine where you change the code; the tests need no network access and no packages.

Scaling

Every operation loops over Python lists. One step costs about B × d × V multiplications for the logits and about twice that for the two products in the backward pass, where B is the batch size, d the embedding size and V the vocabulary size. The gradient check runs two full forward passes per parameter, so its cost grows with the number of parameters times the cost of one forward pass. That is fine for 20 parameters and impractical for a real model.

To scale the same computation, three things change. Array libraries replace the list loops. Automatic differentiation replaces the hand-written backward(). A gradient check runs only on small test cases, as a test of the differentiation code. This build does not measure timings, so it makes no speed claim.

Key takeaways

• One training step is a forward pass to a loss, a backward pass to one gradient per parameter, and an update. Every gradient has the shape of its parameter.
• For softmax with mean cross-entropy, the logit gradient is the probabilities minus the one-hot target, divided by the batch size.
• An embedding lookup copies rows, so its backward pass must add gradients into rows. A repeated token gets the sum of its contributions; an unused token gets zero.

Best practices

• Check a new backward pass against finite differences on a tiny case before you train with it. A falling loss does not prove the gradients are right.
• Validate token IDs, batch lengths and shapes at the boundary, and fail with a message that names the bad value.
• Record the seed, learning rate and starting parameters with any trace you keep, so a later run can reproduce it.

References

Official PyTorch 2.14 documentation, checked October 20, 2026. The program does not use PyTorch; these pages define the library operations that match each part of tiny_step.py.

• torch.nn.Embedding: a lookup table with weight shape (num_embeddings, embedding_dim), like embedding here.
• torch.nn.CrossEntropyLoss: takes unnormalized logits and class-index targets. With no class weights and no ignored targets, its default mean reduction is the batch mean used here.
• torch.Tensor.backward: accumulates gradients in the leaf tensors, which is why training loops reset gradients between steps.
• torch.optim.SGD: without momentum or weight decay, the update is the parameter minus the learning rate times the gradient.
• torch.autograd.gradcheck: compares analytical gradients with finite differences and is designed for double-precision inputs. Python floats are double precision.