nnterp · encyclopedia olmo_hybrid · nnterp.families.olmo_hybrid

family olmo_hybrid

OLMo-Hybrid

Two block classes that norm in opposite places: three blocks in four are pre-norm, with a gated DeltaNet mixer, and every fourth is OLMo-3's post-norm block, whose attention (without rotary on most checkpoints) and MLP read the raw stream and have their outputs normed before the stream receives them.

HybridPost-norms onlyQuery/key normsBlocks without rotary

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block, each tick's color gives its block's shape, and the diagram redraws for the block picked.

embed_tokenstoken_embeddings
→
layers× 32
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + linear_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 OlmoHybridForCausalLMlayoutwhat 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] OlmoHybridLinearAttentionDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
⚠︎attention_norm_output scalar The normed stream entering the mixer

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the mixer's add

the child module post_attention_layernorm's input
⚠︎mlp_norm_output scalar 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].mlp OlmoHybridMLPlayoutwhat 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 post-feedforward norm's output on an attention block, the MLP's own on a linear block

inside the forward, at the operation that by_block finds on each call
model.layers[i].self_attn OlmoHybridAttentionlayoutwhat 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 post-attention norm's output

the sibling module post_attention_layernorm's 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].linear_attn OlmoHybridGatedDeltaNetlayoutwhat it isread at
attention_output Residual [batch seq hidden_size] What the mixer adds to the residual stream

the module's own output
⚠︎state State [batch heads key_dim value_dim] The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')

inside the forward, at the operation that _token_state_op finds on each call
⚠︎states States [batch seq heads key_dim value_dim] The recurrent state after every token of this call; needs route_kernels(family, 'torch')

inside the forward, at the operation that _states finds on each call
attention_queries LinearQK [batch seq heads key_dim] The queries entering the delta rule

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 0
attention_keys LinearQK [batch seq heads key_dim] The keys entering the delta rule

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 1
attention_values LinearV [batch seq heads value_dim] The values entering the delta rule

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 2
decays Gates [batch seq heads] The per-token log decay of the recurrent state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument g
betas Gates [batch seq heads] The per-token write strength into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument beta
state_input State | None [batch heads key_dim value_dim] The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument initial_state
attention_head_outputs LinearV [batch seq heads value_dim] The per-head outputs before the gated norm and the output projection

inside the forward: the output of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 0
state_output State [batch heads key_dim value_dim] The recurrent state leaving the layer

inside the forward: the output of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 1

Printout

What print(model) shows on allenai/Olmo-Hybrid-7B: the native modules, standard names and every standard value with its layout.

