nnterp · encyclopedia phi · nnterp.families.phi

family phi

Phi-1Phi-1.5Phi-2

A parallel block, in which one LayerNorm feeds both the attention and the MLP and the block adds their outputs at once, with a bias on every attention projection and on lm_head, and rotary on only 32 dimensions of each head.

Parallel blockPartial rotaryBiased 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× 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 PhiForCausalLMlayoutwhat 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] PhiDecoderLayerlayoutwhat 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 Unavailable: a parallel block: the attention and the MLP read the same input, so there is no stream between them

nowhere: the family does not serve it
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP: the one norm feeds both sublayers

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 PhiAttentionlayoutwhat 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 PhiMLPlayoutwhat 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 microsoft/phi-2: the native modules, standard names and every standard value with its layout.

PhiForCausalLM(
  (lm_head): Linear(in_features=2560, out_features=51200, bias=True)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(51200, 2560)
  (layers): ModuleList(
    (0-31): 32 x PhiDecoderLayer(
      (self_attn): PhiAttention(
        (q_proj): Linear(in_features=2560, out_features=2560, bias=True)
        (k_proj): Linear(in_features=2560, out_features=2560, bias=True)
        (v_proj): Linear(in_features=2560, out_features=2560, bias=True)
        (o_proj/dense): Linear(in_features=2560, out_features=2560, bias=True)
        (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): PhiMLP(
        (act_fn/activation_fn): NewGELUActivation()
        (fc1): Linear(in_features=2560, out_features=10240, bias=True)
        (down_proj/fc2): Linear(in_features=10240, out_features=2560, bias=True)
        (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((2560,), eps=1e-05, elementwise_affine=True, bias=True)
      (resid_dropout): Dropout(p=0.1, inplace=False)
      (layer_mid): Unavailable: a parallel block: the attention and the MLP read the same input, so there is no stream between them
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP: the one norm feeds both sublayers
      (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((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_size51200
model.intermediate_size10240
architecturesPhiForCausalLM
num_hidden_layers32
hidden_size2560
intermediate_size10240
num_attention_heads32
num_key_value_heads32
vocab_size51200
max_position_embeddings2048
hidden_actgelu_new
layer_norm_eps1e-05
tie_word_embeddingsFalse
dtypefloat16
qk_layernormFalse

The Notes

03 · what to know before you run

The block is parallel, with one LayerNorm

h   = input_layernorm(x)
out = self_attn(h) + mlp(h) + x

input_layernorm is the block's only norm. Its output is both self_attn.input and mlp.input, the same tensor, and the block has no post_attention_layernorm. attention_output is dense's output and mlp_output is fc2's, and nothing normalizes either before the add. Nothing connects the two sublayers of one block: zeroing attention_output leaves that block's mlp.input bit-identical, so block i's attention first reaches an MLP at block i + 1.

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

assert torch.equal(attn + mlp + x, out)   # the block's own order, bit for bit

with model.trace(prompt):
    model.layers[1].self_attn.attention_output[:] = 0
    mlp_in_ablated = model.layers[1].mlp.input.save()

assert torch.equal(mlp_in, mlp_in_ablated)

The block sums the attention and the MLP first and adds the residual last. In the identity's order, x + attn + mlp, the sum differs from layer_output by rounding: by up to 7e-9 on the pinned checkpoint in float32, and by up to 2.0 over phi-2's blocks in bfloat16, whose stream entries reach 800, enough to fail torch.testing.assert_close at its bfloat16 tolerance. The block returns a bare tensor, so model.layers[i].output is layer_output. resid_dropout wraps both contributions and embed_dropout the embeddings (resid_pdrop is 0.1 on phi-2); both are inert in eval.

Every projection in the attention has a bias

self_attn.q_proj, k_proj, v_proj and dense are separate Linears, each with a bias: dense is the attention's output projection, not a layer of the MLP. The MLP is fc1, the activation (gelu_new on all three checkpoints) and fc2, both with a bias. Keys and values have num_heads heads: phi-1 and phi-1.5 leave num_key_value_heads unset, which transformers fills with num_attention_heads, and phi-2 sets it to 32. The class has a qk_layernorm option, a LayerNorm over each query and key head before the rotary; it is false on all three checkpoints, so nothing normalizes the queries and keys.

Rotary turns 32 dimensions of each head

partial_rotary_factor is 0.5 of a 64-dimensional head on phi-1 and phi-1.5, and 0.4 of an 80-dimensional head on phi-2: 32 dimensions on every checkpoint. The rotation is rotate_half on those 32, pairing dimension i with i + 16; the other dimensions of every query and key are the projection's output unchanged and carry no position. attention_queries and attention_keys are read after the rotary; for the queries before it, read the projection in the same trace:

with model.trace(prompt):
    q_raw = model.layers[1].self_attn.source.self_q_proj_0.output.save()
    q = model.layers[1].self_attn.attention_queries.save()

b, s, _ = q_raw.shape
q_pre = q_raw.view(b, s, model.num_heads, model.head_dim).transpose(1, 2)
rot = int(model.head_dim * model.config.rope_parameters["partial_rotary_factor"])   # 32
assert torch.equal(q_pre[..., rot:], q[..., rot:])   # past rot: no rotation

The scale is head_dim ** -0.5, and softmax(attention_scores) equals attention_probabilities.

phi-2 in float16 under eager attention returns NaN

phi-1.5 and phi-2 are stored in float16 (torch_dtype), phi-1 in float32, and a load without dtype keeps it. The eager path multiplies queries by keys and scales the product after, and on phi-2 the unscaled product passes float16's largest value, 65504: on "The Eiffel Tower is in the city of" it reaches 1.1e6 at block 29 and stays above 65504 at blocks 30 and 31. Under attn_implementation="eager" in float16, block 29's attention_output is not finite and every logit is NaN. The default sdpa load in float16 is finite (p(Paris) = 0.94), and so is eager in bfloat16 (0.92). Load the attention interior in bfloat16 or float32:

model = StandardizedTransformer(
    "microsoft/phi-2", dispatch=True, attn_implementation="eager", dtype=torch.bfloat16
)

On phi-1.5 the same prompt's largest product is 8.7e3, inside float16's range. With the dtype set, the eager and sdpa loads compute the same attention: no softcap, sink or window.

Position 0 carries one very large entry

The stream at the first position has a single dimension far larger than anything else. At block 16 of phi-2 on the sample prompt, dimension 743 of position 0 is about -790, and no entry at any other position exceeds 17 in magnitude. phi-1.5 shows the same at block 12: -684 at dimension 725, against at most 19 elsewhere. The tokenizer adds no BOS, so position 0 is the prompt's own first token. A norm or a mean of the stream taken over positions is dominated by it; leave position 0 out, or read it on its own.

The head has a bias, and a constant logit vector

logits equals lm_head.output, and project_on_vocab applied to the last block's layer_output equals logits exactly, final_layernorm's bias and lm_head's included. A logit lens therefore adds the same vector at every block, whatever the input:

bias_logits = model.lm_head.weight @ model.norm.bias + model.lm_head.bias

On phi-2 its standard deviation is 0.33 logits, against 2.3 for a whole lens readout at block 16 on the sample prompt (lm_head.bias alone: 0.04), and its top tokens are \n, -, ., \n\n and ,. embed_tokens and lm_head are separate matrices. The vocabulary is padded: 51200 rows against 50295 tokenizer tokens. The 905 extra rows are not zero: at the sample prompt's last position their logits are below -1 and their total probability is about 1e-6.

The tokenizer is CodeGen's

The three checkpoints ship CodeGen's -mono tokenizer: GPT-2's 50257 tokens plus one token for each run of 2 to 31 spaces and of 2 to 9 tabs. A run of two or more spaces is one token, and the word after it loses its leading space:

model.tokenizer("    return x").input_ids   # [50284, 7783, 2124]: '    ', 'return', ' x'

<|endoftext|> (50256) is its bos_token and eos_token, and the tokenizer adds neither.

Three checkpoints

The sizes, from the configs:

  • phi-1 (1.3B): 24 blocks, hidden 2048, 32 heads × 64, float32 files
  • phi-1.5 (1.3B): 24 blocks, hidden 2048, 32 heads × 64, float16 files
  • phi-2 (2.7B): 32 blocks, hidden 2560, 32 heads × 80, float16 files

All three set a context of 2048 tokens and untied embeddings.

From the family module

04 · nnterp/families/phi.py

Phi 1 / 1.5 / 2 (PhiForCausalLM).

Llama's containers with a parallel block: one input_layernorm feeds both sublayers and x + attn + mlp is summed at the end (resid_dropout on each, inert in eval). The final norm is final_layernorm. There is no post_attention_layernorm.