nnterp · encyclopedia seed_oss · nnterp.families.seed_oss

family seed_oss

Seed-OSS

Llama's block with a bias on the query, key and value projections, and query heads that together are twice as wide as the stream.

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× 64
→
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 SeedOssForCausalLMlayoutwhat 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] SeedOssDecoderLayerlayoutwhat 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 SeedOssAttentionlayoutwhat 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 SeedOssMLPlayoutwhat 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 ByteDance-Seed/Seed-OSS-36B-Base: the native modules, standard names and every standard value with its layout.

SeedOssForCausalLM(
  (lm_head): Linear(in_features=5120, out_features=155136, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(155136, 5120, padding_idx=1)
  (layers): ModuleList(
    (0-63): 64 x SeedOssDecoderLayer(
      (self_attn): SeedOssAttention(
        (q_proj): Linear(in_features=5120, out_features=10240, bias=True)
        (k_proj): Linear(in_features=5120, out_features=1024, bias=True)
        (v_proj): Linear(in_features=5120, out_features=1024, bias=True)
        (o_proj): Linear(in_features=10240, out_features=5120, 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): SeedOssMLP(
        (gate_proj): Linear(in_features=5120, out_features=27648, bias=False)
        (up_proj): Linear(in_features=5120, out_features=27648, bias=False)
        (down_proj): Linear(in_features=27648, out_features=5120, 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): SeedOssRMSNorm((5120,), eps=1e-06)
      (post_attention_layernorm): SeedOssRMSNorm((5120,), 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): SeedOssRMSNorm((5120,), 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_layers64
model.hidden_size5120
model.num_heads80
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size155136
model.intermediate_size27648
architecturesSeedOssForCausalLM
num_hidden_layers64
hidden_size5120
intermediate_size27648
num_attention_heads80
num_key_value_heads8
head_dim128
vocab_size155136
max_position_embeddings524288
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypebfloat16
attention_biasTrue

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))
out = h + mlp(post_attention_layernorm(h))

Llama's block and Llama's names: two RMSNorms, each before its sublayer, none after. The attention and the MLP each end in a residual dropout inside the module (residual_dropout, 0.1 in every config), after o_proj and after down_proj. A loaded model is in eval mode, where the dropout is the identity.

The contributions are the modules' outputs

attention_output is self_attn.output[0] and mlp_output is mlp.output, added to the stream with nothing in between, and the identity holds exactly in float32 on the pinned checkpoint:

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.equal(x + attn + mlp, out)   # True

Under model.train() (fine-tuning an adapter through the model) the residual dropout is active. The contributions are read after it, so attention_output differs from o_proj.output, and the identity still holds: on the pinned checkpoint with residual_dropout=0.5, about half of attention_output's entries are zero and x + attn + mlp == out stays exact.

The attention is twice as wide as the stream

The 36B checkpoints have hidden_size 5120 and 80 query heads of head_dim 128, so q_proj maps 5120 to 10240 and o_proj maps 10240 back to 5120. attention_head_outputs is [batch, seq, 80, 128], and the concatenated heads are twice the stream's width: a head's output is not a slice of a stream vector until o_proj has mapped it.

There are 8 key/value heads, so attention_keys and attention_values are [batch, 8, seq, 128] and query head h reads key/value head h // 10: an edit to key/value head j reaches query heads 10j to 10j + 9.

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

Queries, keys and values carry a bias

q_proj, k_proj and v_proj add a bias (attention_bias); o_proj does not (attention_out_bias is false), nor do the three MLP projections (mlp_bias). The bias is added before the rotary embedding, so it is rotated with the rest of the query and key. With the attention's input zeroed, attention_keys is the rotated k_proj.bias: equal to the bias at position 0 and rotated at every other position (checked on the pinned checkpoint after setting its zero bias to ones).

with model.trace(prompt):
    model.layers[1].self_attn.input[:] = 0
    k = model.layers[1].self_attn.attention_keys.save()   # rope(k_proj.bias), not zero

To remove the attention's effect, ablate attention_output itself.

The default load runs the same attention as eager

A default load runs sdpa, and the six attention interior values report unavailable; attn_implementation="eager" serves them. There is no softcap, window or sink, so the two compute the same function. The query scale is head_dim ** -0.5. The rotary is the plain one with rope_theta 10000000, over a max_position_embeddings of 524288.

The readout and the tokenizer

model.logits is model.lm_head.output, so project_on_vocab is lm_head(norm(hidden)), and lm_head has its own weight (tie_word_embeddings is false). The tokenizer prepends no BOS (<seed:bos>, id 0, is not added): position 0 is the text's first token. The tokenizer has 155121 tokens against a vocab_size of 155136, so the last 15 ids get logits but decode to nothing.

Three checkpoints, one shape

Seed-OSS-36B-Base, Seed-OSS-36B-Base-woSyn and Seed-OSS-36B-Instruct have the same config: 64 blocks, the same widths, heads and rotary. -woSyn is the base model pretrained without synthetic instruction data, released for research that needs a base model untouched by it.

From the family module

04 · nnterp/families/seed_oss.py

Seed-OSS (SeedOssForCausalLM).

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 output and the MLP output each pass a residual dropout inside their module, which is the identity in eval.