mojolearn.Adam#

class mojolearn.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, lr_schedule=None, accumulation_steps=1, maximize=False, resident=None, **kwargs)#

Adam optimizer, as torch.optim.Adam, over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object).

Weight decay is coupled, as in torch’s Adam: it is added to the gradient and so passes through both moments. AdamW is the decoupled form; at weight_decay=0 the two compute the same thing. The bias corrections are computed in float32, not in torch’s float64, so results are close to torch’s but not bit-equal to them. Parameters are updated in place.

Parameters:

params (buffer or list of buffers) – float32 parameter arrays, updated in place.
lr (float, default 1e-3) – Learning rate.
betas (tuple of float, default (0.9, 0.999)) – Decay rates of the first and second moments, each in [0, 1); 1.0 is refused.
eps (float, default 1e-8) – Added to sqrt(v) in the denominator (sqrt(v) + eps).
weight_decay (float, default 0.0) – L2 penalty, added to the gradient.
lr_schedule (object or None, default None) – A schedule with a one-based lr_at(t) method; it sets lr before every step. Anything without lr_at raises TypeError.
accumulation_steps (int, default 1) – Number of microbatches combined by step_accumulated; a power of two >= 1, otherwise ValueError.
maximize (bool, default False) – Maximize the objective (the gradient’s sign is flipped inside the step; the caller’s gradient arrays are not negated). Must be a bool.
resident (bool or None, default None) – Keep the optimizer state on the device between steps when the install supports it (None or True). False keeps it on the host, which suits an optimizer that is rebuilt every step.

Variables:

t (int) – Number of steps taken.
exp_avg (array) – First moment, flat over all parameters.
exp_avg_sq (array) – Second moment, flat over all parameters.

Notes

step(grads, max_norm=None) takes one gradient per parameter, matching params in count, shape and order. With max_norm set, the global L2 gradient norm is clipped first: the caller’s gradient arrays are scaled in place (as torch.nn.utils.clip_grad_norm_) and step returns the pre-clip total norm; without it, step returns None and the gradients are untouched. With clipping on, the order of params is part of the result, because it fixes the summation order of the norm.

lr is read at every step, so you can assign opt.lr between steps, or pass an lr_schedule (ConstantLR, WarmupLinearLR, WarmupCosineLR, or any object with a one-based lr_at(t)). step_accumulated(microbatch_grads, tokens) combines accumulation_steps microbatch gradient lists before one step. zero_grad(grads) writes +0.0 into the given gradient arrays. state_dict() / load_state_dict() save and restore the step count and the optimizer state.

Refused with TypeError: foreach, fused, capturable, differentiable, parameter groups, sparse gradients, float16/bfloat16 and float64 buffers, and any unknown keyword. amsgrad is not supported and raises TypeError.

Examples

>>> import numpy as np, mojolearn
>>> w = np.zeros((64, 8), dtype=np.float32)
>>> opt = mojolearn.Adam([w], lr=1e-3)
>>> opt.step([np.ones_like(w)])

See also

User guide: Neural networks

Methods

`get_params`([deep])
`grads_to_device`(grads)	Upload a gradient into the handle's device gradient buffer, for `step_device()`.
`grads_to_host`(grads)	Download the device gradient (clipped, if the last step clipped) into `grads`, in place.
`load_state_dict`(state)	Restore the optimizer state saved by `state_dict`.
`numeric_mode_used`()	The numeric mode this estimator runs in.
`param_device`(j)	Parameter tensor `j` as a `TrainDeviceArray` (a device copy of the resident parameters; `params_to_device()` first), for the next forward.
`params_resident_available`()	Whether `params_to_device` works on this build.
`params_to_device`()	Upload the parameters into the optimizer's device handle and keep them there: later steps neither upload nor download them.
`params_to_host`()	Download the device parameters into the caller's arrays (in place); they stay resident.
`set_params`(**params)	Set constructor parameters.
`state_dict`()	Everything a resumed run needs, and nothing that is not state.
`step`(grads[, max_norm])	One step.
`step_accumulated`(microbatch_grads, tokens[, ...])	One step from `A = self.accumulation_steps` microbatch gradient lists, combined on device by the clause 9.2 balanced tree in ascending microbatch index, then `step`.
`step_device`([grads, max_norm])	`step` on device-resident parameters (`params_to_device()` first).
`vendor_used`()	The GPU backend this estimator runs on.
`zero_grad`(grads)	`torch.optim.Optimizer.zero_grad`, over arrays the caller owns.

Attributes

`exp_avg`	Adam's first moment (SGD's momentum buffer), a `mojolearn.Array` of `n_total` floats.
`exp_avg_sq`	Adam's second moment; see `exp_avg`.
`resident_`	True once the moments live on the device.
`buf_initialized`	SGD's per-tensor `buf_initialized` flag, contract 7.3b.
`t`	ONE-BASED.
`total_norm_`	The pre-clip global gradient norm of the most recent step, or None when that step ran with `max_norm=None`.
`clip_coef_`	The clamped coefficient the most recent step applied, or None.

