mojolearn.SmallByteLanguageModelTrainer#

class mojolearn.SmallByteLanguageModelTrainer(parameters, *, data_schedule, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, shape=None, resident=None, step_result=None)[source]#

Bases: object

FP32 decoder language model trainer (AdamW, mean cross-entropy).

Trains a decoder-only language model whose dimensions come from LanguageModelConfig (shape=); the defaults are a small byte-level model (batch 2, length 32, width 32, 4 heads, 2 key/value heads, MLP 64, vocabulary 256). The model has no biases, no dropout and no final norm, and the embedding and output head are separate (untied). LanguageModelTrainer is the same class.

You supply everything: the initial parameters (there is no hidden initialization), the token ids of every step (there is no tokenizer, padding, truncation or random number generation inside), and a data_schedule describing them. All arithmetic runs on the GPU in the process numeric mode: "identical" (the default) gives the same bits on NVIDIA and AMD.

Resident sessions. On a GPU install resident defaults to True: the parameters, AdamW moments and last gradient live on the device between steps, and only the ids go in and the loss comes out. They are copied to the host only when you ask (export_state(), export_gradients(), export_checkpoint(), close()). A step that fails is rolled back. If the device context is lost, the steps since the last export are lost and every call raises RuntimeError until load_state_dict(). resident=False runs each step statelessly with the state on the host.

The data_schedule is a small JSON-serializable dict you provide (for example from mojolearn.lm_corpus.TokenBatches.data_schedule()): record the corpus, token order, batch offsets and planned steps there. It is checked only for size and JSON round-trip, not against the ids you pass, but once a checkpoint is written it is covered by the checkpoint’s SHA-256, so it cannot be changed undetected. next_batch_index equals the completed step count; after a restore, pass the matching next batch.

Parameters:

parameters (array-like or dict) – Initial parameters: one flat float32 array of shape.n_total values, or a dict of named tensors as listed by parameter_registry().
data_schedule (dict) – The JSON descriptor described above. When it records a vocabulary, its n_vocab must equal shape.vocab_size.
lr (float, default 1e-3) – AdamW learning rate.
betas (tuple of float, default (0.9, 0.999)) – AdamW moment decay rates.
eps (float, default 1e-8) – AdamW epsilon.
weight_decay (float, default 0.01) – AdamW decoupled weight decay.
shape (LanguageModelConfig or None, default None) – Model dimensions; None uses the defaults.
resident (bool or None, default None) – Keep the training state on the device between steps. None chooses the resident session whenever the loaded GPU binaries support it.
step_result ({“full”, “lean”} or None, default None) – What train_step() returns. "full" includes the gradients; "lean" (requires resident=True) returns the loss, step counters and flags and leaves the gradient on the device. None means "lean" when resident=True is given and "full" otherwise.

Variables:

step_ (int) – Completed training steps.
parameters_ (Array) – A copy of the current flat parameters.
data_schedule (dict) – A copy of the data_schedule descriptor.

Examples

>>> import numpy as np, mojolearn
>>> shape = mojolearn.LanguageModelConfig()
>>> rng = np.random.default_rng(0)
>>> weights = (rng.standard_normal(shape.n_total) * 0.02).astype(np.float32)
>>> trainer = mojolearn.LanguageModelTrainer(weights, data_schedule={"corpus": "demo"})
>>> ids = rng.integers(0, 256, (shape.batch, shape.length + 1), dtype=np.int32)
>>> loss = trainer.train_step(ids)["loss"]

__init__(parameters, *, data_schedule, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, shape=None, resident=None, step_result=None)[source]#

Methods

`attention_stage_report`()	Memory held by the open session's attention stages.
`close`()	Copy the device state to the host and release the device.
`evaluate`(ids)	Mean loss on `ids` without changing any state.
`export_checkpoint`(path)	Write the complete state to `path` as a JSON checkpoint.
`export_checkpoint_binary`(path)	Write the complete state to `path` as a binary checkpoint.
`export_gradients`([named])	The pre-update gradient of the last completed step of a resident session.
`export_state`()	The complete training state as fresh copies.
`from_checkpoint`(path, *[, resident])	Restore a trainer from a JSON checkpoint written by export_checkpoint().
`from_checkpoint_binary`(path, *[, resident])	Restore a trainer from export_checkpoint_binary().
`from_checkpoint_bytes`(encoded, *[, resident])	Restore a trainer from the bytes of a JSON checkpoint.
`load_state_dict`(state)	Replace the complete training state.
`logits`(ids, *[, lengths])	Float32 logits for every position, without changing any state.
`next_bytes`(ids, *[, lengths])	Greedy next token after each row.
`parameter_registry`([shape])	Name, shape, offset and size of every parameter tensor.
`run_metadata`()	Profile, configuration, schedule and binary details of this run.
`save_checkpoint`(path)	Same as export_checkpoint().
`state_dict`()	The complete training state; same as export_state().
`train_step`(ids)	Run one training step: mean cross-entropy loss and an AdamW update.

Attributes

`data_schedule`	A copy of the `data_schedule` descriptor, as checkpoints keep it.
`parameters_`	A copy of the current flat float32 parameters.
`step_`	Number of completed training steps.

property data_schedule#: A copy of the data_schedule descriptor, as checkpoints keep it.

static parameter_registry(shape=None)[source]#

Name, shape, offset and size of every parameter tensor.

Parameters:: shape (LanguageModelConfig or None, default None) – Model dimensions; None uses the defaults.
Returns:: list of dict – One {"name", "shape", "offset", "size"} record per tensor, in the order of the flat parameter vector.

