mojolearn.ByteLanguageModelConfig#

class mojolearn.ByteLanguageModelConfig(batch=2, length=32, d_model=32, n_heads=4, n_kv=2, head_dim=8, intermediate=64, n_layers=2, vocab_size=256, chunked_lm_head_v2=False)[source]#

Bases: object

Dimensions of the decoder language model: batch, sequence, layers, vocabulary.

An immutable description of the model and training-batch shape passed to LanguageModelTrainer as shape=. Each layer has an attention block with grouped key/value heads and a gated MLP, each with its own norm weight; the model has a token embedding and a separate output head. The defaults describe the small published byte-level model (vocabulary 256). LanguageModelConfig is the same class.

Parameters:

  • batch (int, default 2) – Sequences per training step.

  • length (int, default 32) – Tokens per sequence, at most 8192.

  • d_model (int, default 32) – Model width; must equal n_heads * head_dim.

  • n_heads (int, default 4) – Attention heads; must be divisible by n_kv.

  • n_kv (int, default 2) – Key/value heads.

  • head_dim (int, default 8) – Width of each head; must be even.

  • intermediate (int, default 64) – MLP hidden width.

  • n_layers (int, default 2) – Number of decoder layers.

  • vocab_size (int, default 256) – Token vocabulary size (256 for bytes; corpus.n_vocab for a BPE corpus from mojolearn.lm_corpus).

  • chunked_lm_head_v2 (bool, default False) – Compute the loss and its gradient over bounded vocabulary panels, which bounds memory for large vocabularies. Logits from prediction still cover the full vocabulary.

Raises:

ValueError – If a dimension is not an integer in [1, 2**20], the head constraints above are not met, or the shape exceeds signed 32-bit indexing.

Parameters:

  • batch (int)

  • length (int)

  • d_model (int)

  • n_heads (int)

  • n_kv (int)

  • head_dim (int)

  • intermediate (int)

  • n_layers (int)

  • vocab_size (int)

  • chunked_lm_head_v2 (bool)

Examples

>>> import mojolearn
>>> shape = mojolearn.LanguageModelConfig(batch=1, length=128, d_model=64, n_heads=4,
...                                       n_kv=2, head_dim=16, intermediate=128,
...                                       n_layers=2, vocab_size=256)
>>> shape.n_total > 0
True
__init__(batch=2, length=32, d_model=32, n_heads=4, n_kv=2, head_dim=8, intermediate=64, n_layers=2, vocab_size=256, chunked_lm_head_v2=False)#
Parameters:

  • batch (int)

  • length (int)

  • d_model (int)

  • n_heads (int)

  • n_kv (int)

  • head_dim (int)

  • intermediate (int)

  • n_layers (int)

  • vocab_size (int)

  • chunked_lm_head_v2 (bool)

Return type:

None

Methods

to_dict()

The fields as a dict, suitable for ByteLanguageModelConfig(**d).

Attributes

batch

chunked_lm_head_v2

d_model

head_dim

intermediate

length

n_heads

n_kv

n_layers

n_tensors

2 + 9 * n_layers.

n_total

Total number of parameters.

native_shape

The dimensions as a tuple of ints, in field order.

offsets

Start offset of every parameter tensor in the flat parameter vector, plus the total.

parameter_names

The name of every parameter tensor, in packing order.

parameter_shapes

The shape of every parameter tensor, in packing order.

profile

The name of the training profile this shape selects.

vocab_size

property native_shape#

The dimensions as a tuple of ints, in field order.

property profile#

The name of the training profile this shape selects.

property parameter_shapes#

The shape of every parameter tensor, in packing order.

property n_tensors#

2 + 9 * n_layers.

Type:

Number of parameter tensors

property parameter_names#

The name of every parameter tensor, in packing order.

property offsets#

Start offset of every parameter tensor in the flat parameter vector, plus the total.

property n_total#

Total number of parameters.

to_dict()[source]#

The fields as a dict, suitable for ByteLanguageModelConfig(**d).