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 (the FlashAttention dispatch on CUDA). At frontier scale the same module is what a vendored TorchTitan/Megatron trainer shards (FSDP2/TP/PP); here it runs single-process.

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

Parameters:
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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Any)

Return type:

Any

class Block(d_model, n_head)[source]

Bases: Module

Parameters:
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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Any)

Return type:

Any

class CausalLM(vocab, d_model, n_layer, n_head, block, embedding=None)[source]

Bases: Module

Parameters:

  • vocab (int)

  • d_model (int)

  • n_layer (int)

  • n_head (int)

  • block (int)

  • embedding (Any)

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Any)

Return type:

Any

build_causal_lm(vocab, d_model=128, n_layer=3, n_head=4, block=64, embedding=None)[source]

Build a causal decoder-only Transformer LM (token+pos embeddings, pre-norm blocks, weight-tied head).

embedding optionally injects a shared token nn.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).

Parameters:
Return type:

Any