property step_#: Number of completed training steps.

property parameters_#: A copy of the current flat float32 parameters.

state_dict()[source]#: The complete training state; same as export_state().

export_state()[source]#

The complete training state as fresh copies.

On a resident session this downloads the device state once and validates it. Changing the returned arrays does not affect the trainer.

Returns:: dict – Parameters, AdamW moments m and v, flags, step counters, configuration, data_schedule and metadata, as accepted by load_state_dict().

export_gradients(named=True)[source]#

The pre-update gradient of the last completed step of a resident session.

Valid until the next train_step() starts. The stateless path returns its gradients from train_step() instead.

Parameters:: named (bool, default True) – Also return the gradient split into named tensors.
Returns:: dict – {"flat_gradients": Array} and, when named, {"gradients": {name: Array}}.
Raises:: RuntimeError – If no resident session is open, or no step has completed since it was opened, restored or rolled back.

attention_stage_report()[source]#

Memory held by the open session’s attention stages.

The attention work arrays start small and grow the first time a layer needs the unfused attention path, which depends on the data. Read this between steps to see the current device footprint of those arrays. Nothing is launched or downloaded.

Returns:: dict or None – Cell and byte counts per stage (forward_eager_cells, forward_aexp_cells, backward_eager_cells and their *_bytes, eager_bytes, total_bytes), layer counts and per-layer status. None when no resident session is open or the loaded binaries do not report it.

close()[source]#

Copy the device state to the host and release the device.

A later call re-opens the session from that host state. On a lost session nothing can be exported: the device is released and the trainer stays lost until load_state_dict().

load_state_dict(state)[source]#

Replace the complete training state.

The replacement is validated before anything changes; a refused state leaves the trainer untouched. Any resident session is released without export and the next call opens one from the replacement.

Parameters:: state (dict) – A state from export_state() (or state_dict()).

run_metadata()[source]#

Profile, configuration, schedule and binary details of this run.

Reads constants only; runs no model computation. Keep it with any results you record.

Returns:: dict

train_step(ids)[source]#

Run one training step: mean cross-entropy loss and an AdamW update.

Parameters:: ids (array-like of shape (batch, length + 1), int32) – Token ids in [0, shape.vocab_size). The first length positions predict the next length; the loss averages batch * length targets.
Returns:: dict – loss, step, completed_steps and next_batch_index. With step_result="full" also flat_gradients and gradients (by tensor name), the pre-update gradient; with "lean" also flags.

evaluate(ids)[source]#

Mean loss on ids without changing any state.

Parameters:: ids (array-like of shape (batch, length + 1), int32) – Token ids, as for train_step().
Returns:: float

logits(ids, *, lengths=None)[source]#

Float32 logits for every position, without changing any state.

Runs the same GPU forward pass the training loss uses, starting at position 0 of each row.

Parameters:

ids (array-like of shape (batch, length), int32) – Token ids in [0, shape.vocab_size). batch in [1, 1024], length in [1, shape.length], and batch * length * vocab_size at most 268435456.
lengths (sequence of int or None, default None) – Makes the batch ragged: row i is real at positions [0, lengths[i]) and padding after, whatever the padding ids are. Real positions get exactly the logits of the row alone at its own length; padding positions are +0.0.

Returns:

Array of shape (batch, length, vocab_size), float32

Raises:

ValueError – For ids or lengths outside the limits above.
RuntimeError – If a logit is not finite.

next_bytes(ids, *, lengths=None)[source]#

Greedy next token after each row.

Picks the highest logit at each row’s last position (with lengths, its last real position), ties going to the lowest id.

Parameters:

ids (array-like of shape (batch, length), int32) – Token ids, as for logits().
lengths (sequence of int or None, default None) – Ragged row lengths, as for logits().

Returns:

list of int – One token id per row.

save_checkpoint(path)[source]#: Same as export_checkpoint().

export_checkpoint(path)[source]#

Write the complete state to path as a JSON checkpoint.

The file is written atomically, holds no pickle, and is limited to 2 MiB (about 87,000 parameters). For larger models use export_checkpoint_binary().

Parameters:: path (str or path-like) – Destination file.

export_checkpoint_binary(path)[source]#

Write the complete state to path as a binary checkpoint.

For models of any size. The parameters and the AdamW moments are stored as raw little-endian bytes after a JSON header, so a round trip is bit-exact (-0.0, NaN payloads and subnormals survive) and a resumed run continues exactly. The file is written atomically.

Parameters:: path (str or path-like) – Destination file.
Returns:: str – The file’s SHA-256.

classmethod from_checkpoint_binary(path, *, resident=None)[source]#

Restore a trainer from export_checkpoint_binary().

The whole file is verified (header and per-array digests, exact size) before anything is allocated, and the state is validated as load_state_dict() validates it.

Parameters:

path (str or path-like) – The checkpoint file.
resident (bool or None, default None) – As for the constructor.

Returns:

LanguageModelTrainer

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

Restore a trainer from a JSON checkpoint written by export_checkpoint().

Parameters:

path (str or path-like) – The checkpoint file.
resident (bool or None, default None) – As for the constructor.

Returns:

LanguageModelTrainer

classmethod from_checkpoint_bytes(encoded, *, resident=None)[source]#

Restore a trainer from the bytes of a JSON checkpoint.

Runs the same checks as from_checkpoint() without opening a file.

Parameters:

encoded (bytes) – The checkpoint contents (exactly bytes; copy mutable buffers first).
resident (bool or None, default None) – As for the constructor.

Returns:

LanguageModelTrainer