Optimizers

Optimizer state

sgd(m, lr=…) and adamw(m, lr=…) make optimizer state for the model m. Its type comes from the model’s: adamw(m, lr=3e-4) with m: Mlp is an AdamW[Mlp], and sgd(m, lr=0.1) an SGD[Mlp]. It covers every param of the model, so none can be left out, and holds what the optimizer keeps between steps (AdamW’s moments, SGD’s momentum) and how many updates it has made.

fn fit(mut m: Mlp, x: [b D], y: [b] Index[C], steps: int) -> AdamW[Mlp]:
    opt = adamw(m, lr=warmup(100, cosine(3e-4, steps)))
    for i in 0..steps:
        l, g = value_and_grad(loss, m, x, y)
        update(mut m, mut opt, g |> clip_norm(1.0))
    return opt

The model comes first, by position; anything else there is E1501. The hyperparameters are Python numbers, given by name (E0406 for one given by position, an unknown name, or a missing lr), with PyTorch’s defaults:

lr is a learning rate, a number or a Schedule, and has no default. The other parameters are host scalars: a tensor is E0805 and another kind of number E0502, as for any float or bool parameter. Their values are checked by PyTorch when the state is made, so a negative lr raises a ValueError then.

AdamW[M] and SGD[M] are types, for parameters and return types. Optimizer state is matched by optimizer and model: state of another model’s type is E1502, and anything else, such as SGD[Mlp] where AdamW[Mlp] is expected, E1503. AdamW or SGD of something other than a model, as AdamW[M] and the like are E0303. A parameter whose type is optimizer state can’t have a default (E0306), and optimizer state can’t be a field of a model, a data record or a logs record. Like a model, it is changed only through a mut parameter or a value the function made (Changing models with mut); a mut parameter’s state can’t be replaced with new state (E1104).

update

update(mut m, mut opt, g) applies one step of the optimizer state opt to the model m, with the gradient g, at the learning rate opt’s schedule gives for the updates it has made so far, and counts the update. It changes m and opt, so both are written mut (E1103) and must be changeable here (E1102); g isn’t changed (mut g is E1101). It takes its three arguments by position (E0406) and has no result.

The model, the state and the gradient must be of the same model type: opt or g of another model is E1502, comparing with m’s type (or opt’s, when m is a Python value). A first argument that isn’t a model or a third that isn’t a gradient is E1505, and a second that isn’t optimizer state is E1503.

Gradient transforms

A gradient transform is a function from Grad[M] to Grad[M], used with pipes: g |> clip_norm(1.0). clip_norm(g, max_norm) scales the gradient so that the norm of all its values together, norm(g), is at most max_norm: by min(1, max_norm / (norm(g) + 1e-6)), as torch.nn.utils.clip_grad_norm_ does. max_norm is a float, by position or by name. Its first argument must be a gradient (E1505). A transform of one’s own is an ordinary fn, using the gradient arithmetic:

fn halve(g: Grad[Mlp]) -> Grad[Mlp] = g * 0.5

fn step(mut m: Mlp, mut opt: SGD[Mlp], g: Grad[Mlp]):
    update(mut m, mut opt, g |> clip_norm(1.0) |> halve)

Schedules

A learning rate is a number or a schedule, a value of type Schedule that gives one for each step. Schedules are made with:

• cosine(peak, steps, end=0.0): $\mathit{end} + (\mathit{peak} - \mathit{end}) \cdot \frac{1 + \cos(\pi \min(t, \mathit{steps}) / \mathit{steps})}{2}$ at step $t$. It starts at peak, follows a half cosine down to end at step steps, and stays at end after. Up to steps it’s torch.optim.lr_scheduler.CosineAnnealingLR with T_max=steps and eta_min=end, which would rise again after.
• warmup(steps, then): rises in a straight line from 0, $\mathit{then}(0) \cdot t / \mathit{steps}$ for $t < \mathit{steps}$, then is then started afresh, $\mathit{then}(t - \mathit{steps})$. then is a number or a schedule, so warmup(100, cosine(3e-4, 10000)) warms up for 100 steps and decays over 10000 more.

peak and end are floats and steps an int, by position or by name (E0406, E0502, E0805); steps must be at least 1 for cosine and 0 for warmup, which is checked when the schedule is made. then and lr= take a number or a schedule; anything else, such as a [] tensor, is E1504. Schedule is also a type, for parameters and return types, and a number is accepted where a Schedule is expected, as a constant learning rate. A schedule supports no arithmetic.

The step count lives in the optimizer state. opt starts at 0 updates; each update reads the schedule at the number of updates made before it, then adds one. So the first update uses the learning rate for step 0, as a PyTorch scheduler’s first optimizer.step() does, and a schedule is the same value however many optimizers use it. Saving the state saves the count and the schedule.

Optimizers at run time

sgd and adamw compile to rt.sgd and rt.adamw, which make an rt.Optimizer around torch.optim.SGD or torch.optim.AdamW over the model’s params, in field order. update compiles to rt.update(m, opt, g): it sets the learning rate from the schedule, gives each param its gradient as .grad, calls the PyTorch optimizer’s step(), and sets .grad back to None, even if the step fails, so no .grad is left behind and the params stay the same nn.Parameter objects, changed in place. So a step is exactly PyTorch’s. The state is per param of the model’s type, not per model object: given another model of the same type, update moves the state to that model’s params. clip_norm is rt.clip_norm, which makes a new rt.Grad.

AdamW is made with capturable=True where PyTorch supports it, so that a step can be recorded as a graph (FAST-1): PyTorch allows it on CUDA (and XPU, HPU and XLA devices), but not on the CPU or Apple’s MPS, where the state is made without it. A capturable optimizer keeps its learning rate in a device tensor, which update fills before each step. torch.optim.SGD has no capturable option and keeps no step count of its own, so it needs none.

opt.state_dict() gives everything needed to continue: the optimizer ("sgd" or "adamw"), the model’s type and param names, the learning rate (a number, or a schedule as a dict such as {"schedule": "cosine", "peak": 0.0003, "steps": 1000, "end": 0.0}), the update count, and PyTorch’s own state dict, with the learning rate as a number. opt.load_state_dict(state) restores it into state of the same optimizer and model type, and raises ShapeError for any other. opt.steps is the update count, opt.lr the learning rate and opt.torch the PyTorch optimizer. From Python, shapescript.rt.adamw(m, lr=…), shapescript.rt.sgd(m, lr=…), shapescript.rt.cosine(…) and shapescript.rt.warmup(…) make the same values. A parameter whose type is optimizer state checks that its argument is state of that optimizer for that model type, and a Schedule parameter that its argument is a schedule or a number (Runtime checks).