nnterp · encyclopedia gpt_neox_japanese · nnterp.families.gpt_neox_japanese

family gpt_neox_japanese

GPT-NeoX-Japanese

GPT-NeoX's module names on a sequential block, where the last block's attention returns a separate bias that the block adds and attention_output includes.

LayerNorm

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× 32
→
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

model GPTNeoXJapaneseForCausalLMlayoutwhat 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] GPTNeoXJapaneseLayerlayoutwhat 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 GPTNeoXJapaneseAttentionlayoutwhat it isread at
attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of self__attn_0, element 0
attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of self__attn_0, element 1
attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of self__attn_0, element 2
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 self__attn_0
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: its output plus, on the last block, the dense_bias the block adds

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 self_attention_dropout_0 inside self__attn_0
attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of self__attn_0, element 0
model.layers[i].mlp GPTNeoXJapaneseMLPlayoutwhat 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 abeja/gpt-neox-japanese-2.7b: the native modules, standard names and every standard value with its layout.

GPTNeoXJapaneseForCausalLM(
  (lm_head/embed_out): Linear(in_features=2560, out_features=32000, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(32000, 2560)
  (layers): ModuleList(
    (0-31): 32 x GPTNeoXJapaneseLayer(
      (input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True, bias=True)
      (post_attention_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True, bias=True)
      (self_attn/attention): GPTNeoXJapaneseAttention(
        (attention_dropout): Dropout(p=0.1, inplace=False)
        (query_key_value): Linear(in_features=2560, out_features=7680, bias=False)
        (o_proj/dense): Linear(in_features=2560, out_features=2560, bias=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_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output plus, on the last block, the dense_bias the block adds
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
      )
      (mlp): GPTNeoXJapaneseMLP(
        (dense_h_to_4h): Linear(in_features=2560, out_features=10240, bias=False)
        (down_proj/dense_4h_to_h): Linear(in_features=10240, out_features=2560, bias=False)
        (act_fn/act): GELUActivation()
        (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
      )
      (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((2560,), 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_layers32
model.hidden_size2560
model.num_heads32
model.num_kv_heads32
model.head_dim80
model.qk_head_dim80
model.vocab_size32000
model.intermediate_size10240
architecturesGPTNeoXJapaneseForCausalLM
num_hidden_layers32
hidden_size2560
num_attention_heads32
vocab_size32000
max_position_embeddings2048
hidden_actgelu
layer_norm_eps1e-05
tie_word_embeddingsTrue

The Notes

03 · what to know before you run

The block is sequential, under GPT-NeoX's names

h   = x + (attention(input_layernorm(x)) + dense_bias)    # dense_bias: last block only
out = h + mlp(post_attention_layernorm(h))

The modules are input_layernorm, attention (self_attn), post_attention_layernorm and mlp, as on GPT-NeoX, but post_attention_layernorm normalizes the stream after the attention has been added, so each block's MLP reads its attention's contribution. The block returns (hidden_states, attn_weights) on every call: layer_output is the first element.

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mid = model.layers[1].post_attention_layernorm.input.save()

torch.testing.assert_close(mid, x + attn)

The last block's attention bias is part of attention_output

None of the projections has a bias. The last block's attention holds one parameter, dense_bias (hidden_size wide), and returns it as the third element of its output, (attn_output, attn_weights, dense_bias); the block adds it to the attention's output before the residual. attention_output is that sum on the last block and the module's own output on the others, so layers[i].input + attention_output + mlp_output == layer_output holds on every block (bit-identical on 2.7B in float16). An edit of attention_output on the last block, in place or assigned, is carried back to the module's output minus the bias, and a read alone leaves the forward bit-identical.

last = model.layers[-1]
bias = last.self_attn._module.dense_bias

with model.trace(prompt):
    module = last.self_attn.output[0].save()
    attn = last.self_attn.attention_output.save()

assert torch.equal(attn, module + bias)

On 2.7B dense_bias has a norm of 5.4, against a median of 78 for the norm of the last attention's own output rows on a sample prompt. Zeroing self_attn.output[0] on the last block leaves the bias in the stream; zeroing attention_output removes both.

Queries, keys and values come from one projection, laid out per head

query_key_value is a Linear from hidden_size to 3 * hidden_size with no bias. Its output is viewed as [batch, seq, num_heads, 3 * head_dim]: each head's slice holds its queries, then its keys, then its values. The values are read as cut; the queries and keys are read after the rotary.

W = model.layers[1].self_attn.query_key_value.weight
heads, d = model.num_heads, model.head_dim

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

qkv = (h @ W.T).unflatten(-1, (heads, 3 * d)).transpose(1, 2)
q_raw, k_raw, v_again = qkv.split(d, dim=-1)                 # before the rotary
torch.testing.assert_close(v_again, v)

Rotary turns the whole head (rotary_pct is 1.0) by rotate_half: dimension i turns with i + head_dim / 2, at base 10000. Keys and values are num_heads wide, with no grouping.

The interior needs no eager load

The attention has one implementation, its own _attn(query, key, value, mask), whatever attn_implementation says, so a default load serves all six interior values and they carry no load condition. Its arguments are attention_queries, attention_keys and attention_values, heads first. The scores are a baddbmm scaled by 1 / sqrt(head_dim) (0.1118 at 2.7B's head_dim of 80), then the mask is added, so the masked entries hold the dtype's minimum. The pattern is read after attention_dropout, before its cast to the values' dtype, and the first return of _attn is attention_head_outputs.

with model.trace(prompt):
    q = model.layers[1].self_attn.attention_queries.save()
    k = model.layers[1].self_attn.attention_keys.save()
    scores = model.layers[1].self_attn.attention_scores.save()

n = scores.shape[-1]
causal = torch.ones(n, n, dtype=torch.bool, device=scores.device).tril()
again = q @ k.transpose(-1, -2) / model.head_dim ** 0.5
torch.testing.assert_close(again[..., causal], scores[..., causal])

Loading: the config names no dtype

abeja/gpt-neox-japanese-2.7b's config names no dtype, so a load without dtype is float32, about 11 GB; dtype=torch.float16 halves it. The tokenizer (GPTNeoXJapaneseTokenizer, 32000 tokens, with a separate emoji table) adds no BOS: position 0 is the prompt's first token, and <|startoftext|> (31996) is there to prepend by hand.

The first position is a sink with a huge norm

On 2.7B in float16, on a 28-token Japanese prompt, the stream at position 0 has a norm of 6,300 to 7,900 from block 3's output to block 24's, against medians of about 470 to 1,500 at the other positions. Blocks 3 to 29 put 69% to 88% of their pattern (averaged over heads and the later queries) on key 0. Leave position 0 out of norms you scale a steering vector by and of activation statistics.

The readout: tied weights and a LayerNorm with a bias

logits is lm_head.output with nothing applied, and project_on_vocab applied to the last block's layer_output equals logits exactly (difference 0.0 on 2.7B). lm_head is the native embed_out, with no bias, and its weight is embed_tokens.weight, the same tensor. final_layer_norm's bias adds lm_head.weight @ norm.bias to every position's logits whatever the input (a standard deviation of 0.39 logits on 2.7B). The model is 32 blocks × 2560, 32 heads of 80, an MLP four times the width (intermediate_multiple_size), and gelu.

From the family module

04 · nnterp/families/gpt_neox_japanese.py

GPT-NeoX-Japanese (GPTNeoXJapaneseForCausalLM).

gpt_neox_japanese.{embed_in, layers[i].{input_layernorm, post_attention_layernorm, attention, mlp}, final_layer_norm, rotary_emb} and embed_out, the head. GPT-NeoX's names, but a sequential block, not a parallel one: post_attention_layernorm normalizes the stream after the attention has been added. The block adds each sublayer through bias_dropout_add(x, bias, residual), residual + dropout(x + bias), and returns (hidden_states, attn_weights). The attention returns (attn_output, attn_weights, dense_bias): its output projection has no bias, and on the last block only a separate dense_bias is handed to the block to add, so what the attention adds is its output plus that bias: a sum served as attention_output, with an edit carried back to the module's output minus the bias. The attention does its own arithmetic in _attn whatever attn_implementation says (the model has no other): a baddbmm scaled by 1/sqrt(head_dim), the mask added, softmax, attention_dropout, a cast to the values' dtype, @ v; the interior needs no eager load.