mojolearn.training#

Optimizers, schedules, losses and training primitives over float32 arrays.

These functions work on explicit float32 buffers: mojolearn.Array, NumPy arrays, or any buffer-protocol object. NumPy is not required.

SGD, Adam and AdamW update the parameter arrays you pass in place, from gradients you supply in the same order. ConstantLR, WarmupLinearLR and WarmupCosineLR give the learning rate of each step. clip_grad_norm_ scales gradients in place (L2 norm only). cross_entropy takes float32 logits and integer class targets and can return the logits gradient in the same call. embedding_forward, rms_norm_forward, linear_forward, chunked_lm_head_loss and their backward functions are the layer operations of a small language-model stack, and Generator draws reproducible initializations and dropout masks.

These are building blocks, not an autograd framework or a drop-in replacement for PyTorch: you compute and pass the gradients. SmallMLPTrainer and the byte language model trainers provide complete training loops for their fixed model shapes.

numeric_mode selects the compiled tier where a function takes it. No compiled extension is loaded until a function or optimizer step needs it.

Examples

>>> import numpy as np
>>> from mojolearn.training import AdamW, WarmupCosineLR
>>> w = np.zeros((64, 8), dtype=np.float32)
>>> opt = AdamW([w], lr=3e-4, weight_decay=0.01,
...             lr_schedule=WarmupCosineLR(3e-4, warmup_steps=10, total_steps=100))
>>> opt.step([np.ones_like(w)])

Functions

`clip_grad_norm_`(grads, max_norm[, ...])	Clip gradients by their global L2 norm, as `torch.nn.utils.clip_grad_norm_`.
`cross_entropy`(logits, targets[, ...])	Cross-entropy loss over class indices, as `torch.nn.functional.cross_entropy`.
`numeric_mode_used`([numeric_mode])	Return the numeric mode the training binding will run in.
`vendor_used`([numeric_mode])	Return the accelerator API of the training binding: 'metal', 'cuda' or 'hip'.
`accumulate_grads`(microbatch_grads, tokens[, ...])	Combine microbatch gradients into one, on the device.
`accumulation_is_aligned`(tokens, ...[, ...])	Whether a split into microbatches gives exactly the unsplit gradient.
`embedding_forward`(weight, ids[, numeric_mode])	Embedding lookup: `weight[ids]`.
`embedding_backward`(dy, ids, vocab[, ...])	Gradient of `embedding_forward` with respect to the table.
`rms_norm_forward`(x, weight, eps[, numeric_mode])	RMS normalization of each row: `weight * x / sqrt(mean(x**2) + eps)`.
`rms_norm_backward`(dy, x, weight, eps[, ...])	Gradients `(dx, dweight)` of `rms_norm_forward`.
`linear_forward`(a, weight[, numeric_mode])	Linear layer without bias: `a @ weight.T`, as `torch.nn.Linear` with `bias=False`.
`linear_backward`(dc, a, weight[, numeric_mode])	Gradients `(da, dweight)` of `linear_forward`.
`chunked_lm_head_loss`(hidden, weight, targets)	Mean cross-entropy of a language-model head without materializing the logits.
`mlp_inference_sessions`(inputs, weight1, ...)	Run several independent batches of rows through one Linear-ReLU-Linear model.
`residual_dropout`(values, residual, *[, p, ...])	Return `residual + dropout(values)` as one explicit IDENTICAL layer.
`residual_dropout_backward`(gradient, *[, p, ...])	Return `(d_values, d_residual)` for the explicit residual layer.

Classes

