mojolearn.language_model#

Train and run FP32 decoder language models.

LanguageModelConfig: the model and batch dimensions.
LanguageModelTrainer: AdamW training on the GPU. By default the training state stays on the device between steps (resident=True); train_step returns the full result dict, or with step_result="lean" only the loss, step counters and flags. export_state(), export_gradients() and export_checkpoint() copy the state out on demand.
LanguageModelInference: the same model’s forward pass on the CPU, with no GPU (LanguageModelInference.from_checkpoint(path)).

SmallByteLanguageModelTrainer and ByteLanguageModelConfig are the same classes under their earlier names. To train on a BPE-tokenized corpus, see mojolearn.lm_corpus.

Examples

>>> from mojolearn.language_model import LanguageModelConfig, LanguageModelTrainer
>>> shape = LanguageModelConfig(vocab_size=256)

Classes

`LanguageModelConfig`
`LanguageModelTrainer`
`LanguageModelInference`(parameters, *[, ...])	Run a trained decoder language model on the CPU, with no GPU.
`SmallByteLanguageModelTrainer`(parameters, *, ...)	FP32 decoder language model trainer (AdamW, mean cross-entropy).
`ByteLanguageModelConfig`([batch, length, ...])	Dimensions of the decoder language model: batch, sequence, layers, vocabulary.

mojolearn.language_model.LanguageModelConfig#: alias of ByteLanguageModelConfig

mojolearn.language_model.LanguageModelTrainer#: alias of SmallByteLanguageModelTrainer

class mojolearn.language_model.LanguageModelInference(parameters, *, shape=None, threaded=True, threads=None)[source]#

Bases: object

Run a trained decoder language model on the CPU, with no GPU.

Forward pass only: logits, loss and greedy next-token prediction for a model trained with LanguageModelTrainer, giving the same bits as the GPU forward pass. The parameters are copied at construction and never change.

By default (threaded=True) calls use SIMD and up to threads threads (None: one per physical core); threaded=False runs single-threaded reference code. Both give the same bits. Each call can override both settings.

Parameters:

parameters (array-like or dict) – Flat float32 parameters of shape.n_total values (finite), or a dict of named tensors as listed by LanguageModelTrainer.parameter_registry(); packed low-bit tensors from mojolearn.lowbit are expanded exactly.
shape (LanguageModelConfig or None, default None) – Model dimensions; None uses the defaults.
threaded (bool, default True) – Use the multithreaded path.
threads (int or None, default None) – Thread count; None means one per physical core.

Variables:

weight_format (str) – "float32", or the packed format of the parameters given.

Examples

>>> import numpy as np, mojolearn
>>> lm = mojolearn.LanguageModelInference.from_checkpoint("model.json")
>>> ids = np.frombuffer(b"Hello, wor", dtype=np.uint8).astype(np.int32).reshape(1, -1)
>>> next_ids = lm.next_bytes(ids)

classmethod from_checkpoint(path, *, threaded=True, threads=None)[source]#

Load the parameters of a JSON checkpoint written by the trainer.

Reads a file written by LanguageModelTrainer.export_checkpoint(), with the trainer’s own integrity checks.

Parameters:

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

Returns:

LanguageModelInference

property shape#: The model dimensions, a LanguageModelConfig.

property profile#: The name of the arithmetic profile the CPU binding computes.

parameters_sha256()[source]#: SHA-256 of the parameters’ little-endian float32 bytes.

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

Float32 logits for every position.

Parameters:

ids (array-like of shape (batch, length), int32) – Token ids, starting at position 0; length at most shape.length.
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.
threaded (bool or None, default None) – Override the instance setting.
threads (int or None, default None) – Override the instance setting.

Returns:

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

loss_bits(ids, *, threaded=None, threads=None)[source]#

IEEE-754 bit pattern of the mean next-token loss.

Parameters:

ids (array-like of shape (shape.batch, shape.length + 1), int32) – The training batch layout. The last column is only a target; -100 there is ignored.
threaded (bool or None, default None) – Override the instance setting.
threads (int or None, default None) – Override the instance setting.

Returns:

int – The float32 loss’s bits, for exact comparison.

loss(ids, *, threaded=None, threads=None)[source]#: Mean next-token loss as a float; see loss_bits().

next_bytes(ids, *, lengths=None, threaded=None, threads=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.
lengths (sequence of int or None, default None) – Ragged row lengths, as for logits().
threaded (bool or None, default None) – Override the instance setting.
threads (int or None, default None) – Override the instance setting.

Returns:

list of int – One token id per row.

class mojolearn.language_model.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"]

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

class mojolearn.language_model.ByteLanguageModelConfig(batch=2, length=32, d_model=32, n_heads=4, n_kv=2, head_dim=8, intermediate=64, n_layers=2, vocab_size=256, chunked_lm_head_v2=False)[source]#

Bases: object

Dimensions of the decoder language model: batch, sequence, layers, vocabulary.

An immutable description of the model and training-batch shape passed to LanguageModelTrainer as shape=. Each layer has an attention block with grouped key/value heads and a gated MLP, each with its own norm weight; the model has a token embedding and a separate output head. The defaults describe the small published byte-level model (vocabulary 256). LanguageModelConfig is the same class.

Parameters:

batch (int, default 2) – Sequences per training step.
length (int, default 32) – Tokens per sequence, at most 8192.
d_model (int, default 32) – Model width; must equal n_heads * head_dim.
n_heads (int, default 4) – Attention heads; must be divisible by n_kv.
n_kv (int, default 2) – Key/value heads.
head_dim (int, default 8) – Width of each head; must be even.
intermediate (int, default 64) – MLP hidden width.
n_layers (int, default 2) – Number of decoder layers.
vocab_size (int, default 256) – Token vocabulary size (256 for bytes; corpus.n_vocab for a BPE corpus from mojolearn.lm_corpus).
chunked_lm_head_v2 (bool, default False) – Compute the loss and its gradient over bounded vocabulary panels, which bounds memory for large vocabularies. Logits from prediction still cover the full vocabulary.

Raises:

ValueError – If a dimension is not an integer in [1, 2**20], the head constraints above are not met, or the shape exceeds signed 32-bit indexing.

Parameters:

batch (int)
length (int)
d_model (int)
n_heads (int)
n_kv (int)
head_dim (int)
intermediate (int)
n_layers (int)
vocab_size (int)
chunked_lm_head_v2 (bool)

Examples

>>> import mojolearn
>>> shape = mojolearn.LanguageModelConfig(batch=1, length=128, d_model=64, n_heads=4,
...                                       n_kv=2, head_dim=16, intermediate=128,
...                                       n_layers=2, vocab_size=256)
>>> shape.n_total > 0
True

property native_shape#: The dimensions as a tuple of ints, in field order.

property profile#: The name of the training profile this shape selects.

property parameter_shapes#: The shape of every parameter tensor, in packing order.

property n_tensors#

2 + 9 * n_layers.

Type:: Number of parameter tensors

property parameter_names#: The name of every parameter tensor, in packing order.

property offsets#: Start offset of every parameter tensor in the flat parameter vector, plus the total.

property n_total#: Total number of parameters.

to_dict()[source]#: The fields as a dict, suitable for ByteLanguageModelConfig(**d).