OlmoHybridForCausalLM(
  (lm_head): Linear(in_features=3840, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 3840, padding_idx=100277)
  (layers): ModuleList(
    (0-2): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (3): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (4-6): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (7): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (8-10): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (11): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (12-14): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (15): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (16-18): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (19): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (20-22): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (23): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (24-26): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (27): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (28-30): 3 x OlmoHybridLinearAttentionDecoderLayer(
      (mlp): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (linear_attn): OlmoHybridGatedDeltaNet(
        (q_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (k_proj): Linear(in_features=3840, out_features=2880, bias=False)
        (v_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (a_proj): Linear(in_features=3840, out_features=30, bias=False)
        (b_proj): Linear(in_features=3840, out_features=30, bias=False)
        (g_proj): Linear(in_features=3840, out_features=5760, bias=False)
        (o_proj): Linear(in_features=5760, out_features=3840, bias=False)
        (conv1d): Conv1d(11520, 11520, kernel_size=(4,), stride=(1,), padding=(3,), groups=11520, bias=False)
        (o_norm): OlmoHybridRMSNormGated()
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (decays) -> Gates [batch seq heads]: The per-token log decay of the recurrent state
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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
    )
    (31): OlmoHybridAttentionDecoderLayer(
      (self_attn): OlmoHybridAttention(
        (q_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (k_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (v_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (o_proj): Linear(in_features=3840, out_features=3840, bias=False)
        (q_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (k_norm): OlmoHybridRMSNorm((3840,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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): OlmoHybridMLP(
        (gate_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (up_proj): Linear(in_features=3840, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=3840, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output on an attention block, the MLP's own on a linear block
        (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
      )
      (post_attention_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (post_feedforward_layernorm): OlmoHybridRMSNorm((3840,), eps=1e-06)
      (attention_norm_output): The normed stream entering the mixer
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the mixer's add
      (mlp_norm_output): The normed stream entering the MLP
      (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): OlmoHybridRMSNorm((3840,), eps=1e-06)
  (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_size3840
model.num_heads30
model.num_kv_heads30
model.head_dim128
model.qk_head_dim128
model.vocab_size100352
model.intermediate_size11008
architecturesOlmoHybridForCausalLM
num_hidden_layers32
hidden_size3840
intermediate_size11008
num_attention_heads30
num_key_value_heads30
vocab_size100352
max_position_embeddings65536
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
linear_conv_kernel_dim4
linear_num_key_heads30
linear_num_value_heads30
linear_key_head_dim96
linear_value_head_dim192
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

# a DeltaNet block (OlmoHybridLinearAttentionDecoderLayer): pre-norm
h   = x + linear_attn(input_layernorm(x))
out = h + mlp(post_attention_layernorm(h))

# an attention block (OlmoHybridAttentionDecoderLayer): post-norm, no input_layernorm
h   = x + post_attention_layernorm(self_attn(x))
out = h + post_feedforward_layernorm(mlp(h))

config.layer_types is full_attention on every fourth block (3, 7, ..., 31) and linear_attention on the others: 24 DeltaNet and 8 attention blocks on every checkpoint. A block has linear_attn or self_attn, never both. The name post_attention_layernorm is on both classes in two roles: the MLP's input norm on a DeltaNet block, the attention's output norm on an attention block. Pick blocks outside the trace:

kinds = model.config.layer_types
delta = [i for i, t in enumerate(kinds) if t == "linear_attention"]
full = [i for i, t in enumerate(kinds) if t == "full_attention"]      # 3, 7, ..., 31

The contributions follow the block's class

On a DeltaNet block attention_output and mlp_output are the modules' own outputs. On an attention block they are the post-norms' outputs: self_attn.attention_output is post_attention_layernorm's output and mlp.mlp_output is post_feedforward_layernorm's, so mlp.output there is not what the stream receives. One MLP class serves both, and nnterp picks mlp_output's place by the block that holds it. The identity is the plain sum on both, exact in float32:

blk = model.layers[full[0]]
with model.trace(prompt):
    x = blk.input.save()
    attn = blk.self_attn.attention_output.save()      # post_attention_layernorm's output
    mlp = blk.mlp.mlp_output.save()                   # post_feedforward_layernorm's output
    out = blk.layer_output.save()

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

The attention block has no input norm: self_attn.input is the block's input and mlp.input is x + attention_output. Scaling self_attn.output there does not scale what the block adds, since the norm follows; edit attention_output.

Loading

The attention interior needs attn_implementation="eager". The DeltaNet values are read at transformers' pure-torch delta-rule kernels. Where flash-linear-attention is installed, the mixer calls its compiled kernel instead, and the values are unavailable until nnterp.route_kernels(model.family, "torch") binds the pure-torch ones. The per-token state and states need that call in any case, before the first trace: it runs a prompt through the token-by-token kernel, which keeps the state after every token.

import nnterp

nnterp.route_kernels(nnterp.families.olmo_hybrid, "torch")      # for state / states
model = StandardizedTransformer("allenai/Olmo-Hybrid-7B", attn_implementation="eager")

The DeltaNet mixer

Each mixer has 30 heads, keys 96 wide and values 192 wide, one value head per key head: attention_queries and attention_keys are [batch, seq, 30, 96], attention_values and attention_head_outputs [batch, seq, 30, 192], decays and betas [batch, seq, 30], and the state [batch, 30, 96, 192], 2.2 MB per block per prompt in float32. The queries and keys are read after the width-4 convolution and its SiLU, before the kernel's l2 norm. After the kernel, y passes a gated RMSNorm (o_norm, gated by g_proj) and o_proj; attention_head_outputs is y before both.

Every checkpoint sets linear_allow_neg_eigval: the mixer doubles the sigmoid of b_proj before the kernel, so betas lie in (0, 2); write betas in that range.

mix = model.layers[delta[0]].linear_attn
with model.trace(prompt):
    betas = mix.betas.save()               # [batch, seq, heads], in (0, 2)
    states = mix.states.save()             # [batch, seq, heads, key_dim, value_dim]; needs route_kernels
    final = mix.state_output.save()

assert torch.equal(states[:, -1], final)

Attention: normed across heads, no rotary

q_norm and k_norm are RMSNorms over the whole projection (30 heads × 128 = 3840 wide), applied before the heads are split, so one RMS covers every head and an edit to one head's slice of q_proj.output reaches the others; edit a single head at attention_queries. 30 query heads over 30 key/value heads: there is no grouping. The released configs set rope_theta to null, and the model then builds no rotary embedding: attention_queries is q_norm's output split into heads, unchanged. Token order reaches the attention only through the causal mask and through the DeltaNet blocks before it.

attn = model.layers[full[0]].self_attn
with model.trace(prompt):
    qn = attn.q_norm.output.save()
with model.trace(prompt):
    q = attn.attention_queries.save()          # [batch, num_heads, seq, head_dim]

b, s, _ = qn.shape
assert torch.equal(qn.view(b, s, -1, model.head_dim).transpose(1, 2), q)

Olmo-Hybrid-Think-SFT-7B is the exception. Its config sets rope_parameters itself to null, not rope_theta, and transformers fills a null rope_parameters with a theta of 10000. Loaded through transformers, that checkpoint's attention blocks therefore apply rotary embeddings: its model.rotary_emb is an OlmoHybridRotaryEmbedding, where the other checkpoints' is None.

The readout, the embeddings and the tokenizer

logits is lm_head.output, with no cap or scale, and project_on_vocab on the last block's layer_output equals logits exactly. lm_head and embed_tokens are separate weights, and nothing is added after embed_tokens, so token_embeddings equals layers[0].input. The tokenizer is a 100278-token BPE (the embedding has 100352 rows); it prepends nothing, so position 0 is the text's first token, and <|endoftext|> (100257) ends a text.

From the family module

04 · nnterp/families/olmo_hybrid.py

OLMo-Hybrid (OlmoHybridForCausalLM).

Llama's containers over two block classes, chosen by config.layer_types:

  • OlmoHybridLinearAttentionDecoderLayer (linear_attention) is pre-norm, as Llama: input_layernorm -> linear_attn (a gated DeltaNet, see nnterp.LinearAttention), then post_attention_layernorm -> mlp, each added to the stream as the module returns it.
  • OlmoHybridAttentionDecoderLayer (full_attention) is OLMo-3's post-norm sandwich: self_attn -> post_attention_layernorm and mlp -> post_feedforward_layernorm, the post-norms' outputs added to the stream; it has no input_layernorm, so the block input enters the attention.

Each block has self_attn or linear_attn, never both. The attention runs the shared eager forward (query and key norms before it; no rotary embedding when the checkpoint sets no rope_theta), so its contribution is OLMo-3's override at the post-attention norm. Both block classes hold one MLP class, OlmoHybridMLP, whose contribution is its own output on a linear block and the post-feedforward norm's output on an attention block, so mlp_output is a EProperty whose key is chosen by the parent block's layer type. The DeltaNet calls transformers' torch_chunk_gated_delta_rule / torch_recurrent_gated_delta_rule as Qwen3-Next does, so the base LinearAttention holds; with linear_allow_neg_eigval the betas are doubled before the kernel and lie in (0, 2).