`SGD`(params[, lr, momentum, dampening, ...])	Stochastic gradient descent, as `torch.optim.SGD`, over float32 buffers (NumPy arrays, `array.array` or `mojolearn.Array`; any writable buffer-protocol object).
`Adam`(params[, lr, betas, eps, weight_decay, ...])	Adam optimizer, as `torch.optim.Adam`, over float32 buffers (NumPy arrays, `array.array` or `mojolearn.Array`; any writable buffer-protocol object).
`AdamW`(params[, lr, betas, eps, ...])	AdamW optimizer, as `torch.optim.AdamW`, over float32 buffers (NumPy arrays, `array.array` or `mojolearn.Array`; any writable buffer-protocol object).
`ConstantLR`(peak_lr[, warmup_steps, ...])	Constant learning rate, after an optional linear warmup.
`WarmupLinearLR`(peak_lr[, warmup_steps, ...])	Linear warmup, then linear decay to `min_lr`.
`WarmupCosineLR`(peak_lr[, warmup_steps, ...])	Linear warmup, then cosine decay to `min_lr`.
`Generator`(seed[, numeric_mode])	Seeded random stream for neural-network initialization and dropout.
`SambaConfig`(vocab, d_model, layers[, ...])	Shape of a `SambaStack`: vocabulary, width and the sequence of layer kinds.
`SambaStack`(config[, weights, generator, lr, ...])	A Samba-style language model with its own training loop.

class mojolearn.training.SGD(params, lr=0.001, momentum=0.0, dampening=0.0, weight_decay=0.0, nesterov=False, lr_schedule=None, accumulation_steps=1, maximize=False, resident=None, **kwargs)[source]#

Bases: _Optimizer

Stochastic gradient descent, as torch.optim.SGD, over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object).

Supports momentum, dampening, Nesterov momentum and weight decay. Weight decay is coupled, as in torch: it is added to the gradient and so passes through the momentum buffer. 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. torch’s SGD has no default; this one is 1e-3.
momentum (float, default 0.0) – Momentum factor.
dampening (float, default 0.0) – Dampening of the momentum.
weight_decay (float, default 0.0) – L2 penalty, added to the gradient.
nesterov (bool, default False) – Nesterov momentum. Needs momentum > 0 and dampening == 0, otherwise ValueError (as torch).
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.

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.

Examples

>>> import numpy as np, mojolearn
>>> w = np.zeros((64, 8), dtype=np.float32)
>>> opt = mojolearn.SGD([w], lr=1e-2, momentum=0.9)
>>> g = np.ones_like(w)
>>> norm = opt.step([g], max_norm=1.0)

class mojolearn.training.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)[source]#

Bases: _Optimizer

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

class mojolearn.training.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)])

mojolearn.training.clip_grad_norm_(grads, max_norm, norm_type=2.0, error_if_nonfinite=True, numeric_mode=None, **kwargs)[source]#

Clip gradients by their global L2 norm, as torch.nn.utils.clip_grad_norm_.

Works over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object). Computes a norm per tensor, then the norm of those norms, and scales every gradient in place by max_norm / total when the total exceeds max_norm. Returns the pre-clip total norm, as torch does. The order of the tensors is part of the result, because it fixes the summation order.

Parameters:

grads (buffer or list of buffers) – float32 gradient arrays, scaled in place.
max_norm (float) – Maximum total norm; must be > 0, otherwise ValueError.
norm_type (float, default 2.0) – Only 2.0 is supported; any other value (including inf) raises NotImplementedError.
error_if_nonfinite (bool, default True) – Only True is supported (False raises NotImplementedError): a non-finite total norm always raises.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

total_norm (float) – The total gradient norm before clipping.

Notes

Not supported: foreach, per-tensor norms in the result, float16/bfloat16 and float64. Each tensor costs one device synchronization, so many small tensors are slower than a few large ones.

Examples

>>> import numpy as np, mojolearn
>>> g1 = np.ones((64, 8), dtype=np.float32); g2 = np.ones(8, dtype=np.float32)
>>> total = mojolearn.clip_grad_norm_([g1, g2], max_norm=1.0)

mojolearn.training.cross_entropy(logits, targets, ignore_index=-100, reduction='mean', label_smoothing=0.0, num_items=None, return_grad=False, numeric_mode=None, **kwargs)[source]#

Cross-entropy loss over class indices, as torch.nn.functional.cross_entropy.

logits is (N, V) float32 and targets is (N,) int32 or int64 holding a class index or ignore_index. Returns a float for reduction='mean' or 'sum' and an (N,) float32 array for 'none'. With return_grad=True the forward and the gradient with respect to the logits come from one call: (loss, dlogits), dlogits of shape (N, V).

Parameters:

