nnterp · encyclopedia persimmon · nnterp.families.persimmon

family persimmon

Persimmon

Llama's sequential block with LayerNorms; the attention cuts queries, keys and values per head from one fused projection and norms each query and key head before the rotary.

Biased q, k, vQuery/key normsPartial rotarySquared ReLULayerNorm

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 36
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model PersimmonForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] PersimmonDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn PersimmonAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp PersimmonMLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output

Printout

What print(model) shows on adept/persimmon-8b-base: the native modules, standard names and every standard value with its layout.

PersimmonForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=262144, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(262144, 4096)
  (layers): ModuleList(
    (0-35): 36 x PersimmonDecoderLayer(
      (self_attn): PersimmonAttention(
        (query_key_value): Linear(in_features=4096, out_features=12288, bias=True)
        (o_proj/dense): Linear(in_features=4096, out_features=4096, bias=True)
        (q_layernorm): LayerNorm((64,), eps=1e-05, elementwise_affine=True, bias=True)
        (k_layernorm): LayerNorm((64,), eps=1e-05, elementwise_affine=True, bias=True)
        (attention_dropout): Dropout(p=0.0, inplace=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): PersimmonMLP(
        (dense_h_to_4h): Linear(in_features=4096, out_features=16384, bias=True)
        (down_proj/dense_4h_to_h): Linear(in_features=16384, out_features=4096, bias=True)
        (act_fn/act): ReLUSquaredActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=True)
      (post_attention_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=True)
      (dropout): Dropout(p=0.6, inplace=False)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True, bias=True)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers36
model.hidden_size4096
model.num_heads64
model.num_kv_heads64
model.head_dim64
model.qk_head_dim64
model.vocab_size262144
model.intermediate_size16384
architecturesPersimmonForCausalLM
num_hidden_layers36
hidden_size4096
intermediate_size16384
num_attention_heads64
num_key_value_heads64
vocab_size262144
max_position_embeddings16384
hidden_actrelu2
rms_norm_eps1e-06
layer_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16
qk_layernormTrue

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))
out = h + dropout(mlp(post_attention_layernorm(h)))     # dropout is the identity in eval

The block returns a bare tensor, and nothing but the eval-mode dropout sits between a sublayer and the stream: attention_output is the attention's output after dense, mlp_output is the MLP's output after dense_4h_to_h, and layers[i].input + attention_output + mlp_output == layer_output holds exactly. Every norm is a LayerNorm with a weight and a bias, and every projection has a bias.

Queries, keys and values are cut per head from one projection

self_attn.query_key_value is one biased nn.Linear from the width to three times the width. Its output is laid out head by head, [num_heads, 3, head_dim] on the last axis: head h's query, key and value are three consecutive head_dim blocks, not three blocks of the whole width. The rows of the weight follow the same layout:

attn = model.layers[1].self_attn
H, D = model.num_heads, model.head_dim
W = attn.query_key_value.weight.view(H, 3, D, -1)    # [num_heads, q/k/v, head_dim, hidden_size]
b = attn.query_key_value.bias.view(H, 3, D)

with model.trace(prompt):
    h = model.layers[1].self_attn.input.save()
    v = model.layers[1].self_attn.attention_values.save()

v_again = torch.einsum("bsx,hdx->bhsd", h, W[:, 2]) + b[:, 2, None]
torch.testing.assert_close(v_again, v)

Every head has its own keys and values (num_kv_heads == num_heads, 64 heads of 64 on the 8B checkpoints). The attention runs through the shared interface: the interior values need attn_implementation="eager", and the scores are scaled by head_dim ** -0.5 (0.125) inside it.

Queries and keys are normed per head, then half of each is turned

Under qk_layernorm (set on both checkpoints) q_layernorm and k_layernorm are LayerNorms over head_dim, one weight shared by every head, applied to the split queries and keys before rotary. Rotary then turns the first partial_rotary_factor of each head (0.5: 32 of 64 dimensions, with rope_theta 25000); the other half carries no position. attention_queries and attention_keys are read after both, so their last half is the normed projection unchanged and their first half is rotated:

rot = int(D * model.config.partial_rotary_factor)
with model.trace(prompt):
    h = model.layers[1].self_attn.input.save()
    q = model.layers[1].self_attn.attention_queries.save()

q_raw = torch.einsum("bsx,hdx->bhsd", h, W[:, 0]) + b[:, 0, None]
q_normed = attn.q_layernorm._module(q_raw.transpose(1, 2)).transpose(1, 2)
torch.testing.assert_close(q_normed[..., rot:], q[..., rot:])

An edit to the query slices of query_key_value.output passes through the norm, so it is rescaled and re-centered per head before it reaches the scores. An edit to attention_queries is not.

The MLP is squared ReLU with no gate

mlp is dense_h_to_4h, the activation and dense_4h_to_h, 4096 → 16384 → 4096 on the 8B checkpoints, with no gate projection. The checkpoints set hidden_act to relu2, relu(x) ** 2: a neuron, mlp.act.output[..., n], is exactly zero wherever its pre-activation is negative. The pinned tiny checkpoint uses GELU, so a check of that property needs the real weights.

The readout: LayerNorm with a bias, untied weights

logits is lm_head.output with nothing applied, and project_on_vocab is lm_head(final_layernorm(x)), so a logit lens at the last block equals logits. lm_head has its own weight and no bias. The tokenizer prepends |ENDOFTEXT|, its BOS token, to every prompt.

The family's checkpoints

persimmon-8b-base and persimmon-8b-chat share one shape: 36 blocks of width 4096, 64 heads of 64, an MLP of 16384, a 262144-token vocabulary and 16384 positions, stored in bfloat16. Fuyu-8B's language model is a Persimmon, but nnterp does not bind the fuyu wrapper.

From the family module

04 · nnterp/families/persimmon.py

Persimmon (PersimmonForCausalLM).

Llama's tree with Persimmon's final norm: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, final_layernorm} and lm_head; the residual added in the block, attention through the shared eager forward. The norms are nn.LayerNorm. The attention projects queries, keys and values with one fused query_key_value, norms the queries and keys per head (q_layernorm, k_layernorm, under qk_layernorm) and applies the rotary embedding to partial_rotary_factor of each head, all before the interface. The MLP is dense_h_to_4h/dense_4h_to_h with no gate; the block's own dropout after it is the identity in eval.