family dots1
Llama's pre-norm block with each query and key head RMS-normed, a dense MLP on block 0, and on every block after it a mixture whose experts are chosen by a sigmoid plus a selection bias that the weights leave out, beside a shared expert every token runs through.
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.
identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output
⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.
What print(model) shows on rednote-hilab/dots.llm1.base: the native modules, standard names and every standard value with its layout.
Dots1ForCausalLM(
(lm_head): Linear(in_features=4096, out_features=152064, bias=False)
(generator): Generator(
(streamer): Streamer()
)
(embed_tokens): Embedding(152064, 4096)
(layers): ModuleList(
(0): Dots1DecoderLayer(
(self_attn): Dots1Attention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(q_norm): Dots1RMSNorm((128,), eps=1e-05)
(k_norm): Dots1RMSNorm((128,), eps=1e-05)
(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): Dots1MLP(
(gate_proj): Linear(in_features=4096, out_features=10944, bias=False)
(up_proj): Linear(in_features=4096, out_features=10944, bias=False)
(down_proj): Linear(in_features=10944, out_features=4096, 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): Dots1RMSNorm((4096,), eps=1e-05)
(post_attention_layernorm): Dots1RMSNorm((4096,), eps=1e-05)
(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-61): 61 x Dots1DecoderLayer(
(self_attn): Dots1Attention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(q_norm): Dots1RMSNorm((128,), eps=1e-05)
(k_norm): Dots1RMSNorm((128,), eps=1e-05)
(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): Dots1MoE(
(experts): Dots1Experts(
(act_fn): SiLUActivation()
)
(router/gate): Dots1TopkRouter()
(shared_experts): Dots1MLP(
(gate_proj): Linear(in_features=4096, out_features=2816, bias=False)
(up_proj): Linear(in_features=4096, out_features=2816, bias=False)
(down_proj): Linear(in_features=2816, out_features=4096, 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): Dots1RMSNorm((4096,), eps=1e-05)
(post_attention_layernorm): Dots1RMSNorm((4096,), eps=1e-05)
(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): Dots1RMSNorm((4096,), eps=1e-05)
(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
)What print(model) shows on rednote-hilab/dots.llm1.inst: the native modules, standard names and every standard value with its layout.
Dots1ForCausalLM(
(lm_head): Linear(in_features=4096, out_features=152064, bias=False)
(generator): Generator(
(streamer): Streamer()
)
(embed_tokens): Embedding(152064, 4096)
(layers): ModuleList(
(0): Dots1DecoderLayer(
(self_attn): Dots1Attention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(q_norm): Dots1RMSNorm((128,), eps=1e-05)
(k_norm): Dots1RMSNorm((128,), eps=1e-05)
(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): Dots1MLP(
(gate_proj): Linear(in_features=4096, out_features=10944, bias=False)
(up_proj): Linear(in_features=4096, out_features=10944, bias=False)
(down_proj): Linear(in_features=10944, out_features=4096, 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): Dots1RMSNorm((4096,), eps=1e-05)
(post_attention_layernorm): Dots1RMSNorm((4096,), eps=1e-05)
(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-61): 61 x Dots1DecoderLayer(
(self_attn): Dots1Attention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(q_norm): Dots1RMSNorm((128,), eps=1e-05)
(k_norm): Dots1RMSNorm((128,), eps=1e-05)
(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): Dots1MoE(
(experts): Dots1Experts(
(act_fn): SiLUActivation()
)
(router/gate): Dots1TopkRouter()
(shared_experts): Dots1MLP(
(gate_proj): Linear(in_features=4096, out_features=2816, bias=False)
(up_proj): Linear(in_features=4096, out_features=2816, bias=False)
(down_proj): Linear(in_features=2816, out_features=4096, 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): Dots1RMSNorm((4096,), eps=1e-05)
(post_attention_layernorm): Dots1RMSNorm((4096,), eps=1e-05)
(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): Dots1RMSNorm((4096,), eps=1e-05)
(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
)h = x + self_attn(input_layernorm(x)) # q_norm, k_norm per head inside
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 both checkpoints: block
0 has a dense MLP 10944 wide, and the other 61 a mixture. Nothing norms a sublayer's output, so
attention_output is o_proj's output and mlp_output the mixture's, shared expert included.
The identity is the plain sum on both kinds of 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)
q_norm and k_norm are RMSNorms over each head's 128 dimensions, applied before the rotary;
attention_queries and attention_keys are read after both, and an edit there reaches the
attention as written. An edit to one head's slice of q_proj.output stays in that head, and the
norm rescales it. There are 32 key/value heads for 32 query heads, so no grouping. Every block
attends over the whole prefix (sliding_window is null), with rotary at rope_theta 1e7. The
attention interior needs attn_implementation="eager".
128 routed experts, 6 per token. router_logits are the router's float32 logits, before the
sigmoid and before the selection bias (router.e_score_correction_bias). The bias is added to
the sigmoids only to choose; n_group and topk_group are 1, so every expert competes.
expert_weights are the chosen experts' sigmoids without the bias, renormalized to sum to one
(norm_topk_prob) and multiplied by routed_scaling_factor 2.5, so they sum to 2.5 at every token:
moe = model.layers[1].mlp
with model.trace(prompt):
logits = moe.router_logits.save()
w = moe.expert_weights.save()
idx = moe.expert_indices.save()
chosen = logits.sigmoid().gather(-1, idx)
torch.testing.assert_close(w, chosen / chosen.sum(-1, keepdim=True) * 2.5)
A written router_logits moves the choice and the weights; a written bias moves only the choice.
n_shared_experts is 2, merged into one Dots1MLP (mlp.shared_experts) of width 2 × 1408. It
reads the mixture's input and runs after the routed experts, so in one trace read routed_output
before shared_expert_output. Their sum is mlp_output:
with model.trace(prompt):
routed = moe.routed_output.save()
shared = moe.shared_expert_output.save()
out = moe.mlp_output.save()
torch.testing.assert_close(routed + shared, out)
Zero the weight of every slot that chose it. The token's other weights keep their values, so its
weights then sum to less than 2.5, and mlp_output changes on exactly the tokens that chose it:
with model.trace(prompt):
idx = moe.expert_indices.save()
clean = moe.mlp_output.save()
with model.trace(prompt):
moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == 3, 0)
ablated = moe.mlp_output.save()
assert torch.equal((ablated != clean).any(-1), (idx == 3).any(-1))
Block 0 has a dense MLP: support() reports every mixture value missing there.
dots.llm1 (Dots1ForCausalLM).
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.
The attention norms its queries and keys per head (q_norm, k_norm)
before the rotary embedding. The first first_k_dense_replace blocks have a
dense MLP, the rest a mixture of experts that adds a shared expert
(mlp.shared_experts, itself a Dots1MLP) before it returns; the experts
are moe_intermediate_size wide.