mojolearn.SambaStack

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

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

__init__(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)

Methods

allocate_state(batch_size, max_tokens)	Return a zero decode state for batch_size sequences of up to max_tokens positions.
forward(inputs[, state, lengths])	Compute logits for a batch of token ids, without dropout.
from_checkpoint(path[, numeric_mode])	Build a stack from a checkpoint written by save_checkpoint.
load_state_dict(state)	Restore a state made by state_dict; returns self.
load_weights(weights)	Copy parameters in from a dict keyed by registry name.
loss(inputs, targets)	Mean cross-entropy of targets given inputs, both (B, L) ids; no dropout, no gradient.
loss_and_grads(inputs, targets[, num_items, ...])	Run one forward and the full backward pass.
parameters()	The parameters as {name: array}, views of the flat buffer in registry order.
save_checkpoint(path)	Write the training state to path as a checksummed checkpoint.
state_dict(*[, _copy_arrays])	The complete training state as a dict.
step(inputs, state)	Decode one token per row, updating state in place.
train_step(inputs, targets)	Run one optimizer step on a batch.

load_weights(weights)

Copy parameters in from a dict keyed by registry name.

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

parameters()

The parameters as {name: array}, views of the flat buffer in registry order.

forward(inputs, state=None, *, lengths=None)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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.