mixle.models.language_model module

LM – a declarative autoregressive language model with fit, generation, and scoring helpers.

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 into a compact 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, tp_size=1, pp_size=1, cp_size=1)[source]

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

tp_size/pp_size/cp_size (only meaningful with distributed=True) are the F1 N-D parallelism dimensions – tensor/pipeline/context parallel, ORTHOGONAL to the data-parallel axis (DDP on CPU, FSDP2/ZeRO-3 on CUDA) this handle already runs. They default to 1 (off): the plan is validated against the model’s real dimensions (n_head % tp_size, pp_size <= n_layer, block % cp_size) so a bad plan fails fast, using the sharding/reconstruction mechanism in mixle.utils.parallel.tensor_pipeline_context_parallel (tested there at small scale against the dense forward). Composing the validated plan into real per-axis multi-GPU process groups is the piece that needs actual hardware this environment does not have – see that module’s docstring.

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