logits (buffer, shape (N, V)) – float32 logits, row-major.
targets (buffer, shape (N,)) – Integer class indices; a value that is neither a valid class nor ignore_index raises an error naming its row.
ignore_index (int, default -100) – Targets equal to this are ignored, even when it is a valid class index (as torch).
reduction ({‘mean’, ‘sum’, ‘none’}, default ‘mean’) – Reduction over rows. 'none' has no gradient: return_grad=True with it is refused.
label_smoothing (float, default 0.0) – Label smoothing in [0, 1).
num_items (int or None, default None) – Divisor of the mean, as Hugging Face’s num_items_in_batch. None uses the number of non-ignored targets.
return_grad (bool, default False) – Also return the gradient with respect to logits.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

loss (float or ndarray) – The loss; (N,) per-row values under reduction='none'.
dlogits (array, shape (N, V)) – Only with return_grad=True.

Notes

Refused: per-class weight, the deprecated size_average and reduce, class-probability (soft) targets, float16/bfloat16 and float64. Memory use is several (N, V) float32 buffers; for a large vocabulary, split the rows into chunks (with reduction='none' the chunked result equals the unsplit one).

Examples

>>> import numpy as np, mojolearn
>>> logits = np.random.default_rng(0).random((32, 10), dtype=np.float32)
>>> targets = np.arange(32, dtype=np.int64) % 10
>>> loss, dlogits = mojolearn.cross_entropy(logits, targets, return_grad=True)

mojolearn.training.numeric_mode_used(numeric_mode=None)[source]#

Return the numeric mode the training binding will run in.

Read from the compiled binding itself, not from the requested string.

Parameters:: numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.
Returns:: str

mojolearn.training.vendor_used(numeric_mode=None)[source]#

Return the accelerator API of the training binding: ‘metal’, ‘cuda’ or ‘hip’.

Read from the binding’s compile-time constant, not from the platform.

Parameters:: numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.
Returns:: str

class mojolearn.training.ConstantLR(peak_lr, warmup_steps=0, total_steps=None, min_lr=0.0)[source]#

Bases: _Schedule

Constant learning rate, after an optional linear warmup.

Returns peak_lr at every step after warmup_steps steps of linear warmup starting from peak_lr / warmup_steps. total_steps and min_lr are accepted and ignored.

Parameters:

peak_lr (float) – Learning rate after warmup; finite and >= 0 (rounded to float32).
warmup_steps (int, default 0) – Steps of linear warmup; >= 0.
total_steps (int or None, default None) – Step at which the decay reaches min_lr; >= max(1, warmup_steps).
min_lr (float, default 0.0) – Final learning rate; finite and >= 0.

Notes

lr_at(t) takes the optimizer’s one-based step: lr_at(1) is the first step’s rate. Each value is the float32 nearest the exact expression, so a schedule gives the same bits on every machine. config() returns the schedule’s settings as a dict, which from_config turns back into a schedule. Invalid arguments raise ValueError. Pass a schedule as lr_schedule= to SGD, Adam or AdamW.

Examples

>>> import mojolearn
>>> sched = mojolearn.ConstantLR(1e-3, warmup_steps=100)
>>> lr = sched.lr_at(1)

class mojolearn.training.WarmupLinearLR(peak_lr, warmup_steps=0, total_steps=None, min_lr=0.0)[source]#

Bases: _Schedule

Linear warmup, then linear decay to min_lr.

Rises linearly over warmup_steps, falls in a straight line from peak_lr to min_lr at total_steps, then stays at min_lr.

Parameters:

peak_lr (float) – Learning rate after warmup; finite and >= 0 (rounded to float32).
warmup_steps (int, default 0) – Steps of linear warmup; >= 0.
total_steps (int or None, default None) – Step at which the decay reaches min_lr; >= max(1, warmup_steps).
min_lr (float, default 0.0) – Final learning rate; finite and >= 0.

Notes

lr_at(t) takes the optimizer’s one-based step: lr_at(1) is the first step’s rate. Each value is the float32 nearest the exact expression, so a schedule gives the same bits on every machine. config() returns the schedule’s settings as a dict, which from_config turns back into a schedule. Invalid arguments raise ValueError. Pass a schedule as lr_schedule= to SGD, Adam or AdamW.

