nnterp · encyclopedia bloom · nnterp.families.bloom

family bloom

BLOOMBLOOMZ

Each sublayer adds the residual inside its own module, so the contributions are the tensors before that add; a LayerNorm sits between the embedding and block 0.

ALiBiBiased q, k, vLayerNorm

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 + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

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

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

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

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

inside the forward: the input of F_softmax_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: the tensor entering dropout_add

inside the forward: the input of dropout_add_0
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
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_0
model.layers[i].mlp BloomMLPlayoutwhat 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 tensor entering dropout_add

inside the forward: the input of dropout_add_0

Printout

What print(model) shows on bigscience/bloom-560m: the native modules, standard names and every standard value with its layout.

BloomForCausalLM(
  (lm_head): Linear(in_features=1024, out_features=250880, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(250880, 1024)
  (layers): ModuleList(
    (0-23): 24 x BloomBlock(
      (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True, bias=True)
      (self_attn/self_attention): BloomAttention(
        (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)
        (o_proj/dense): Linear(in_features=1024, out_features=1024, 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_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the tensor entering dropout_add
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
      )
      (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True, bias=True)
      (mlp): BloomMLP(
        (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)
        (act_fn/gelu_impl): BloomGelu()
        (down_proj/dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the tensor entering dropout_add
        (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
      )
      (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((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_size250880
model.intermediate_size4096
architecturesBloomForCausalLM
hidden_size1024
vocab_size250880
tie_word_embeddingsTrue
n_layer24
n_head16
layer_norm_epsilon1e-05

The Notes

03 · what to know before you run

Each module adds the residual itself

h   = self_attn(input_layernorm(x), residual=x)            # returns x + attention_output
out = mlp(post_attention_layernorm(h), residual=h)         # returns h + mlp_output

The block hands each sublayer the stream as an argument, and the module ends in dropout_add, which adds it. self_attn.output[0] is therefore the stream after the attention's add, and mlp.output is the block's output, equal to layer_output. attention_output and mlp_output are the first argument of each dropout_add: the attention's dense output and the MLP's dense_4h_to_h output, each with its bias. The sum is exact in every dtype.

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()
with model.trace(prompt):
    raw = model.layers[1].self_attn.output.save()

assert torch.equal(raw[0], x + attn)          # the module returns the stream
assert torch.equal(x + attn + mlp, out)

Ablate the contribution, not the module output

Zeroing attention_output removes the attention's contribution and leaves the stream: the MLP's norm then reads layers[i].input unchanged. Zeroing self_attn.output[0] zeroes the stream itself, and the MLP and every later block see zeros. The same holds for mlp_output against mlp.output. Steering, patching and attribution on a contribution go through the standard names for the same reason.

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

assert torch.equal(h, x)

The contribution is an operation inside the module, so in one trace read it before the module's own .output. Asking for self_attn.output first and attention_output after raises OutOfOrderError naming self_attention.source.dropout_add_0.input.

Queries, keys and values come out of one projection, head by head

query_key_value is one Linear with a bias, hidden_size to 3 * hidden_size. Its output is laid out by head: each head's query, key and value sit side by side, [q_h | k_h | v_h], so it splits as view(batch, seq, heads, 3, head_dim), not into thirds. The weight's rows follow the same order. The three values are views of that output; in-place edits of them land.

attn = model.layers[1].self_attn
with model.trace(prompt):
    h = model.layers[1].input_layernorm.output.save()
    q = attn.attention_queries.save()

fused = torch.nn.functional.linear(h, attn.query_key_value.weight, attn.query_key_value.bias)
per_head = fused.view(*fused.shape[:2], model.num_heads, 3, model.head_dim)
torch.testing.assert_close(per_head[..., 0, :].transpose(1, 2), q)
W = attn.query_key_value.weight.view(model.num_heads, 3, model.head_dim, -1)
W_q = W[:, 0]                                  # [num_heads, head_dim, hidden_size]

ALiBi biases replace rotary

There is no rotary embedding and no position embedding: queries and keys carry no position. The same token at two positions gives the same attention_queries and attention_keys at block 0. Position enters only through a bias on the scores: attention_scores is q @ kᵀ / sqrt(head_dim) plus slope_h * j for key position j, then the causal mask. The bias grows with the key's position and is the same on every query row, so after the softmax it acts as a penalty of slope_h per token of distance. With a power-of-two head count, head h has the slope 2 ** (-8 * (h + 1) / heads): on 560m, with 16 heads, 0.71 for head 0 down to 0.0039 for head 15. The 176B checkpoints have 112 heads: heads 0 to 63 take the same formula with 64 for heads, and head 64 + k takes 2 ** (-(2k + 1) / 16), halfway between two of the first 64. An edit to the queries or keys leaves the bias as it was. On a float32 load the scores recompute to rounding:

from transformers.models.bloom.modeling_bloom import build_alibi_tensor

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

seq = q.shape[2]
alibi = build_alibi_tensor(torch.ones(1, seq), model.num_heads, q.dtype).view(1, model.num_heads, 1, seq)
again = q @ k.transpose(-1, -2) / model.head_dim ** 0.5 + alibi
causal = torch.ones(seq, seq, dtype=torch.bool).tril()
torch.testing.assert_close(again[..., causal], scores[..., causal])

The attention needs no flag

transformers has one attention for BLOOM, its own arithmetic, so a load without attn_implementation runs it and all six interior values are served: support() reports none missing. The scores are read where the softmax takes them; the softmax runs in float32 and the pattern is cast back to the model's dtype. The head outputs are the bmm of the pattern and the values, [batch * num_heads, seq, head_dim] in the forward, served [batch, seq, num_heads, head_dim]. bigscience/bloom-560m stores float16 weights, but its config names no dtype, so a load without dtype is float32.

A LayerNorm sits between the embedding and block 0

token_embeddings is word_embeddings' output, and layers[0].input is that tensor through word_embeddings_layernorm, so the two differ. The norm has no standard name; it is model.transformer.word_embeddings_layernorm. On 560m the embedding rows are small (mean norm 0.37) and the first block's input has a norm of about 10 per token. An edit at token_embeddings passes through that LayerNorm, which rescales it with the rest of the row; edit layers[0].input to set what block 0 reads.

with model.trace(prompt):
    emb = model.token_embeddings.save()
    first = model.layers[0].input.save()

assert torch.equal(model.transformer.word_embeddings_layernorm(emb), first)

The readout

lm_head is word_embeddings' matrix (tied) with no bias, and logits equals lm_head.output. ln_f is a LayerNorm with a bias; project_on_vocab on the last block's layer_output equals logits. The tokenizer has 250680 entries and the matrix 250880 rows; the last 200 ids are padding. The tokenizer adds no BOS: position 0 is the prompt's first token.

From the family module

04 · nnterp/families/bloom.py

BLOOM (BloomForCausalLM).

transformer.{word_embeddings, word_embeddings_layernorm, h[i].{input_layernorm, self_attention, post_attention_layernorm, mlp}, ln_f} and lm_head. Both sublayers take the residual as an argument and add it inside the module (dropout_add), so the module outputs are residual-stream states; the contributions are the first argument of each dropout_add call. The attention does its own arithmetic whatever attn_implementation says, so the pattern is its dropout's output and needs no eager load. The embedding norm has no standard name.