nnterp · encyclopedia bitnet · nnterp.families.bitnet

family bitnet

BitNet b1.58

Llama's tree whose projections run ternary weights on 8-bit inputs, whose sublayers each RMS-norm the tensor entering their output projection (attn_sub_norm, ffn_sub_norm), and whose MLP gates with a squared ReLU.

Quantized projectionsSub-normsSquared ReLUGrouped-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× 30
→
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 BitNetForCausalLMlayoutwhat 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] BitNetDecoderLayerlayoutwhat 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 BitNetAttentionlayoutwhat 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 BitNetMLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output

Printout

What print(model) shows on microsoft/bitnet-b1.58-2B-4T: the native modules, standard names and every standard value with its layout.

BitNetForCausalLM(
  (lm_head): Linear(in_features=2560, out_features=128256, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(128256, 2560)
  (layers): ModuleList(
    (0-29): 30 x BitNetDecoderLayer(
      (self_attn): BitNetAttention(
        (q_proj): Linear(in_features=2560, out_features=2560, bias=False)
        (k_proj): Linear(in_features=2560, out_features=640, bias=False)
        (v_proj): Linear(in_features=2560, out_features=640, bias=False)
        (o_proj): Linear(in_features=2560, out_features=2560, bias=False)
        (attn_sub_norm): BitNetRMSNorm((2560,), 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): BitNetMLP(
        (gate_proj): Linear(in_features=2560, out_features=6912, bias=False)
        (up_proj): Linear(in_features=2560, out_features=6912, bias=False)
        (down_proj): Linear(in_features=6912, out_features=2560, bias=False)
        (act_fn): ReLUSquaredActivation()
        (ffn_sub_norm): BitNetRMSNorm((6912,), eps=1e-05)
        (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): BitNetRMSNorm((2560,), eps=1e-05)
      (post_attention_layernorm): BitNetRMSNorm((2560,), 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): BitNetRMSNorm((2560,), 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
)

Config

model.num_layers30
model.hidden_size2560
model.num_heads20
model.num_kv_heads5
model.head_dim128
model.qk_head_dim128
model.vocab_size128256
model.intermediate_size6912
architecturesBitNetForCausalLM
num_hidden_layers30
hidden_size2560
intermediate_size6912
num_attention_heads20
num_key_value_heads5
vocab_size128256
max_position_embeddings4096
hidden_actrelu2
rms_norm_eps1e-05
tie_word_embeddingsTrue
dtypebfloat16
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

a   = attention(q_proj, k_proj, v_proj of input_layernorm(x))   # heads, before o_proj
h   = x + o_proj(attn_sub_norm(a))
m   = relu(gate_proj(z)) ** 2 * up_proj(z),  z = post_attention_layernorm(h)
out = h + down_proj(ffn_sub_norm(m))

Four RMSNorms per block. input_layernorm and post_attention_layernorm are Llama's pre-norms; self_attn.attn_sub_norm (over hidden_size, all heads at once) and mlp.ffn_sub_norm (over intermediate_size) sit inside the modules, each right before the output projection. Every projection inside the block (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj) is quantized; embed_tokens and lm_head are not.

The contributions are the output projections' outputs

attention_output is self_attn.output[0], which is o_proj.output, and mlp_output is mlp.output, down_proj.output: the sub-norms come before them, not after, so the contributions are the modules' outputs and the identity holds exactly. attention_head_outputs is read at the attention interface, before attn_sub_norm: flattened over its heads it equals attn_sub_norm.input.

attn = model.layers[1].self_attn
with model.trace(prompt):
    heads = attn.attention_head_outputs.save()   # [batch, seq, num_heads, head_dim]
    normed = attn.attn_sub_norm.input.save()

torch.equal(heads.flatten(-2), normed)   # True

Scaling inside a sublayer is undone by the sub-norm

attn_sub_norm divides by the RMS over every head's output together, so scaling all of attention_head_outputs changes nothing: on 2B-4T-bf16, doubling them at block 5 leaves attention_output bit-identical, and so does doubling ffn_sub_norm.input for mlp_output. A change to one head moves the others, which share the RMS: zeroing head 3 at block 5 changes the other heads' slices of attn_sub_norm.output by 0.5%. Scale or ablate a sublayer's effect at attention_output or mlp_output; a head's effect is not a separate term of attention_output.

with model.trace(prompt):
    model.layers[1].self_attn.attention_head_outputs[:] *= 2   # no effect

The projections quantize their inputs and weights

On a real load every projection in the block is an AutoBitLinear. Its forward rounds its input per token to 255 levels (round(x · 127 / max|x|)) and multiplies it by ternary weights. On the bf16 checkpoint the weights are stored in bfloat16 and rounded to {−1, 0, 1} · mean|W| in every forward. On the packed checkpoint they are stored ternary, four to a byte, with a weight_scale. q_proj.input is the unrounded normed stream. On 2B-4T-bf16 q_proj.output equals the linear map of the rounded input by the ternarized weight exactly, and differs from F.linear(input, weight) by 2.4 times its own norm. An edit at a projection's input is rounded with the rest of the row, and a weight edit on the bf16 checkpoint is re-ternarized: adding 0.1 · mean|W| to one q_proj weight at block 1 leaves the logits bit-identical. The page's printout comes from a meta build, where the projections are plain Linear: transformers replaces them with AutoBitLinear when the weights load.

The stream is large, and the input norms' gains small

On 2B-4T-bf16 the residual stream at block 10 has a norm of 20,000 to 35,000 per position and 161,000 at position 0, <|begin_of_text|>. The pre-norms' gains are small to match: input_layernorm.weight at block 5 averages 0.015, the final norm.weight 0.10. A steering vector sized from another family's stream is lost in this one; take its scale from layer_output.norm(dim=-1) on this checkpoint.

The MLP gates with a squared ReLU

mlp.act_fn is relu2 on the gate: a neuron is relu(gate_proj(z))² · up_proj(z), exactly zero wherever the gate's pre-activation is negative. At block 5 of 2B-4T-bf16, 62% of gate_proj.output entries are at or below zero on a test prompt. ffn_sub_norm.input is the gated product, [batch, seq, 6912].

Loading

The six attention interior values need attn_implementation="eager". There is no softcap, window or sink, and the default sdpa load computes the same function up to rounding: on 2B-4T-bf16 in bfloat16 the two loads' logits differ by up to 0.75 with the same top token at every position (" Paris" 0.822 eager, 0.807 sdpa). 20 query heads share 5 key/value heads, 128 wide, so an edit to key/value head j of attention_keys reaches query heads 4j to 4j + 3. The rotary is Llama's rotate_half with rope_theta 500,000, over 4096 positions.

The two checkpoints

bitnet-b1.58-2B-4T stores its weights packed (quantization_mode offline) and bitnet-b1.58-2B-4T-bf16 in bfloat16 for training (online); the configs differ only there. The packed one multiplies by its stored ternary weights and weight_scale, the bf16 one ternarizes its weights in each forward. The tokenizer is Llama 3's: <|begin_of_text|> is prepended and " Paris" is one token. The chat template writes User: ...<|eot_id|>Assistant:.

From the family module

04 · nnterp/families/bitnet.py

BitNet b1.58 (BitNetForCausalLM).

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 per-head outputs (attn_sub_norm) after the interface and before o_proj, so attention_head_outputs is before that norm; the MLP norms its gated product (ffn_sub_norm) before down_proj. The released checkpoints (quant_method: bitnet) load their projections as transformers' AutoBitLinear; a meta build and the pinned tiny show plain nn.Linear.