Examples

>>> import mojolearn
>>> sched = mojolearn.WarmupLinearLR(1e-3, warmup_steps=100, total_steps=1000)
>>> lr = sched.lr_at(500)

class mojolearn.training.WarmupCosineLR(peak_lr, warmup_steps=0, total_steps=None, min_lr=0.0)[source]#

Bases: _Schedule

Linear warmup, then cosine decay to min_lr.

Rises linearly over warmup_steps, then follows min_lr + (peak_lr - min_lr) * (1 + cos(pi * p)) / 2 with p = (t - warmup_steps) / (total_steps - warmup_steps), then stays at min_lr.

Parameters:

peak_lr (float) – Learning rate after warmup; finite and >= 0 (rounded to float32).
warmup_steps (int, default 0) – Steps of linear warmup; >= 0.
total_steps (int or None, default None) – Step at which the decay reaches min_lr; >= max(1, warmup_steps).
min_lr (float, default 0.0) – Final learning rate; finite and >= 0.

Notes

lr_at(t) takes the optimizer’s one-based step: lr_at(1) is the first step’s rate. Each value is the float32 nearest the exact expression, so a schedule gives the same bits on every machine. config() returns the schedule’s settings as a dict, which from_config turns back into a schedule. Invalid arguments raise ValueError. Pass a schedule as lr_schedule= to SGD, Adam or AdamW.

Examples

>>> import mojolearn
>>> sched = mojolearn.WarmupCosineLR(3e-4, warmup_steps=200, total_steps=10000, min_lr=3e-5)
>>> lr = sched.lr_at(5000)

class mojolearn.training.Generator(seed, numeric_mode=None)[source]#

Bases: object

Seeded random stream for neural-network initialization and dropout.

Every value drawn is a function of (seed, stream id, element index) only, so the same call at the same counter gives the same bits on every device. Each draw uses the next stream id. state_dict() returns {'seed', 'counter'}, and load_state_dict() resumes the stream exactly.

Parameters:

seed (int) – Seed in [0, 2**64); otherwise ValueError.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Examples

>>> import mojolearn
>>> gen = mojolearn.Generator(0)
>>> w = gen.kaiming_uniform((128, 64), fan_in=64)
>>> y, key = gen.dropout(w, p=0.1)

state_dict()[source]#: Return the generator state as {'seed': int, 'counter': int}.

load_state_dict(state)[source]#

Restore a state from state_dict.

The counter must fit in 32 bits unsigned, otherwise ValueError.

next_stream()[source]#

Consume and return the next stream id.

Use it for operations that draw several slices of one stream, such as dropout over microbatches. Raises RuntimeError when the 32-bit counter is exhausted.

uniform(shape, low=0.0, high=1.0)[source]#

Draw float32 values uniform in [low, high).

Parameters:

shape (int or tuple of int) – Output shape; every dimension must be positive.
low (float, default 0.0) – Lower bound.
high (float, default 1.0) – Upper bound; must be >= low.

normal(shape, mean=0.0, std=1.0)[source]#

Draw float32 normal values (Box-Muller).

Parameters:

shape (int or tuple of int) – Output shape.
mean (float, default 0.0) – Mean.
std (float, default 1.0) – Standard deviation; must be >= 0.

kaiming_uniform(shape, fan_in, a=None)[source]#

Kaiming-uniform initialization, as torch.nn.init.kaiming_uniform_.

The bound is sqrt(2 / (1 + a**2)) * sqrt(3 / fan_in). The default a = sqrt(5) is torch’s nn.Linear default and gives the bound 1 / sqrt(fan_in).

Parameters:

shape (int or tuple of int) – Output shape.
fan_in (int) – Fan-in of the layer.
a (float or None, default None) – Negative slope; None means sqrt(5).

xavier_uniform(shape, fan_in, fan_out, gain=1.0)[source]#: Xavier-uniform initialization with bound gain * sqrt(6 / (fan_in + fan_out)).

kaiming_normal(shape, fan_in, gain=None)[source]#: Kaiming-normal initialization with std gain / sqrt(fan_in); gain None means sqrt(2).

