nnterp · encyclopedia dbrx · nnterp.families.dbrx

family dbrx

DBRX

A pre-norm block whose attention sits with both norms inside norm_attn_norm, queries, keys and values from one clamped Wqkv projection, and a mixture of 16 experts, 4 per token, on every block.

LayerNormMixture of expertsGrouped-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 40
→
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 DbrxForCausalLMlayoutwhat 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] DbrxBlocklayoutwhat 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].mlp DbrxFFNlayoutwhat 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

the child module router.layer's output
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 scalar Unavailable: DBRX's experts loop over the experts in their own forward; no tensor holds each slot's output

nowhere: the family does not serve it
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output
model.layers[i].self_attn DbrxAttentionlayoutwhat 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

Printout

What print(model) shows on Undi95/dbrx-base: the native modules, standard names and every standard value with its layout.

DbrxForCausalLM(
  (lm_head): Linear(in_features=6144, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 6144)
  (layers): ModuleList(
    (0-39): 40 x DbrxBlock(
      (norm_attn_norm): DbrxNormAttentionNorm(
        (norm_1): LayerNorm((6144,), eps=1e-05, elementwise_affine=True, bias=False)
        (attn): DbrxAttention(
          (Wqkv): Linear(in_features=6144, out_features=8192, bias=False)
          (o_proj/out_proj): Linear(in_features=6144, out_features=6144, 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
        )
        (norm_2): LayerNorm((6144,), eps=1e-05, elementwise_affine=True, bias=False)
      )
      (mlp/ffn): DbrxFFN(
        (router): DbrxRouter(
          (layer): Linear(in_features=6144, out_features=16, bias=False)
        )
        (experts): DbrxExperts(
          (mlp): DbrxExpertGLU(
            (activation_fn): SiLUActivation()
          )
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_outputs): Unavailable: DBRX's experts loop over the experts in their own forward; no tensor holds each slot's output
        (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
        (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
        (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
      )
      (self_attn): DbrxAttention(
        (Wqkv): Linear(in_features=6144, out_features=8192, bias=False)
        (o_proj/out_proj): Linear(in_features=6144, out_features=6144, 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
      )
      (input_layernorm): LayerNorm((6144,), eps=1e-05, elementwise_affine=True, bias=False)
      (post_attention_layernorm): LayerNorm((6144,), eps=1e-05, elementwise_affine=True, bias=False)
      (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): LayerNorm((6144,), eps=1e-05, elementwise_affine=True, bias=False)
  (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_layers40
model.hidden_size6144
model.num_heads48
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size100352
model.intermediate_size10752
architecturesDbrxForCausalLM
num_key_value_heads8
vocab_size100352
tie_word_embeddingsFalse
dtypebfloat16
d_model6144
n_heads48
n_layers40
max_seq_len32768

The Notes

03 · what to know before you run

The block, in order

h, a = norm_attn_norm(x)    # a = norm_2(h), h = x + attn(norm_1(x))
out  = h + ffn(a)           # 4 of 16 experts per token

The native tree is transformer.blocks[i].{norm_attn_norm.{norm_1, attn, norm_2}, ffn}, and nnterp names them input_layernorm, self_attn, post_attention_layernorm and mlp on the block, so a recipe reads DBRX like Llama. norm_attn_norm norms, runs the attention, adds the residual, norms again, and returns two tensors: the stream after the attention and the FFN's input. The block adds the FFN's output to that stream. The attention module's own output is its contribution, because norm_attn_norm adds the residual after it, and the identity is the plain sum:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mid = model.layers[1].norm_attn_norm.output.save()   # (h, a)
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

torch.testing.assert_close(mid[0], x + attn)
torch.testing.assert_close(x + attn + mlp, out)

One projection gives queries, keys and values, clamped

self_attn.Wqkv maps the normed stream to [queries | keys | values], hidden_size, num_kv_heads × head_dim and num_kv_heads × head_dim wide (6144, 1024, 1024: 48 query heads and 8 key/value heads of 128). The forward clamps the whole output to [-clip_qkv, clip_qkv] (attn_config.clip_qkv, 8) and then splits it, so Wqkv.output is before the clamp and attention_queries, attention_keys and attention_values after it (queries and keys also after rotary, rope_theta 500,000). Splitting the clamped output by those widths gives the values exactly:

attn = model.layers[1].self_attn
with model.trace(prompt):
    qkv = attn.Wqkv.output.save()           # before the clamp
with model.trace(prompt):
    v = attn.attention_values.save()        # [batch, num_kv_heads, seq, head_dim]

H, KV, D = model.num_heads, model.num_kv_heads, model.head_dim
clip = model.config.attn_config.clip_qkv
q, k, val = qkv.clamp(-clip, clip).split([H * D, KV * D, KV * D], dim=-1)
torch.testing.assert_close(val.view(*val.shape[:2], KV, D).transpose(1, 2), v)

An edit to a slice of Wqkv.output passes through the clamp, so a value written beyond ±8 arrives as ±8. Edit attention_queries, attention_keys or attention_values instead: no clamp follows them. Query head h reads key/value head h // 6: an edit to one key/value head reaches 6 query heads. The output projection is out_proj, and no projection has a bias.

Load with eager for the attention interior only

The attention runs on transformers' shared interface: the six interior values need attn_implementation="eager", and the default sdpa load computes the same attention without them. The attention has no softcap, sink or window.

The norms are LayerNorms without bias

input_layernorm, post_attention_layernorm and norm are nn.LayerNorm with bias=False: each subtracts the mean over the hidden axis, divides by the standard deviation and multiplies by its weight. A shift of the stream along the all-ones direction reaches no sublayer and no logit.

The router renormalizes its top 4 to sum to one

The router is a bias-free linear map, mlp.router.layer, whose output is router_logits. The FFN takes a softmax over all 16 logits, keeps the 4 largest, and divides them by their sum (an L1 norm, ffn_config.moe_normalize_expert_weights 1), so expert_weights sum to one at every token. There is no shared expert, so routed_output is mlp_output.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()       # [batch, seq, num_experts]
    w = moe.expert_weights.save()           # [batch, seq, top_k]
    idx = moe.expert_indices.save()

top = logits.softmax(-1).topk(moe.top_k, dim=-1)
torch.testing.assert_close(top.values / top.values.sum(-1, keepdim=True), w)
assert torch.equal(top.indices, idx)

The experts are one module, mlp.experts, whose GLU weights w1 (gate), v1 (up) and w2 (down) are stacked parameters on experts.mlp, [num_experts × ffn_hidden_size, hidden_size] each, with SiLU. It loops over the experts in its own forward, so no tensor holds a slot's output and expert_outputs is unavailable. intermediate_size is one expert's width, ffn_config.ffn_hidden_size (10752); the config has no dense width. In training only, moe_jitter_eps multiplies the router's input by uniform noise in [1 - eps, 1 + eps].

Ablating an expert touches only the tokens routed to it

Zeroing expert e's weight where expert_indices == e removes that slot's term and leaves the token's other 3 weights as they were, so the token's weights then sum to less than one. mlp_output changes on exactly the tokens that chose e:

moe, e = model.layers[1].mlp, 3
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 == e, 0)
    ablated = moe.mlp_output.save()

changed = (ablated != clean).any(-1)                       # [batch, seq]
assert torch.equal(changed, (idx == e).any(-1))            # only the tokens that chose e

The readout and the checkpoints

logits is lm_head.output, with no cap or scale, and project_on_vocab on the last block's layer_output equals logits. databricks/dbrx-base and databricks/dbrx-instruct are not on the Hub; the page reads the configs of Undi95/dbrx-base and alpindale/dbrx-instruct, whose sizes agree. Their tokenizer is remote code over tiktoken: load it with trust_remote_code=True and tiktoken installed, and hand it to StandardizedTransformer as tokenizer=.

From the family module

04 · nnterp/families/dbrx.py

DBRX (DbrxForCausalLM).

transformer.{wte, blocks[i].{norm_attn_norm.{norm_1, attn, norm_2}, ffn}, norm_f} and lm_head. The attention sits inside a norm_attn_norm module that norms before it, adds the residual after it, and norms again for the FFN; the block adds the FFN's output to the state norm_attn_norm returns. So the attention module's own output is its contribution (the add happens outside it, in norm_attn_norm) and the base classes hold; the aliases reach through norm_attn_norm so the block reads like any other. The attention runs on the shared interface. The FFN is a mixture of experts (Moe) returning the hidden states. Its router's projection is a child nn.Linear, router.layer, whose output is the logits; the FFN's own route_tokens_to_experts takes a softmax top-k and p-normalizes the weights (moe_normalize_expert_weights). The experts are DBRX's own loop over experts, taking and returning [batch, seq, hidden_size], so the per-slot outputs are unavailable.