LlamalikeTransformerConfig

LlamalikeTransformerConfig#

class penzai.experimental.v2.models.transformer.variants.llamalike_common.LlamalikeTransformerConfig[source]#

Bases: object

Common configuration parameters for a “llama-like” transformer.

These are held in a single configuration object to simplify argument passing during construction of the model.

Variables:

  • num_kv_heads (int) – The number of key-value attention heads or head groups.

  • query_head_multiplier (int) – The number of query heads for each KV head.

  • embedding_dim (int) – Dimension of the embedding vectors and residual stream.

  • projection_dim (int) – Dimension of the query, key, and value projections. Usually embedding_dim // num_heads.

  • mlp_hidden_dim (int) – Dimensionality of the hidden layer of the MLP blocks in each layer (the “neurons” axis).

  • num_decoder_blocks (int) – Number of transformer decoder blocks in the model.

  • vocab_size (int) – Number of tokens in the vocabulary.

  • mlp_variant (Literal['geglu_approx', 'swiglu']) – Gated linear unit variant for MLPs.

  • tie_embedder_and_logits (bool) – Whether to tie the weights of the input token embedding and output logit layers. If True, also scales down input token embeddings by sqrt(embedding_dim). (This is used by Gemma.)

  • rope_wavelength (float) – Wavelength for RoPE layers.

  • rms_norm_eps (float) – Epsilon for RMSNorm layers.

  • attention_type (AttentionType | Sequence[AttentionType]) – A single attention type or sequence of per-layer attention types. If a sequence, its length should evenly divide the number of decoder blocks, and will be repeated to match the number of blocks.

  • parameter_dtype (jax.typing.DTypeLike) – Floating dtype to use for all parameters.

  • activation_dtype (jax.typing.DTypeLike) – Floating dtype to use for activations and KV cache tables.

  • use_layer_stack (bool) – Whether to stack the blocks together using a LayerStack.

Methods

__init__(*, num_kv_heads, ...[, ...])

Attributes

attention_type

rms_norm_eps

rope_wavelength

use_layer_stack

num_kv_heads

query_head_multiplier

embedding_dim

projection_dim

mlp_hidden_dim

num_decoder_blocks

vocab_size

mlp_variant

tie_embedder_and_logits

Contents