xavier_normal(shape, fan_in, fan_out, gain=1.0)[source]#: Xavier-normal initialization with std gain * sqrt(2 / (fan_in + fan_out)).

dropout(x, p, offset=0, stream=None)[source]#

Inverted dropout of a float32 array.

Keeps each element whose uniform draw is >= p and scales it by 1 / (1 - p). Returns (y, key); pass key to dropout_backward to replay the same mask.

Parameters:

x (array of float32) – Input.
p (float) – Drop probability in [0, 1); otherwise ValueError.
offset (int, default 0) – Element index of x[0] within the stream, so that a microbatch of a larger batch draws the same values it would draw unsplit.
stream (int or None, default None) – Stream id; None takes the next one.

Returns:

y (array of float32)
key (dict)

dropout_backward(dy, key)[source]#: Gradient of dropout: apply the mask and scale recorded in key to dy.

mojolearn.training.accumulate_grads(microbatch_grads, tokens, numeric_mode=None)[source]#

Combine microbatch gradients into one, on the device.

Adds the microbatch gradients by a balanced tree in ascending microbatch order (not a running sum), so the result does not depend on the device. When tokens is given, a split that accumulation_is_aligned rejects raises ValueError; pass tokens=None to combine without that check.

Parameters:

microbatch_grads (list) – One entry per microbatch, each a float32 array or a list of arrays with the same shapes across entries.
tokens (int or None) – Token count of the whole step, or None.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

array or list of arrays – The combined gradient, in the structure of one entry.

mojolearn.training.accumulation_is_aligned(tokens, accumulation_steps, numeric_mode=None)[source]#

Whether a split into microbatches gives exactly the unsplit gradient.

True when accumulation_steps microbatches over tokens tokens, combined by accumulate_grads, give the same weight-gradient bits as one unsplit batch.

Parameters:

tokens (int) – Token count of the whole step.
accumulation_steps (int) – Number of microbatches.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

bool

class mojolearn.training.SambaConfig(vocab, d_model, layers, n_heads=None, n_kv_heads=None, head_dim=None, intermediate=None, tie_embeddings=True, norm_eps=1e-05, dropout=0.0)[source]#

Bases: object

Shape of a SambaStack: vocabulary, width and the sequence of layer kinds.

Parameters:

vocab (int) – Vocabulary size, >= 2.
d_model (int) – Model dimension, >= 1. Must be a multiple of 32 when any layer is ‘mamba3’, and equal n_heads * head_dim when any layer is ‘attention’.
layers (sequence of {‘mamba3’, ‘attention’}) – Layer kinds in stack order; at least one.
n_heads (int or None, default None) – Attention heads. Required when any layer is ‘attention’.
n_kv_heads (int or None, default None) – Key/value heads for grouped-query attention. None uses n_heads.
head_dim (int or None, default None) – Attention head dimension. None uses d_model // n_heads.
intermediate (int or None, default None) – Hidden size of the attention layers’ gated MLP. Required when any layer is ‘attention’.
tie_embeddings (bool, default True) – Share the embedding matrix with the output head.
norm_eps (float, default 1e-5) – Epsilon of the final RMSNorm (>= 0). Each block keeps its own constants.
dropout (float, default 0.0) – Dropout rate after the embedding during SambaStack.train_step, in [0, 1).
Invalid values raise `ValueError`.

Examples

>>> import mojolearn
>>> cfg = mojolearn.SambaConfig(vocab=256, d_model=64, layers=("mamba3", "attention"),
...                             n_heads=4, intermediate=128)

to_dict()[source]#: The configuration as a plain dict (the inverse of from_dict).

classmethod from_dict(d)[source]#: Build a configuration from a dict made by to_dict.

block_shapes(kind)[source]#

Names and shapes of one block’s tensors.

Parameters:: kind ({‘mamba3’, ‘attention’}) – Layer kind.
Returns:: list of (str, tuple) – (name, shape) pairs in the block’s weight order.

registry()[source]#

Names and shapes of every parameter tensor of the stack, in order.

The order is embed.weight, then layers.{i}.<block tensor> for each layer, then norm_f.weight, then lm_head.weight when embeddings are not tied. It is also the summation order of gradient clipping.

