nnterp · encyclopedia deepseek_v2 · nnterp.families.deepseek_v2

family deepseek_v2

DeepSeek-V2DeepSeek-Coder-V2

Latent attention with queries and keys wider than the values and one rotary key shared by every head, and a softmax-routed mixture with two shared experts on every block but the first.

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× 27
→
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 DeepseekV2ForCausalLMlayoutwhat 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] DeepseekV2DecoderLayerlayoutwhat 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 DeepseekV2Attentionlayoutwhat 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 DeepseekV2MLP / DeepseekV2Moelayoutwhat 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

router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

inside router's forward: the output of F_linear_0

expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2

expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1

expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1

routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output

shared_expert_output Residual [batch seq hidden_size] The shared expert's output

the child module shared_experts's output

Printout

What print(model) shows on deepseek-ai/DeepSeek-V2-Lite: the native modules, standard names and every standard value with its layout.

DeepseekV2ForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=102400, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(102400, 2048)
  (layers): ModuleList(
    (0): DeepseekV2DecoderLayer(
      (self_attn): DeepseekV2Attention(
        (q_proj): Linear(in_features=2048, out_features=3072, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2048, out_features=576, bias=False)
        (kv_a_layernorm): DeepseekV2RMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=4096, bias=False)
        (o_proj): Linear(in_features=2048, out_features=2048, 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): DeepseekV2MLP(
        (gate_proj): Linear(in_features=2048, out_features=10944, bias=False)
        (up_proj): Linear(in_features=2048, out_features=10944, bias=False)
        (down_proj): Linear(in_features=10944, out_features=2048, 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): DeepseekV2RMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): DeepseekV2RMSNorm((2048,), eps=1e-06)
      (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
    )
    (1-26): 26 x DeepseekV2DecoderLayer(
      (self_attn): DeepseekV2Attention(
        (q_proj): Linear(in_features=2048, out_features=3072, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2048, out_features=576, bias=False)
        (kv_a_layernorm): DeepseekV2RMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=4096, bias=False)
        (o_proj): Linear(in_features=2048, out_features=2048, 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): DeepseekV2Moe(
        (experts): DeepseekV2Experts(
          (act_fn): SiLUActivation()
        )
        (router/gate): DeepseekV2TopkRouter()
        (shared_experts): DeepseekV2MLP(
          (gate_proj): Linear(in_features=2048, out_features=2816, bias=False)
          (up_proj): Linear(in_features=2048, out_features=2816, bias=False)
          (down_proj): Linear(in_features=2816, out_features=2048, 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
        )
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV2RMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): DeepseekV2RMSNorm((2048,), eps=1e-06)
      (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): DeepseekV2RMSNorm((2048,), 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_layers27
model.hidden_size2048
model.num_heads16
model.num_kv_heads16
model.head_dim128
model.qk_head_dim192
model.vocab_size102400
model.intermediate_size10944
architecturesDeepseekV2ForCausalLM
num_hidden_layers27
hidden_size2048
intermediate_size10944
num_attention_heads16
num_key_value_heads16
head_dim64
vocab_size102400
max_position_embeddings163840
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypebfloat16
kv_lora_rank512
qk_nope_head_dim128
qk_rope_head_dim64
v_head_dim128
first_k_dense_replace1
moe_intermediate_size1408
n_routed_experts64
num_experts_per_tok6
n_shared_experts2
routed_scaling_factor1.0
norm_topk_probFalse
attention_biasFalse
n_group1
topk_group1

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))      # multi-head latent attention
out = h + mlp(post_attention_layernorm(h))   # dense on block 0, a mixture after

Llama's pre-norm block and Llama's names. first_k_dense_replace is 1 on every released checkpoint: block 0's mlp is a dense MLP, every later block's a mixture. The block adds each sublayer's output to the stream and returns a tensor, so the identity is the plain sum, on the dense block as on a mixture block:

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)

Latent attention: queries and keys are wider than values

attention_queries and attention_keys are [batch, num_heads, seq, qk_head_dim], with qk_head_dim = qk_nope_head_dim + qk_rope_head_dim (128 + 64 = 192); attention_values and attention_head_outputs are head_dim wide, which nnterp reads from v_head_dim (128). transformers sets the config's head_dim to qk_rope_head_dim (64), the rotary part's width, which no served value has. The attention interior needs attn_implementation="eager".

attn = model.layers[1].self_attn
with model.trace(prompt):
    q = attn.attention_queries.save()          # [1, 16, seq, 192] on V2-Lite
    v = attn.attention_values.save()           # [1, 16, seq, 128]

Keys and values are expanded from one compressed latent per token (kv_lora_rank, 512) by kv_b_proj, so attention_keys has num_heads heads, as many as the queries. The last qk_rope_head_dim dimensions of a key come from one projection per token, shared by every head: attention_keys[:, h, :, -64:] is the same tensor for every h, and an edit there has to be made on every head to act like an edit to the shared part.

Query compression differs between the sizes

DeepSeek-V2, V2.5 and Coder-V2 (236B) compress the queries too (q_lora_rank 1536): q_a_proj, q_a_layernorm and q_b_proj, with no q_proj module. V2-Lite and Coder-V2-Lite (16B) have q_lora_rank null: one q_proj, and no q_a_proj, q_a_layernorm or q_b_proj. attention_queries is read at the attention interface on both, the same width.

Rotary turns adjacent pairs in place

The rotary part is rotated as complex numbers over adjacent pairs of dimensions (2i, 2i + 1), in place, so the last 64 dimensions of attention_queries and attention_keys keep the order of the projection's output.

The query scale carries YaRN's factor

Every released checkpoint extends the context with YaRN (rope_scaling factor 40), and the attention multiplies the scores by qk_head_dim ** -0.5 times a factor set by mscale_all_dim. Where mscale_all_dim is 0.707 (V2, V2-Chat and the Lite models) the factor is 1.59 and self_attn.scaling is 0.1147. Where it is 1.0 (V2-Chat-0628, V2.5, V2.5-1210 and the 236B Coder-V2 models) the factor is 1.87 and self_attn.scaling is 0.1352. Recompute attention_scores from the queries and keys with self_attn.scaling, not with qk_head_dim ** -0.5.

The mixture: a softmax router whose weights are not renormalized

Every block after the first routes each token to num_experts_per_tok (6) of n_routed_experts experts (64 on Lite, 160 on the 236B models) and runs it through n_shared_experts (2) shared experts, one MLP 2 × moe_intermediate_size wide. router_logits are the logits before the softmax. expert_weights are the chosen experts' softmax probabilities times routed_scaling_factor, never renormalized over the chosen six. On the Lite models the factor is 1.0, so a token's weights sum to less than one. On the 236B models it is 16.0. The 236B models route group_limited_greedy (topk_method): the experts sit in n_group 8 groups, each group scored by its best expert, and a token chooses its six among the experts of its best topk_group 3 groups. The Lite models route greedy, over every expert.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()
    w = moe.expert_weights.save()
    idx = moe.expert_indices.save()
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()
    out = moe.mlp_output.save()

scale = model.config.routed_scaling_factor
torch.testing.assert_close(w, logits.softmax(-1).gather(-1, idx) * scale)   # greedy and group-limited
torch.testing.assert_close(routed + shared, out)

The shared experts run after the routed experts. Ablating an expert is zeroing the weight of every slot that chose it. The shared experts are untouched:

with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == 3, 0)
    ablated = model.logits.save()

Block 0's MLP is dense: support() reports every mixture value missing there.

From the family module

04 · nnterp/families/deepseek_v2.py

DeepSeek-V2 (DeepseekV2ForCausalLM).

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. Multi-head latent attention: queries and keys are qk_head_dim wide (qk_nope_head_dim + qk_rope_head_dim), values v_head_dim, and the interface sees num_heads key/value heads whatever num_key_value_heads says. The first first_k_dense_replace blocks have a dense MLP, the rest a mixture of experts.