nnterp · encyclopedia xglm · nnterp.families.xglm

family xglm

XGLM

OPT's pre-norm block with no MLP module, so the MLP path's term is fc2.output; the embedding is scaled by sqrt(d_model), and sinusoidal positions are added after it.

No MLP moduleScaled embeddingsPosition embeddingsLayerNormBiased q, k, v

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× 24
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + fc2.output == layer_output

The API

02 · support()

Values, by host

model XGLMForCausalLMlayoutwhat 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] XGLMDecoderLayerlayoutwhat 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 final_layer_norm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the feed-forward path

the child module final_layer_norm's output
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 XGLMAttentionlayoutwhat it isread at
attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention, already scaled by 1/sqrt(head_dim)

inside the forward: the output of query_states_reshape_0
attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the output of key_states_reshape_0
attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the output of value_states_reshape_0
attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward, at the operation that nn_functional_softmax_0|nn_functional_softmax_1 finds on each call
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
attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of torch_bmm_1

Printout

What print(model) shows on facebook/xglm-564M: the native modules, standard names and every standard value with its layout.

XGLMForCausalLM(
  (lm_head): Linear(in_features=1024, out_features=256008, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): XGLMScaledWordEmbedding(256008, 1024, padding_idx=1)
  (layers): ModuleList(
    (0-23): 24 x XGLMDecoderLayer(
      (self_attn): XGLMAttention(
        (k_proj): Linear(in_features=1024, out_features=1024, bias=True)
        (v_proj): Linear(in_features=1024, out_features=1024, bias=True)
        (q_proj): Linear(in_features=1024, out_features=1024, bias=True)
        (o_proj/out_proj): Linear(in_features=1024, out_features=1024, bias=True)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention, already scaled by 1/sqrt(head_dim)
        (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_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
        (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
      )
      (activation_fn): GELUActivation()
      (input_layernorm/self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True, bias=True)
      (fc1): Linear(in_features=1024, out_features=4096, bias=True)
      (fc2): Linear(in_features=4096, out_features=1024, bias=True)
      (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True, bias=True)
      (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 feed-forward path
      (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
      (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((1024,), 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_layers24
model.hidden_size1024
model.num_heads16
model.num_kv_heads16
model.head_dim64
model.qk_head_dim64
model.vocab_size256008
model.intermediate_size4096
architecturesXGLMForCausalLM
vocab_size256008
max_position_embeddings2048
tie_word_embeddingsTrue
activation_functiongelu
d_model1024
ffn_dim4096
scale_embeddingTrue

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(self_attn_layer_norm(x))     # self_attn_layer_norm = input_layernorm
out = h + fc2(act(fc1(final_layer_norm(h))))    # act: gelu, relu on xglm-4.5B

The block returns a bare tensor. attention_output is the attention module's output after out_proj. The block's own final_layer_norm is the MLP path's input norm and keeps its native name; the decoder's last norm is layer_norm, which is model.norm. Every norm is a LayerNorm with a weight and a bias, and every projection has a bias.

The MLP is fc1 and fc2 on the block

No block has an mlp module, so mlp, mlp_output and their kin do not exist on XGLM and support() lists no mlp.* key. What the MLP path adds is layers[i].fc2.output, [batch, seq, hidden], and the neurons are layers[i].activation_fn.output, [batch, seq, ffn_dim]:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    neurons = model.layers[1].activation_fn.output.save()
    mlp_out = model.layers[1].fc2.output.save()
    out = model.layers[1].layer_output.save()

torch.equal(x + attn + mlp_out, out)              # True

intermediate_size is the config's ffn_dim: four times the width on every size but xglm-4.5B, whose MLP is eight times the width (16384 on 2048) and uses ReLU where the others use GELU.

The embedding is scaled, the positions are sinusoidal

embed_tokens is an XGLMScaledWordEmbedding: it multiplies its lookup by sqrt(d_model) (32 on xglm-564M), and token_embeddings is the product. model.embed_positions is a fixed sinusoidal table with no parameters, kept in the buffer weights; position t reads row t + 2, and row 1, the padding index, is zeros. Its output is added after the scaled lookup:

with model.trace(prompt):
    ids = model.input_ids.save()
    tokens = model.token_embeddings.save()
    positions = model.model.embed_positions.output.save()
    x0 = model.layers[0].input.save()

lookup = model.embed_tokens.weight[ids]
torch.allclose(tokens, lookup * model.config.d_model ** 0.5)   # True
W = model.model.embed_positions._module.weights
torch.equal(tokens + positions, x0)                            # True
torch.equal(positions[0], W[2 : 2 + x0.shape[1]])              # True

On xglm-564M the raw lookup has a median norm of 4.1 per token, the scaled one 132, and every position's row 22.6. The table is rebuilt longer when a prompt runs past it, so no length raises. A left-padded prompt gives the same logits as the prompt alone (largest difference 1.3e-7 on the pinned checkpoint).

The attention is its own arithmetic

XGLM has one attention implementation, so the interior values need no eager load and carry no ⚠. The queries are scaled by 1/sqrt(head_dim) as they are projected, so attention_queries is q_proj's output times that scale. attention_scores is attention_queries @ attention_keys.transpose(-1, -2) plus the mask, clamped below at the dtype's minimum. Every value is a view of the [batch * num_heads, ...] tensor the forward holds, served heads first, so an in-place edit reaches the model. Every head has its own keys and values.

Block 0 writes a large shared vector, and position 0 is a sink

On xglm-564M block 0's two terms add a vector of norm about 860 to every position, almost the same at each one and largest in coordinates 35, 148 and 874: layer_output of block 0 has a median norm of 858, and its mean over positions has a norm of 858, against 132 for token_embeddings. The tokenizer prepends </s> (id 2), and the stream at that position grows to a norm of 13,000 from block 8 to block 16, mostly in coordinates 660 and 524, while the median elsewhere is 900 to 1,100. From block 7 to block 21, 59% to 82% of the attention pattern, averaged over heads and later queries, lands on key 0. Subtract the mean over positions, or leave position 0 out, before taking norms or activation statistics.

The readout: LayerNorm with a bias, tied weights

logits is lm_head.output with nothing applied, and project_on_vocab is lm_head(layer_norm(x)), so a logit lens at the last block equals logits. lm_head has no bias and its weight is embed_tokens.weight, the same tensor; the embedding's scale is applied in the forward, so the unembedding is the unscaled matrix.

The family's checkpoints

Blocks × width: 24 × 1024 (xglm-564M, 16 heads of 64), 24 × 2048 (xglm-1.7B), 48 × 2048 (xglm-2.9B, xglm-4.5B), 32 × 4096 (xglm-7.5B, 32 heads); head_dim is 128 from 1.7B on. Every size has 2048 trained positions and the same 256008-token multilingual vocabulary.

From the family module

04 · nnterp/families/xglm.py

XGLM (XGLMForCausalLM).

model.{embed_tokens, embed_positions, layers[i].{self_attn, self_attn_layer_norm, fc1, fc2, final_layer_norm}, layer_norm} and lm_head. embed_tokens is a scaled embedding (times sqrt(d_model) under scale_embedding), so token_embeddings is the scaled lookup; the sinusoidal positions are added after it, outside any module with a standard name. As on OPT there is no MLP module: fc1 and fc2 sit on the block, so mlp and mlp_output do not exist here and StandardizedTransformer.support says so. The block's own final_layer_norm is the pre-MLP norm and keeps its native name. The block returns a bare tensor. The attention does its own arithmetic whatever attn_implementation says (the model has no other): queries scaled by 1/sqrt(head_dim) as they are projected, queries, keys and values flattened to [batch * num_heads, seq, head_dim] for torch.bmm, the mask added and clamped, softmax, dropout, bmm with the values; the interior needs no eager load.