nnterp · encyclopedia cohere2 · nnterp.families.cohere2

family cohere2

Command R7BCommand A

Command R's parallel block on one bias-free LayerNorm, with three sliding-window blocks to each full one; only the sliding blocks apply rotary, and the logits are the head's output times logit_scale.

Parallel blockLayerNormSliding windowBlocks without rotaryInterleaved rotaryScaled 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. The slider picks a block; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 32
→
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 Cohere2ForCausalLMlayoutwhat 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] Cohere2DecoderLayerlayoutwhat 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 Cohere2Attentionlayoutwhat 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 Cohere2MLPlayoutwhat 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 mlx-community/c4ai-command-r7b-12-2024-bf16: the native modules, standard names and every standard value with its layout.

Cohere2ForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=256000, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(256000, 4096, padding_idx=0)
  (layers): ModuleList(
    (0-31): 32 x Cohere2DecoderLayer(
      (self_attn): Cohere2Attention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, 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): Cohere2MLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, 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): Cohere2LayerNorm()
      (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): Cohere2LayerNorm()
  (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_layers32
model.hidden_size4096
model.num_heads32
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size256000
model.intermediate_size14336
architecturesCohere2ForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size14336
num_attention_heads32
num_key_value_heads8
head_dim128
vocab_size256000
max_position_embeddings8192
sliding_window4096
hidden_actsilu
layer_norm_eps1e-05
tie_word_embeddingsTrue
dtypebfloat16
attention_biasFalse
logit_scale0.25

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)

The block is Command R's: input_layernorm is its only norm, a LayerNorm with a weight and no bias, computed in float32. 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: an in-place edit of self_attn.input also changes what the MLP reads, and assigning input_layernorm.output changes what both read.

The contributions are the two modules' outputs

attention_output is o_proj's output and mlp_output is down_proj's. The block computes residual + attn + mlp, so that sum reproduces layer_output bit for bit; summed in another order it differs by rounding. Zeroing attention_output leaves that block's mlp.input unchanged.

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

Three sliding blocks to one full block, and only the sliding blocks apply the rotary

config.layer_types is full_attention on every fourth block (3, 7, ..., 31 of Command R7B's 32 blocks) and sliding_attention elsewhere. A sliding block attends over a window of 4096 tokens; on a prompt no longer than the window both masks are the same causal mask. To tell the kinds apart outside a trace, read model.layers[i].self_attn._module.sliding_window: the window on a sliding block, None on a full one.

Only the sliding blocks apply the rotary. A full block has no position encoding: its attention_queries and attention_keys are q_proj's and k_proj's outputs split into heads, and only the causal mask orders the tokens there. On a sliding block they are read after the rotary, which turns adjacent pairs of dimensions (2i with 2i + 1) at base 50,000.

full = model.config.layer_types.index("full_attention")
with model.trace(prompt):
    k_proj = model.layers[full].self_attn.k_proj.output.save()
with model.trace(prompt):
    keys = model.layers[full].self_attn.attention_keys.save()

heads = k_proj.view(*k_proj.shape[:2], -1, model.head_dim).transpose(1, 2)
assert torch.equal(heads, keys)                   # no rotary on a full block

Attention

Command R7B has 32 query heads and 8 key/value heads, Command A 96 and 8, so an edit to a key or value head reaches 4 or 12 query heads; attention_keys and attention_values are served before repeat_kv. The query scale is head_dim ** -0.5 (128). Both kinds of block make the same attention-interface call, so every interior value resolves on every block, under attn_implementation="eager" (the values marked ⚠). There are no query or key norms.

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.25 on Command R7B and Command A. logits is the scaled tensor and lm_head.output the unscaled one. The family imports Command R's project_on_vocab (the final norm, lm_head, the scale), so a logit lens on the last block equals logits exactly.

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)

Aya Vision and Command A Vision wrap this text model

aya_vision and cohere2_vision checkpoints load with task="image-text-to-text". nnterp picks this family from config.text_config and binds embed_tokens, layers and norm under model.language_model, with lm_head at the root. The vision encoder is SigLIP at model.vision_tower. The image features replace the image tokens' embeddings after embed_tokens, so at those positions token_embeddings is the placeholder token's embedding and layers[0].input holds the projected image features; at text positions the two are equal. embed_tokens runs before the vision encoder, so read token_embeddings before any vision value in a trace.

From the family module

04 · nnterp/families/cohere2.py

Cohere 2 (Cohere2ForCausalLM, Command-R7B / Command-A).

Cohere's tree and parallel block (x + attn + mlp from one LayerNorm), with sliding-window and full attention layers interleaved (config.layer_types). Only the sliding layers apply the rotary; the full layers use no position encoding. Both kinds make the same attention_interface call, so every interior value resolves on every block, and on a full layer attention_queries / attention_keys are the unrotated projections.

The model multiplies the head's output by config.logit_scale, as Cohere's does, so the family takes Cohere's project_on_vocab.

Aya Vision and Cohere2-Vision put a SigLIP tower at model.vision_tower, which is vision (its post_layernorm over the patches is vision.norm), and a pixel-shuffling projector model.multi_modal_projector, which is projector; vision.image_features is read at the scatter (ImageScatter). Aya Vision's projector reads the last block's layer_output before vision.norm (vision_feature_layer -1, strategy "full") and norms after its pixel shuffle; Cohere2-Vision's reads vision.tower_output and has no norm. Aya Vision 32B's text config is cohere, so it loads through that family, which binds no tower.