mojolearn.AdamW#

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

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

Identical to Adam except that weight decay is decoupled: the parameter is multiplied by 1 - lr * weight_decay and the gradient is left untouched, so the decay does not pass through the moments. Parameters, attributes and refusals are Adam’s.

Note that the default weight_decay here is 0.0, while torch’s AdamW default is 0.01; at 0.0 AdamW and Adam compute the same thing. Pass weight_decay=0.01 to match torch’s default.

Examples

>>> import numpy as np, mojolearn
>>> w = np.zeros((64, 8), dtype=np.float32)
>>> opt = mojolearn.AdamW([w], lr=1e-3, weight_decay=0.01)
>>> 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.