mixle.models.transformer module¶
A causal decoder-only Transformer as a torch module – the engine behind the declarative AR-LM surface.
Built lazily by the mixle.ppl.core.Transformer predictor token. forward(x) takes a (batch, block)
context of token ids (accepted as float so it rides the SoftmaxNeuralLeaf float path, cast to long inside)
and returns next-token logits (batch, vocab) from the last position. So
Categorical(logits=Transformer(out=V)) is exactly next-token prediction p(token | context), fit by the
standard estimate() loop whose cross-entropy is -log p – no new training machinery.
Attention is F.scaled_dot_product_attention with CUDA FlashAttention dispatch when available. This module runs
single-process here, while larger training stacks can shard the same architecture externally.
The nn.Module subclasses are defined at MODULE level (not nested inside build_causal_lm) so a trained
LM pickles/saves: a function-local class has no importable qualname and torch.save/pickle cannot find it.
- build_causal_lm(vocab, d_model=128, n_layer=3, n_head=4, block=64, embedding=None, gradient_checkpointing=False)[source]
Build a causal decoder-only Transformer LM (token+pos embeddings, pre-norm blocks, weight-tied head).
embeddingoptionally injects a shared tokennn.Embedding(vocab x d_model) to use in place of a fresh one – so several language models can tie the same word embedding and train it jointly (the weight-tied head follows it). Its shape must match(vocab, d_model).gradient_checkpointing=Truerecomputes block activations during backward instead of storing them – identical gradients (pinned by test) for a large activation-memory cut on deep stacks or long blocks. The flag is a plain module attribute, so it can also be toggled on an existing model – including to a per-block list/tuple of bools (one pern_layer) rather than a single all-or-nothing bool, for F6’s cost-model-driven selective policy (mixle.models.memory_efficient_training.SelectiveRecomputePolicy).