Returns:: list of (str, tuple)

class mojolearn.training.SambaStack(config, weights=None, generator=None, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, lr_schedule=None, max_norm=None, accumulation_steps=1, numeric_mode=None)[source]#

Bases: object

A Samba-style language model with its own training loop.

Embedding, then a stack of Mamba3Block and TransformerBlock layers (each with its own pre-norm and residual), a final RMSNorm, a tied or untied LM head and the cross-entropy loss, with a full backward pass. Training uses AdamW over the parameter registry (SambaConfig.registry), optional gradient-norm clipping, an optional learning-rate schedule and gradient accumulation. All parameters live in one flat float32 buffer. Attention layers use full causal attention.

Pass exactly one of weights or generator to supply the initial values; otherwise ValueError.

With accumulation_steps > 1 the batch rows are split into microbatches whose gradients are combined in a fixed tree. A split the accumulation rule does not allow is refused with ValueError before any gradient is computed. An allowed split reproduces the unsplit step’s block, final-norm and LM-head gradients bit for bit, but not the embedding gradient’s, which is summed per token occurrence: after a step with accumulation the embedding can differ from the unsplit step in the last bits.

Parameters:

config (SambaConfig) – Model shape. Anything else raises TypeError.
weights (dict of str to float32 buffer or None, default None) – Initial parameters keyed by the registry names, with exact names and shapes.
generator (Generator or None, default None) – Seeded mojolearn.training.Generator used to initialize the parameters and to draw dropout masks.
lr (float, default 1e-3) – Learning rate, used when lr_schedule is None.
betas (tuple of float, default (0.9, 0.999)) – AdamW moment decays.
eps (float, default 1e-8) – AdamW denominator term.
weight_decay (float, default 0.01) – AdamW decoupled weight decay.
lr_schedule (schedule or None, default None) – For example WarmupCosineLR or WarmupLinearLR.
max_norm (float or None, default None) – Clip the global gradient L2 norm to this value; None turns clipping off.
accumulation_steps (int, default 1) – Number of microbatches per optimizer step; must divide the batch rows.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Variables:

config (SambaConfig)
names (list of str) – Registry names in order.
optimizer (AdamW) – The optimizer.

Examples

>>> import numpy as np, mojolearn
>>> from mojolearn.training import Generator, WarmupCosineLR
>>> cfg = mojolearn.SambaConfig(vocab=256, d_model=64, layers=("mamba3", "mamba3"))
>>> stack = mojolearn.SambaStack(cfg, generator=Generator(7), lr=1e-3,
...                              lr_schedule=WarmupCosineLR(1e-3, 8, 64),
...                              max_norm=1.0, accumulation_steps=2)
>>> ids = np.random.default_rng(0).integers(0, 256, (4, 32)).astype(np.int32)
>>> out = stack.train_step(ids[:, :-1], ids[:, 1:])
>>> stack.save_checkpoint("run.ckpt")

load_weights(weights)[source]#

Copy parameters in from a dict keyed by registry name.

The key set and every shape must match exactly; values must be float32.

parameters()[source]#: The parameters as {name: array}, views of the flat buffer in registry order.

forward(inputs, state=None, *, lengths=None)[source]#

Compute logits for a batch of token ids, without dropout.

Parameters:

inputs (int array, shape (B, L)) – Token ids.
state (SambaState or None, default None) – None runs every block from a zero state. Pass a state from allocate_state to carry it: each block’s state is read at entry and updated in place, so a later forward or step continues the sequence.
lengths (sequence of int or None, default None) – Makes the batch ragged: B integers in [1, L]. Row i is real at positions [0, lengths[i]) and padding after. Each real position’s logits equal the row run alone at its own length, and every padding position’s logits are exactly +0.0. Not allowed together with state.

Returns:

logits (mojolearn.Array, shape (B, L, vocab))

allocate_state(batch_size, max_tokens)[source]#

Return a zero decode state for batch_size sequences of up to max_tokens positions.

The state holds one block state per layer in stack order (Mamba3State or TransformerState), owned by the caller.

Parameters:

batch_size (int)
max_tokens (int)

