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)[source]#

Bases: Adam

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)])

__init__(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)[source]#

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.