nnterp · encyclopedia helium · nnterp.families.helium

family helium

Helium

Llama's block, with one key/value head per query head and a rotary embedding that turns adjacent pairs of dimensions (2i, 2i + 1) rather than i with i + head_dim / 2.

Interleaved rotary

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

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

model HeliumForCausalLMlayoutwhat 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] HeliumDecoderLayerlayoutwhat 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 HeliumAttentionlayoutwhat 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 HeliumMLPlayoutwhat 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 kyutai/helium-1-preview-2b: the native modules, standard names and every standard value with its layout.

HeliumForCausalLM(
  (lm_head): Linear(in_features=2560, out_features=48000, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(48000, 2560, padding_idx=3)
  (layers): ModuleList(
    (0-23): 24 x HeliumDecoderLayer(
      (self_attn): HeliumAttention(
        (q_proj): Linear(in_features=2560, out_features=2560, bias=False)
        (k_proj): Linear(in_features=2560, out_features=2560, bias=False)
        (v_proj): Linear(in_features=2560, out_features=2560, bias=False)
        (o_proj): 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_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): HeliumMLP(
        (gate_proj): Linear(in_features=2560, out_features=7040, bias=False)
        (up_proj): Linear(in_features=2560, out_features=7040, bias=False)
        (down_proj): Linear(in_features=7040, out_features=2560, bias=False)
        (act_fn): SiLUActivation()
        (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): HeliumRMSNorm((2560,), eps=1e-08)
      (post_attention_layernorm): HeliumRMSNorm((2560,), eps=1e-08)
      (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): HeliumRMSNorm((2560,), eps=1e-08)
  (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_size2560
model.num_heads20
model.num_kv_heads20
model.head_dim128
model.qk_head_dim128
model.vocab_size48000
model.intermediate_size7040
architecturesHeliumForCausalLM
num_hidden_layers24
hidden_size2560
intermediate_size7040
num_attention_heads20
num_key_value_heads20
head_dim128
vocab_size48000
max_position_embeddings4096
hidden_actsilu
rms_norm_eps1e-08
tie_word_embeddingsFalse
dtypebfloat16
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))
out = h + mlp(post_attention_layernorm(h))
mlp(z) = down_proj(silu(gate_proj(z)) * up_proj(z))

Llama's block, names and SwiGLU MLP. The RMSNorms have eps 1e-8 and multiply by their gain in float32 before casting back to the input's dtype.

The contributions are the modules' outputs

attention_output is self_attn.output[0] and mlp_output is mlp.output, added to the stream with nothing in between. The identity holds exactly in float32 on the pinned checkpoint:

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

torch.testing.assert_close(x + attn + mlp, out)

Rotary turns adjacent pairs

Helium's rotary pairs dimension 2i with 2i + 1: its rotate_half takes the even and odd dimensions, and apply_rotary_pos_emb repeats each frequency twice in a row (repeat_interleave) before rotating. Llama pairs i with i + head_dim / 2. The rotary turns the whole head, with rope_theta 100,000. attention_queries and attention_keys are read after the rotation; on the pinned checkpoint they equal the interleaved rotation of q_proj.output and k_proj.output exactly, and differ from the rotate_half one. Code written for a rotate_half family's queries needs the head reordered, evens then odds:

hd = model.head_dim
perm = torch.cat([torch.arange(0, hd, 2), torch.arange(1, hd, 2)])

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

q_half = q[..., perm]          # in rotate_half's order

One key/value head per query head

helium-1-preview-2b has 20 query heads and 20 key/value heads, each 128 wide, in a 2560-wide stream: attention_keys and attention_values are [batch, 20, seq, 128], and an edit to head j reaches query head j alone. No projection has a bias.

The default load runs the same attention as eager

The six attention interior values need attn_implementation="eager". There is no softcap, window or sink: on helium-1-preview-2b in float32 the default sdpa load's logits differ from the eager load's by at most 2e-5.

<s> at position 0 is the sink

The tokenizer prepends <s> (id 1) though tokenizer.bos_token is None, so code that tests bos_token to decide whether a BOS is present gets the wrong answer. On helium-1-preview-2b the stream at position 0 has a norm of 930 after block 10, against 9 to 12 at the other positions of a test prompt, and the heads of block 10 put on average 0.84 of their attention on it. Drop [:, 1:] before averaging over positions.

Target tokens carry the word marker

The tokenizer is SentencePiece: "Paris" encodes to ▁Paris (2291), and " Paris" to two tokens, ▁ and ▁Paris. After "The Eiffel Tower is in the city of" the model puts 0.986 on ▁Paris. Take the target id from the word without its leading space:

ids = model.tokenizer("Paris", add_special_tokens=False).input_ids   # [2291]

The readout

model.logits is model.lm_head.output: no cap and no scale, and project_on_vocab is lm_head(norm(hidden)). lm_head and embed_tokens are separate weights.

What loads as this family

Only kyutai/helium-1-preview-2b is model_type helium. Helium 1 (kyutai/helium-1-2b and its per-domain variants -books, -wiki, -science and the others) and Sequential_Helium_6B are model_type llama and load as the Llama family, with Llama's rotate_half rotary; Helium1-VL-2B is a remote-code architecture.

From the family module

04 · nnterp/families/helium.py

Helium (HeliumForCausalLM).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, attention through the shared eager forward.