nnterp · encyclopedia cohere · nnterp.families.cohere

family cohere

Command RCommand R+

A parallel block whose one LayerNorm, without a bias, feeds both the attention and the MLP; rotary turns adjacent pairs of dimensions, Command R+ norms queries and keys per head, and the logits are the head's output times logit_scale.

Parallel blockLayerNormInterleaved rotaryQuery/key normsScaled logitsGrouped-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
→
logitsscaled
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 CohereForCausalLMlayoutwhat 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] CohereDecoderLayerlayoutwhat 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 scalar Unavailable: a parallel block: the attention and the MLP read the same input, so there is no stream between them

nowhere: the family does not serve it
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP: the one norm feeds both sublayers

the child module input_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 CohereAttentionlayoutwhat 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 CohereMLPlayoutwhat 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 unsloth/c4ai-command-r-08-2024: the native modules, standard names and every standard value with its layout.

CohereForCausalLM(
  (lm_head): Linear(in_features=8192, out_features=256000, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(256000, 8192, padding_idx=0)
  (layers): ModuleList(
    (0-39): 40 x CohereDecoderLayer(
      (self_attn): CohereAttention(
        (q_proj): Linear(in_features=8192, out_features=8192, bias=False)
        (k_proj): Linear(in_features=8192, out_features=1024, bias=False)
        (v_proj): Linear(in_features=8192, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=8192, 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): CohereMLP(
        (gate_proj): Linear(in_features=8192, out_features=24576, bias=False)
        (up_proj): Linear(in_features=8192, out_features=24576, bias=False)
        (down_proj): Linear(in_features=24576, out_features=8192, 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): CohereLayerNorm()
      (layer_mid): Unavailable: a parallel block: the attention and the MLP read the same input, so there is no stream between them
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP: the one norm feeds both sublayers
      (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
      (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): CohereLayerNorm()
  (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_size8192
model.num_heads64
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size256000
model.intermediate_size24576
architecturesCohereForCausalLM
num_hidden_layers40
hidden_size8192
intermediate_size24576
num_attention_heads64
num_key_value_heads8
vocab_size256000
max_position_embeddings131072
hidden_actsilu
layer_norm_eps1e-05
tie_word_embeddingsTrue
dtypebfloat16
attention_biasFalse
logit_scale0.0625
use_qk_normFalse

The Notes

03 · what to know before you run

One LayerNorm feeds both sublayers

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

input_layernorm is the block's only norm: there is no post_attention_layernorm. Its one output tensor is passed to the attention and then to the MLP, so self_attn.input and mlp.input are the same tensor, not two equal ones. An in-place edit of self_attn.input therefore also changes what the MLP reads; an assignment replaces only the attention's argument. To change what both read, assign input_layernorm.output.

with model.trace(prompt):
    model.layers[1].self_attn.input[:, -1] = 0      # in place
    mlp_in = model.layers[1].mlp.input.save()       # its last row is zero too

The contributions are the two modules' outputs

attention_output is o_proj's output and mlp_output is down_proj's; neither has a bias. The block computes residual + attn + mlp, so that sum reproduces layer_output bit for bit, in float32 and in bfloat16. Summed in another order (attn + mlp + x) it differs by rounding. Nothing inside a block connects the two sublayers: zeroing attention_output leaves that block's mlp.input bit-identical, and the first MLP that reads block i's attention is block i + 1's.

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()

assert torch.equal(x + attn + mlp, out)           # the block's own order

The norms are LayerNorms with a weight and no bias

CohereLayerNorm casts to float32, subtracts the mean, divides by the standard deviation (layer_norm_eps, 1e-5) and multiplies by weight; there is no bias term, and the result is cast back to the model's dtype. input_layernorm and the final norm are this class. Folding a norm into the next matrix takes weight alone, and the mean subtraction stays as a projection that removes the all-ones direction.

Rotary turns adjacent pairs

Cohere's rotary pairs dimension 2i with 2i + 1 (its rotate_half takes the even and odd dimensions, and each frequency in the cos/sin table appears twice in a row), where Llama pairs i with i + head_dim / 2. The rotation covers the whole head, in float32, cast back to the model's dtype. attention_queries and attention_keys are read after it. A comparison with a rotate_half family's queries, or rotary code written for one, needs the head reordered, evens then odds. The base (rope_theta) is 8,000,000 on Command R v01, 4,000,000 on Command R 08-2024, 75,000,000 on Command R+ and 8,000,000 on Command R+ 08-2024.

hd = model.head_dim
perm = torch.cat([torch.arange(0, hd, 2), torch.arange(1, hd, 2)])

with model.trace(prompt):
    q = model.layers[1].self_attn.attention_queries.save()

q_half = q[..., perm]          # in rotate_half's order

Command R+ norms queries and keys per head

Where use_qk_norm is true (Command R+ and Command R+ 08-2024), the attention has q_norm and k_norm, CohereLayerNorms over each head's head_dim vector. Their weights are [num_heads, head_dim] and [num_kv_heads, head_dim]: every head has its own gain. They act on the projections before the rotary, on [batch, seq, num_heads, head_dim] tensors, and attention_queries and attention_keys are read after both. Command R and Command R 08-2024 have no q_norm or k_norm.

with model.trace(prompt):
    normed = model.layers[1].self_attn.q_norm.output.save()
    # [batch, seq, num_heads, head_dim], before the rotary

Attention

Command R v01 has 64 query heads and 64 key/value heads (no grouping); Command R 08-2024 has 64 and 8, and Command R+ 96 and 8, so on those two an edit to a key or value head reaches 8 or 12 query heads. attention_keys and attention_values are served before repeat_kv. The query scale is head_dim ** -0.5 (128 on every size). The attention runs the shared attention interface, so the values read inside it (marked ⚠) need attn_implementation="eager".

The logits are the head's output times logit_scale

lm_head is the embedding matrix (tie_word_embeddings), and the model multiplies its output by config.logit_scale: 0.0625 on Command R v01 and 08-2024, 0.8333 on both Command R+ releases. logits is the scaled tensor and lm_head.output the unscaled one. The family's project_on_vocab is the final norm, lm_head and the scale, so a logit lens on the last block equals logits exactly. A direct logit attribution through lm_head alone is off by that constant factor.

with model.trace(prompt):
    resid = model.layers[-1].layer_output.save()
    raw = model.lm_head.output.save()
    logits = model.logits.save()

assert torch.equal(logits, raw * model.config.logit_scale)
assert torch.equal(model.project_on_vocab(resid), logits)

The embeddings are not scaled

embed_tokens is a plain lookup over 256,000 tokens, so token_embeddings equals layers[0].input. Because the weight is shared with lm_head, an edit to embed_tokens.weight also edits the unembedding.

From the family module

04 · nnterp/families/cohere.py

Cohere (CohereForCausalLM, Command-R).

model.{embed_tokens, layers[i].{input_layernorm, self_attn, mlp}, norm} and lm_head: Llama's names. The block is parallel: one LayerNorm (not RMS) feeds both sublayers and the block returns x + attn + mlp, so the contributions are the two modules' outputs and the base holds. There is no post_attention_layernorm. The attention runs the shared interface, after an optional per-head q_norm / k_norm (use_qk_norm) and an interleaved rotary.

The model multiplies the head's output by config.logit_scale to make the logits, so the family defines project_on_vocab with that step, and the logit lens on the last block equals logits.

Aya Vision 32B (aya_vision with a cohere text config) loads through this family. It binds no tower keys, so that load has no vision or projector.