Returns:

SambaState

step(inputs, state)[source]#

Decode one token per row, updating state in place.

The stateful forward at L = 1. No dropout.

Parameters:

inputs (int array, shape (B,) or (B, 1)) – Token ids.
state (SambaState) – Decode state from allocate_state.

Returns:

logits (mojolearn.Array, shape (B, vocab))

loss(inputs, targets)[source]#: Mean cross-entropy of targets given inputs, both (B, L) ids; no dropout, no gradient.

loss_and_grads(inputs, targets, num_items=None, dropout_stream=None, token_offset=0)[source]#

Run one forward and the full backward pass.

Parameters:

inputs, targets (int array, shape (B, L)) – Token ids and next-token targets. Shapes must match.
num_items (int or None, default None) – Divisor of the summed loss and gradients. None uses this call’s own number of targets; pass the whole step’s count to make a microbatch’s gradient its slice of the unsplit one.
dropout_stream (int or None, default None) – Generator stream for the dropout mask; None applies no dropout.
token_offset (int, default 0) – Position offset of this batch in the dropout stream.

Returns:

loss (float)
grads (list of mojolearn.Array) – One gradient per parameter, in registry order.

train_step(inputs, targets)[source]#

Run one optimizer step on a batch.

The batch rows are split into accumulation_steps microbatches; each one’s gradient uses the whole step’s divisor and the pieces are combined in a fixed tree before clipping and the AdamW update. B must be divisible by accumulation_steps, otherwise ValueError.

Parameters:: inputs, targets (int array, shape (B, L)) – Token ids and next-token targets.
Returns:: dict – loss, lr, step and total_norm (the gradient norm before clipping, or None when max_norm is None).

state_dict(*, _copy_arrays=True)[source]#

The complete training state as a dict.

Holds the configuration, parameter registry and values, optimizer moments and step, schedule, generator state and numeric mode, so load_state_dict resumes training exactly.

load_state_dict(state)[source]#

Restore a state made by state_dict; returns self.

Raises ValueError when the schema, configuration, numeric profile or sizes do not match this stack.

save_checkpoint(path)[source]#

Write the training state to path as a checksummed checkpoint.

The file is written atomically. Saving reads the live parameter and optimizer storage, so do not train or modify the stack while it saves.

classmethod from_checkpoint(path, numeric_mode=None)[source]#

Build a stack from a checkpoint written by save_checkpoint.

Parameters:

path (str or path-like)
numeric_mode (str or None, default None) – None uses the mode stored in the checkpoint.

Returns:

SambaStack
Raises ValueError when the checkpoint fails its integrity or schema checks.

mojolearn.training.embedding_forward(weight, ids, numeric_mode=None)[source]#

Embedding lookup: weight[ids].

Parameters:

weight (array, shape (V, D)) – float32 embedding table.
ids (array, shape (N,)) – Integer row indices.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

array, shape (N, D)

mojolearn.training.embedding_backward(dy, ids, vocab, numeric_mode=None)[source]#

Gradient of embedding_forward with respect to the table.

Returns a fresh (vocab, D) gradient from dy. Rows with equal ids are summed in a fixed ascending order (no atomics), so the result is reproducible.

Parameters:

dy (array, shape (N, D)) – float32 gradient of the output.
ids (array, shape (N,)) – The forward’s ids; the row counts must match, otherwise ValueError.
vocab (int) – Number of table rows V.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

array, shape (vocab, D)

mojolearn.training.rms_norm_forward(x, weight, eps, numeric_mode=None)[source]#

RMS normalization of each row: weight * x / sqrt(mean(x**2) + eps).

Parameters:

x (array, shape (M, D)) – float32 input.
weight (array, shape (D,)) – float32 scale.
eps (float) – Added to the mean square.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

array, shape (M, D)

mojolearn.training.rms_norm_backward(dy, x, weight, eps, numeric_mode=None)[source]#

Gradients (dx, dweight) of rms_norm_forward.

Parameters:

dy (array, shape (M, D)) – Gradient of the output.
x (array, shape (M, D)) – The forward input.
weight (array, shape (D,)) – The forward’s scale.
eps (float) – The forward’s epsilon.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

