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, dense=True, distributed=True, precision="bf16") # packed DDP/FSDP2 training
text = lm.generate(prompt_ids, n=200, temperature=0.8)            # autoregressive sampling
nll  = lm.nll(held_out_ids)                                       # nats/token on held-out data

fit retains the non-buffering streaming estimator for compatibility. Packed dense=True training uses the distributed-gradient backend when requested, with deterministic data sharding, microbatches, accumulation, and complete DCP training state. fit_pairs is the SFT stage: dense all-position teacher forcing on (prompt, completion) pairs with the loss masked to completions. The rest of the multi-stage pipeline (CPT-with-EWC, DPO) is mixle.models.continual / mixle.models.dpo_leaf.

class LMFitReceipt(objective, input_tokens, consumed_tokens_per_epoch, supervised_tokens_per_epoch, discarded_cross_row_transitions_per_epoch, dropped_tail_tokens, epochs)[source]

Bases: object

Token-accounting receipt for one packed dense corpus layout.

Parameters:

  • objective (str)

  • input_tokens (int)

  • consumed_tokens_per_epoch (int)

  • supervised_tokens_per_epoch (int)

  • discarded_cross_row_transitions_per_epoch (int)

  • dropped_tail_tokens (int)

  • epochs (int)

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 LM’s inference artifact to path via torch.save: architecture config (vocab/d_model/n_layer/n_head/block) + learned weights, nothing else.

This is not a training checkpoint. It deliberately does not store the optimizer, LR scheduler, step count, RNG state, data-loader position, or gradient-scaler state, so load() returns a model ready for scoring and generation, not one you can resume training from without loss of state. For resumable training use the checkpoint path in mixle.utils.parallel.fault_tolerant_training / mixle.utils.parallel.dcp_checkpoint, which stores model + optimizer + data-loader position.

Parameters:

path (str)

Return type:

None

classmethod load(path)[source]

Load an LM inference artifact written by save() (config + weights).

Returns a model ready for scoring/generation; this is not a training-resume checkpoint (see save()), so optimizer/step/RNG/data-loader state is not restored.

Loaded with weights_only=True: to_dict()’s payload is plain config scalars plus a tensor state_dict, never an arbitrary object graph, so this does not execute code from the file (unlike a full-module pickle – see mixle.models._neural_serial for that case).

Parameters:

path (str)

Return type:

LM

fit(token_ids, *, epochs=1, batch_size=64, lr=3e-3, distributed=False, precision='fp32', shuffle=True, dense=False, seed=0, log=None, tp_size=1, pp_size=1, cp_size=1, dp_replicate=1, dp_shard=1, ep_size=1, etp_size=1, microbatches=1, gradient_accumulation_steps=1, distributed_backend='torch_native', optimizer=None, max_grad_norm=None, compile=False, checkpoint_path=None, resume=False)[source]

Small-scale reference pretraining (or continuation) on a token-id array.

Two paths, chosen by dense:

  • dense=False (default, unchanged): the streaming estimator; the corpus is never buffered. It scores one next-token target per block-length window – the right shape for an unbounded stream, but only a fraction (~``1/block``) of the token supervision that packed dense teacher forcing extracts from the same tokens. A reference / continuation path for small-scale runs, not a frontier pretraining engine.

  • dense=True: packed dense teacher forcing on a BOUNDED corpus. The token stream is cut into non-overlapping rows of block + 1 tokens (a partial tail row is dropped; there is no padding, so no attention-mask or EOS special-casing is needed – document boundaries are whatever the ids encode). Every row consumes block+1 tokens but contributes block shifted targets: the corpus’s first token and each cross-row transition are not supervised, and the partial tail is not consumed. last_fit_receipt reports those counts explicitly. Token ids stay integer end to end (no float conversion). seed drives row shuffling and log(epoch, mean_loss) reports per-epoch training loss, both as in fit_pairs(). With distributed=True this objective is executed by the selected distributed-gradient backend.

Parallel sizes are executable contracts. torch_native supports DDP/HSDP, FSDP2, MLP tensor parallelism, and CUDA context parallelism; unsupported combinations fail capability validation. Full TP/PP/CP/EP transformer and MoE training is exposed by the megatron backend.

Parameters:

  • token_ids (Any)

  • epochs (int)

  • batch_size (int)

  • lr (float)

  • distributed (bool)

  • precision (str)

  • shuffle (bool)

  • dense (bool)

  • seed (int)

  • log (Any)

  • tp_size (int)

  • pp_size (int)

  • cp_size (int)

  • dp_replicate (int)

  • dp_shard (int)

  • ep_size (int)

  • etp_size (int)

  • microbatches (int)

  • gradient_accumulation_steps (int)

  • distributed_backend (str)

  • optimizer (Any)

  • max_grad_norm (float | None)

  • compile (bool)

  • checkpoint_path (str | None)

  • resume (bool)

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