property exp_avg#: Adam’s first moment (SGD’s momentum buffer), a mojolearn.Array of n_total floats. With the moments resident this is a download on access, so read it once per use, not per element.

property exp_avg_sq#: Adam’s second moment; see exp_avg.

grads_to_device(grads)#: Upload a gradient into the handle’s device gradient buffer, for step_device(). Returns self.

grads_to_host(grads)#: Download the device gradient (clipped, if the last step clipped) into grads, in place. Returns grads.

load_state_dict(state)#

Restore the optimizer state saved by state_dict.

Parameters:: state (dict) – The dictionary returned by state_dict. Its arrays may be any buffer or nested list and are converted to float32 (int32 for buf_initialized). The optimizer keeps its own copy, so later steps never modify state.

numeric_mode_used()#

The numeric mode this estimator runs in.

Returns:: str – 'identical' or 'fast', read from the compiled binary the estimator uses, not from the argument it was given.

param_device(j)#: Parameter tensor j as a TrainDeviceArray (a device copy of the resident parameters; params_to_device() first), for the next forward. Its shape is the host tensor’s.

params_resident_available()#: Whether params_to_device works on this build.

params_to_device()#: Upload the parameters into the optimizer’s device handle and keep them there: later steps neither upload nor download them. Call it again after editing the host arrays. Returns self.

params_to_host()#: Download the device parameters into the caller’s arrays (in place); they stay resident. Returns the parameter list.

property resident_#: True once the moments live on the device.

set_params(**params)#: Set constructor parameters. The estimator is rebuilt through its constructor, so every constructor check runs again and any fitted state is dropped, as a new setting invalidates it.

state_dict()#

Everything a resumed run needs, and nothing that is not state.

t, exp_avg, exp_avg_sq and buf_initialized. Leaving buf_initialized out is the quiet way to break a resume: SGD’s first step after the reload would take the COPY arm again and overwrite a momentum buffer it should have continued, which is the failure the optimizer contract’s clause (d) control 2 is built to catch.

step(grads, max_norm=None)#

One step. grads matches params in count, shape and order.

max_norm runs the global-norm clip BEFORE the elementwise update, over the whole registry, and SCALES THE CALLER’S GRADIENT ARRAYS IN PLACE, which is what torch.nn.utils.clip_grad_norm_ does and why it carries a trailing underscore. None turns it off; the gradient is then untouched.

WITH CLIPPING ON, ONE PARAMETER’S UPDATE IS NOT INDEPENDENT OF THE REST OF THE MODEL. The coefficient is a function of every gradient in the registry, by the reference’s own semantics and not by a defect (optimizer contract 3.5). That is why the profile’s parameter-count invariance clause is stated for max_norm=None.

Returns the pre-clip total gradient norm as a float when the clip ran, and None when it did not. +0.0 would have been the shorter spelling and it is the wrong one: a coefficient of 1.0 says the clip ran and found nothing to do, which is a different fact from the clip not running, and optimizer_step_oracle draws the same line by leaving its clip.* stages empty rather than filling them.

step_accumulated(microbatch_grads, tokens, max_norm=None)#: One step from A = self.accumulation_steps microbatch gradient lists, combined on device by the clause 9.2 balanced tree in ascending microbatch index, then step. tokens is the FULL step’s token count T; a split that clause 9.2 does not admit is refused by name (no gradient is combined). Returns what step returns.

step_device(grads=None, max_norm=None)#: step on device-resident parameters (params_to_device() first). grads given: uploaded (one transfer), and scaled in place on the host when max_norm clips, as step does. grads None: the gradient already in the handle (grads_to_device) is used and, with max_norm, clipped there; nothing crosses the bus but the three info words. Returns what step returns.

vendor_used()#

The GPU backend this estimator runs on.

Returns:: str – 'metal', 'cuda' or 'hip', read from the compiled binary the estimator uses.

zero_grad(grads)#: torch.optim.Optimizer.zero_grad, over arrays the caller owns. Present so that a training loop written against torch reads the same; it writes +0.0 and nothing else – as all-zero bytes through memset on the buffer’s own address (DEVIATION 2406), which is +0.0 in float32 and needs the buffer C-contiguous and writable.

buf_initialized#: SGD’s per-tensor buf_initialized flag, contract 7.3b. 0 means this tensor’s momentum buffer has never been written, so the first step COPIES the gradient into it instead of running the recurrence. It is CARRIED STATE and belongs in a checkpoint beside exp_avg; the flag flips once per TENSOR after that tensor’s kernel, never per element.

t#: ONE-BASED. The first step of a run is t = 1, and training/estimator.mojo refuses t < 1 by name because at t = 0 the bias correction 1 - beta^0 is exactly zero and the host scalars divide by it.

total_norm_#: The pre-clip global gradient norm of the most recent step, or None when that step ran with max_norm=None.

clip_coef_#: The clamped coefficient the most recent step applied, or None.