mixle.models.language_model module

LM – a declarative autoregressive language model tying the frontier-LLM stack into one usable object.

A causal Transformer trained on a token stream:

lm = LM(vocab=V, d_model=256, n_layer=6, n_head=8, block=128)
lm.fit(token_ids, epochs=3, batch_size=64, device="mps")          # pretrain (single process)
lm.fit(token_ids, distributed=True, precision="bf16")             # or distributed under torchrun (FSDP2 on CUDA)
text = lm.generate(prompt_ids, n=200, temperature=0.8)            # autoregressive sampling
nll  = lm.nll(held_out_ids)                                       # bits/token on held-out data

fit runs the non-buffering streaming estimator (mixle.models.streaming_transformer_leaf); with distributed=True it dispatches through StreamingTokenEncodedData (per-rank shard, in-backward all-reduce, FSDP2/ZeRO-3 + bf16 + DCP on CUDA). fit_pairs is the SFT stage: dense all-position teacher forcing on (prompt, completion) pairs with the loss masked to completions – the shape needed to distill verified trajectories (e.g. execution-checked page -> parser code pairs from a data factory) into a tiny model. The rest of the multi-stage pipeline (CPT-with-EWC, DPO) is mixle.models.continual / mixle.models.dpo_leaf.

class LM(vocab, *, d_model=256, n_layer=6, n_head=8, block=128, device='cpu', embedding=None)[source]

Bases: object

A causal-Transformer language model with a small declarative surface: fit / generate / nll.

Parameters:

  • vocab (int)

  • d_model (int)

  • n_layer (int)

  • n_head (int)

  • block (int)

  • device (str)

  • embedding (Any)

to_dict()[source]

Serialize the hyperparameters + trained weights so the LM survives a process boundary.

The token embedding may be tied across LMs (embedding=); from_dict rebuilds an untied module and loads the saved state_dict into it, so a round-tripped LM is standalone (any external tie is dropped).

Return type:

dict

classmethod from_dict(payload)[source]

Rebuild an LM from to_dict() output (fresh module, saved weights loaded in).

Parameters:

payload (dict)

Return type:

LM

save(path)[source]

Persist the trained LM to path via torch.save (hyperparameters + weights).

Parameters:

path (str)

Return type:

None

classmethod load(path)[source]

Load an LM previously written by save().

Parameters:

path (str)

Return type:

LM

fit(token_ids, *, epochs=1, batch_size=64, lr=3e-3, distributed=False, precision='fp32', shuffle=True)[source]

Pretrain (or continue) on a token-id array via the streaming estimator; the corpus is never buffered.

Parameters:
Return type:

LM

fit_pairs(pairs, *, epochs=1, batch_size=32, lr=3e-3, mask_prompt=True, pad_id=0, seed=0, log=None)[source]

Supervised fine-tuning on (prompt_ids, completion_ids) pairs with a dense per-position loss.

The streaming fit path scores ONE next-token target per window (the right shape for an unbounded pretraining stream); for a pair corpus that wastes a factor of block in compute. Here every position of every pair contributes cross-entropy in a single forward, and mask_prompt restricts the loss to completion positions – the standard SFT objective. Sequences longer than block keep the completion and drop the oldest prompt tokens; shorter ones are left-padded with pad_id (excluded from the loss). Include your end-of-sequence token in each completion so generate(stop_id=...) knows where to stop.

Parameters:
Return type:

LM

generate(prompt_ids, n=200, *, temperature=1.0, greedy=False, seed=0, stop_id=None)[source]

Autoregressively extend prompt_ids by n tokens (greedy, or temperature-sampled).

stop_id ends generation early when that token is produced (it is included in the return value, so callers can strip it – and its presence distinguishes ‘finished’ from ‘ran out of budget’).

Parameters:
Return type:

list

nll(token_ids)[source]

Mean next-token negative log-likelihood (nats/token) on a token-id array.

Parameters:

token_ids (Any)

Return type:

float