mojolearn.SambaConfig

class mojolearn.SambaConfig(vocab, d_model, layers, n_heads=None, n_kv_heads=None, head_dim=None, intermediate=None, tie_embeddings=True, norm_eps=1e-05, dropout=0.0)

Bases: object

Shape of a SambaStack: vocabulary, width and the sequence of layer kinds.

Parameters:

• vocab (int) – Vocabulary size, >= 2.

• d_model (int) – Model dimension, >= 1. Must be a multiple of 32 when any layer is ‘mamba3’, and equal n_heads * head_dim when any layer is ‘attention’.

• layers (sequence of {‘mamba3’, ‘attention’}) – Layer kinds in stack order; at least one.

• n_heads (int or None, default None) – Attention heads. Required when any layer is ‘attention’.

• n_kv_heads (int or None, default None) – Key/value heads for grouped-query attention. None uses n_heads.

• head_dim (int or None, default None) – Attention head dimension. None uses d_model // n_heads.

• intermediate (int or None, default None) – Hidden size of the attention layers’ gated MLP. Required when any layer is ‘attention’.

• tie_embeddings (bool, default True) – Share the embedding matrix with the output head.

• norm_eps (float, default 1e-5) – Epsilon of the final RMSNorm (>= 0). Each block keeps its own constants.

• dropout (float, default 0.0) – Dropout rate after the embedding during SambaStack.train_step, in [0, 1).

• Invalid values raise `ValueError`.

Examples

>>> import mojolearn
>>> cfg = mojolearn.SambaConfig(vocab=256, d_model=64, layers=("mamba3", "attention"),
...                             n_heads=4, intermediate=128)

__init__(vocab, d_model, layers, n_heads=None, n_kv_heads=None, head_dim=None, intermediate=None, tie_embeddings=True, norm_eps=1e-05, dropout=0.0)

Methods

block_shapes(kind)	Names and shapes of one block's tensors.
from_dict(d)	Build a configuration from a dict made by to_dict.
registry()	Names and shapes of every parameter tensor of the stack, in order.
to_dict()	The configuration as a plain dict (the inverse of from_dict).

to_dict()

The configuration as a plain dict (the inverse of from_dict).

classmethod from_dict(d)

Build a configuration from a dict made by to_dict.

block_shapes(kind)

Names and shapes of one block’s tensors.

Parameters:

kind ({‘mamba3’, ‘attention’}) – Layer kind.

Returns:

list of (str, tuple) – (name, shape) pairs in the block’s weight order.

registry()

Names and shapes of every parameter tensor of the stack, in order.

The order is embed.weight, then layers.{i}.<block tensor> for each layer, then norm_f.weight, then lm_head.weight when embeddings are not tied. It is also the summation order of gradient clipping.

Returns:

list of (str, tuple)