nnterp · encyclopedia falcon · nnterp.families.falcon

family falcon

Falcon

A parallel block on every checkpoint but Falcon-RW: the block adds the attention's output into the MLP's output tensor in place, so mlp_output is a copy taken as the MLP returns.

Parallel blockLayerNormMulti-query attention

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

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

model FalconForCausalLMlayoutwhat 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] FalconDecoderLayerlayoutwhat 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 scalar 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 one norm feeds both sublayers, its own on the 40B layout

the child module input_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 FalconAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward, at the operation that apply_rotary_pos_emb_0|query_layer_0 finds on each call
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward, at the operation that apply_rotary_pos_emb_0|key_layer_0 finds on each call
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

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

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

inside the forward, at the operation that attn_output_1|flatten_0 finds on each call
model.layers[i].mlp FalconMLPlayoutwhat 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 (a copy, since the block adds the attention into the live tensor in place)

the module's own output

Printout

What print(model) shows on tiiuae/falcon-7b: the native modules, standard names and every standard value with its layout.

FalconForCausalLM(
  (lm_head): Linear(in_features=4544, out_features=65024, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(65024, 4544)
  (layers): ModuleList(
    (0-31): 32 x FalconDecoderLayer(
      (self_attn/self_attention): FalconAttention(
        (query_key_value): FalconLinear(in_features=4544, out_features=4672, bias=False)
        (o_proj/dense): FalconLinear(in_features=4544, out_features=4544, bias=False)
        (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_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
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
      )
      (mlp): FalconMLP(
        (dense_h_to_4h): FalconLinear(in_features=4544, out_features=18176, bias=False)
        (act_fn/act): GELUActivation()
        (down_proj/dense_4h_to_h): FalconLinear(in_features=18176, out_features=4544, bias=False)
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream (a copy, since the block adds the attention into the live tensor in place)
        (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
      )
      (input_layernorm): LayerNorm((4544,), eps=1e-05, elementwise_affine=True, bias=True)
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid): The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP: the one norm feeds both sublayers, its own on the 40B layout
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (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((4544,), 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_size4544
model.num_heads71
model.num_kv_heads1
model.head_dim64
model.qk_head_dim64
model.vocab_size65024
model.intermediate_size18176
architecturesFalconForCausalLM
num_hidden_layers32
hidden_size4544
num_attention_heads71
vocab_size65024
max_position_embeddings2048
tie_word_embeddingsTrue
dtypebfloat16
layer_norm_epsilon1e-05

The Notes

03 · what to know before you run

On 7B and 11B one LayerNorm feeds both sublayers

h   = input_layernorm(x)
m   = mlp(h)
m  += self_attn(h)          # in place, into the MLP's output tensor
out = m + x

input_layernorm is the block's only norm. Its output tensor is passed to the attention and then to the MLP, so self_attn.input and mlp.input are the same tensor: an in-place edit of self_attn.input reaches the MLP too, and an assignment replaces only the attention's argument. The MLP never reads its own block's attention: the first MLP that reads block i's attention is block i + 1's. ffn_hidden_size is the MLP's width, which nnterp serves as intermediate_size.

mlp_output is a copy, and the live tensor is not the MLP's

The block adds the attention's output into the MLP's output tensor in place, so once the block has run, mlp.output holds mlp + attn. mlp_output is a copy taken as the MLP returns and holds the MLP's contribution alone. nnterp hands an edit of the copy back to the model, so an in-place edit and an assignment both land.

with model.trace(prompt):
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    live = model.layers[1].mlp.output.save()

assert torch.equal(live, mlp + attn)            # the live tensor after the in-place add

The block's own order is (mlp + attn) + x, so that sum reproduces layer_output bit for bit; x + attn + mlp differs from it by rounding in bfloat16.

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()

assert torch.equal((mlp + attn) + x, out)

One key/value head on 7B, eight broadcast on 11B

query_key_value is one Linear with no bias. On 7B (multi_query) its output is every head's query, then one key, then one value: view(batch, seq, num_heads + 2, head_dim), whose last two heads are the key and the value. attention_keys and attention_values have one head, [batch, 1, seq, 64], and an edit to either reaches all 71 query heads.

11B groups the output by key/value head, [q × 4 | k | v] for each of 8 groups. The attention copies each key and value to its four query heads before the rotary, so attention_keys and attention_values are num_heads (32) wide while model.num_kv_heads is 8. An edit to one of the four copies reaches only that query head.

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

fused = h @ attn.query_key_value.weight.T
rows = fused.view(*fused.shape[:2], model.num_heads + 2, model.head_dim)    # the 7B layout
torch.testing.assert_close(rows[:, :, -1:].transpose(1, 2), v)

Without ALiBi, read the values before the queries and keys

The attention does its own arithmetic. On the rotary checkpoints the queries and keys are the two returns of the rotary embedding, and the values are bound before it, so inside one trace read attention_values first. Asking for the queries or keys first raises OutOfOrderError naming value_layer_0.

attn = model.layers[1].self_attn
with model.trace(prompt):
    v = attn.attention_values.save()       # binds first
    q = attn.attention_queries.save()
    k = attn.attention_keys.save()

Rotary turns the whole head in rotate_half pairs, with rope_theta 10000 on 7B and 500042 on 11B. The pattern is the softmax's output (no dropout follows it), and the head outputs are the pattern @ values product, heads first in the forward and served [batch, seq, num_heads, head_dim].

The interior needs an eager load

transformers' FalconAttention runs sdpa unless told otherwise, and then neither scores nor pattern exist: the six interior values carry ⚠ and report load with attn_implementation='eager'. attention_output and mlp_output do not depend on the implementation.

model = StandardizedTransformer("tiiuae/falcon-7b", attn_implementation="eager")

Two more layouts load as FalconForCausalLM

  • Falcon-40B and Falcon-180B (new_decoder_architecture, two norms in parallel) give each sublayer its own LayerNorm of the block input, ln_attn and ln_mlp; neither has a standard name. The fused projection and the broadcast keys and values are 11B's, with 16 query heads per key/value head on 40B.
  • Falcon-RW (tiiuae/falcon-rw-1b, falcon-rw-7b) is sequential, with input_layernorm and post_attention_layernorm, biased projections, and ALiBi in place of rotary. Its fused projection is laid out by head, [q_h | k_h | v_h]. With ALiBi the read order is the forward's: queries, then keys, then values; values first raises OutOfOrderError naming query_layer_0.
  • On an ALiBi checkpoint the eager attention adds the bias twice: transformers builds it into the attention mask and adds it again to the scores, so attention_scores is q @ kᵀ / sqrt(head_dim) + 2 * slope_h * j / sqrt(head_dim) for key position j, plus the causal mask. Head h's slope is 2 ** (-8 * (h + 1) / num_heads), and transformers computes slope_h * j in bfloat16 on every load. The default sdpa load adds the bias once. On falcon-rw-1b in float32 the two loads' logits differ by up to about 1.5 on a sample prompt, so an eager trace there runs a different model from the default one.

From the family module

04 · nnterp/families/falcon.py

Falcon (FalconForCausalLM): the 7B layout (parallel attention), with or without alibi, and the 40B layout.

transformer.{word_embeddings, h[i].{input_layernorm, self_attention, mlp}, ln_f} and lm_head. One norm feeds both sublayers and the block sums x + attn + mlp, but it does so by adding the attention output into the MLP's output tensor in place, so mlp_output reads a copy taken as the MLP returns, and a transform carries edits to it back into the model.

The attention does its own arithmetic, and config.alibi picks one of two branches of it with different operations, so each interior value names its operation by that flag. Without alibi the queries and keys leave apply_rotary_pos_emb and the pattern is the first softmax, which has no dropout after it; with alibi there is no rotary, the pattern is the dropout after the second softmax, and the head outputs are flattened over batch and heads. The 40B layout (ln_attn / ln_mlp, new_decoder_architecture) runs the same forward with its key/value heads already broadcast. falcon-11B has the 40B layout with num_ln_in_parallel_attn 1: one input_layernorm feeds both sublayers, and 8 key/value heads serve its 32 query heads.