dx (array, shape (M, D))
dweight (array, shape (D,))

mojolearn.training.linear_forward(a, weight, numeric_mode=None)[source]#

Linear layer without bias: a @ weight.T, as torch.nn.Linear with bias=False.

Parameters:

a (array, shape (M, K)) – float32 input.
weight (array, shape (N, K)) – float32 weight, in torch’s layout.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

array, shape (M, N)

mojolearn.training.linear_backward(dc, a, weight, numeric_mode=None)[source]#

Gradients (da, dweight) of linear_forward.

Parameters:

dc (array, shape (M, N)) – Gradient of the output.
a (array, shape (M, K)) – The forward input.
weight (array, shape (N, K)) – The forward weight.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

da (array, shape (M, K))
dweight (array, shape (N, K))

mojolearn.training.chunked_lm_head_loss(hidden, weight, targets, return_grad=False, numeric_mode=None)[source]#

Mean cross-entropy of a language-model head without materializing the logits.

Computes the mean over rows of the cross-entropy of hidden @ weight.T against targets, and optionally its gradients, without holding the (rows, vocab) logits. The result does not depend on how the vocabulary is chunked internally. There is no bias, ignore index, label smoothing or other reduction; use linear_forward and cross_entropy when you need those. Non-finite inputs and out-of-range targets raise before any output is written.

Parameters:

hidden (array, shape (rows, width)) – float32 hidden states.
weight (array, shape (vocab, width)) – float32 output embedding, vocab >= 2.
targets (array, shape (rows,)) – Integer class indices in [0, vocab).
return_grad (bool, default False) – Also return the gradients.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

loss (float)
d_hidden (array, shape (rows, width)) – Only with return_grad=True.
d_weight (array, shape (vocab, width)) – Only with return_grad=True.

Examples

>>> import numpy as np, mojolearn
>>> rng = np.random.default_rng(0)
>>> hidden = rng.random((64, 32), dtype=np.float32)
>>> weight = rng.random((1000, 32), dtype=np.float32)
>>> targets = rng.integers(0, 1000, 64)
>>> loss, dh, dw = mojolearn.chunked_lm_head_loss(hidden, weight, targets, return_grad=True)

mojolearn.training.mlp_inference_sessions(inputs, weight1, bias1, weight2, bias2, numeric_mode='identical')[source]#

Run several independent batches of rows through one Linear-ReLU-Linear model.

Computes relu(x @ weight1.T + bias1) @ weight2.T + bias2 for every session and returns the logits of all sessions concatenated in the order given. Empty sessions are allowed. The model is stateless (no caches, no dropout).

Parameters:

inputs (list of arrays) – float32 sessions, each (rows_i, width).
weight1 (array, shape (hidden, width)) – First layer weight, in torch’s layout.
bias1 (array, shape (hidden,)) – First layer bias.
weight2 (array, shape (out_width, hidden)) – Second layer weight.
bias2 (array, shape (out_width,)) – Second layer bias.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’}, default ‘identical’) – Numeric mode.

Returns:

array, shape (sum of rows_i, out_width)

Raises:

ValueError – When the weight, bias and session shapes do not agree.

mojolearn.training.residual_dropout(values, residual, *, p=0.5, seed=0, stream=0, offset=0)[source]#

Return residual + dropout(values) as one explicit IDENTICAL layer.

Inputs must have identical shapes and C-contiguous float32 storage. The native implementation admits finite values and p in [0,1), computes scale=1/(1-p), and preserves the rounded dropout intermediate before add. RNG coordinates are (64-bit seed, 32-bit stream, offset+flat_index). NN59 selects the fused schedule; OFF/ALL_OFF use two native stages. No existing model is implicitly changed and no validation is claimed.

mojolearn.training.residual_dropout_backward(gradient, *, p=0.5, seed=0, stream=0, offset=0)[source]#

Return (d_values, d_residual) for the explicit residual layer.

Pass exactly the forward p/seed/stream/offset to regenerate its mask. d_values applies the same rounded keep-and-scale operation to gradient; d_residual preserves the upstream gradient words. This is an explicit backward operation, not an automatic differentiation registration.