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.
- class CausalAttention(d_model, n_head)[source]
Bases:
Module- forward(x)[source]
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class Block(d_model, n_head)[source]
Bases:
Module- forward(x)[source]
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class CausalLM(vocab, d_model, n_layer, n_head, block, embedding=None)[source]
Bases:
Module- forward(x)[source]
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- 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).