Coverage for transformer_lens/model_bridge/supported_architectures/cohere.py: 92%
85 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Cohere architecture adapter.
3Supports CohereForCausalLM models (Command-R family) with:
4- Parallel attention+MLP sharing a single input_layernorm (no post_attention_layernorm)
5- True LayerNorm (CohereLayerNorm) with weight but no bias
6- GQA (grouped-query attention) with separate Q/K/V/O projections
7- Gated SwiGLU MLP (gate_proj, up_proj, down_proj)
8- Logit scaling: output logits multiplied by config.logit_scale (default 1/16)
9- Tied embed/unembed weights by default (tie_word_embeddings=True)
10- Interleaved RoPE via CohereRotaryEmbedding (delegated to HF module)
11"""
13from typing import Any
15import torch
17from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
18from transformer_lens.model_bridge.generalized_components import (
19 EmbeddingBridge,
20 LinearBridge,
21 NormalizationBridge,
22 ParallelBlockBridge,
23 PositionEmbeddingsAttentionBridge,
24 RotaryEmbeddingBridge,
25 UnembeddingBridge,
26)
29class CohereArchitectureAdapter(ArchitectureAdapter):
30 """Architecture adapter for Cohere models (CohereForCausalLM).
32 Architectural quirks vs. standard decoder-only models:
33 - Single input_layernorm per block; NO post_attention_layernorm.
34 Attention and MLP both read the SAME normed hidden states (parallel).
35 - CohereLayerNorm is true LayerNorm (mean-subtracting), NOT RMSNorm.
36 It has a weight parameter but NO bias parameter.
37 - Logit scale: CohereForCausalLM.forward multiplies logits by logit_scale
38 (default 0.0625 = 1/16). Folded into unembed.weight via preprocess_weights.
39 - Rotary embeddings use repeat_interleave instead of cat-split (delegated to HF).
41 Optional parameters (absent from state_dict by default):
42 - blocks.{i}.attn.b_Q/b_K/b_V/b_O — no bias on projections (attention_bias=False)
43 - blocks.{i}.mlp.b_gate/b_in/b_out — no bias on MLP projections
44 - blocks.{i}.ln1.b — CohereLayerNorm has no bias
45 - ln_final.b — CohereLayerNorm has no bias
46 """
48 _testing_eager = None
50 def __init__(self, cfg: Any) -> None:
51 """Initialize the Cohere architecture adapter."""
52 super().__init__(cfg)
54 # --- Normalization ---
55 # CohereLayerNorm is true LayerNorm (subtracts mean), NOT RMSNorm.
56 # uses_rms_norm=False tells NormalizationBridge to subtract the mean.
57 self.cfg.normalization_type = "LN"
58 self.cfg.uses_rms_norm = False
59 self.cfg.final_rms = False
61 # --- Position embeddings and MLP ---
62 self.cfg.positional_embedding_type = "rotary"
63 self.cfg.gated_mlp = True
64 self.cfg.attn_only = False
66 # --- Parallel block: single norm, no post_attention_layernorm ---
67 self.cfg.parallel_attn_mlp = True
69 # --- Tokenizer: BOS is prepended by default ---
70 # CohereTokenizerFast has add_bos_token=False but HF's __call__ with
71 # add_special_tokens=True (the default) prepends BOS. Verified against
72 # trl-internal-testing/tiny-CohereForCausalLM.
73 self.cfg.default_prepend_bos = True
75 # --- GQA: n_key_value_heads ---
76 # sources/transformers.py copies num_key_value_heads generically.
77 # Re-read here to ensure it's set on cfg for _qkvo_weight_conversions.
78 n_kv = getattr(cfg, "n_key_value_heads", None)
79 if n_kv is not None: 79 ↛ 86line 79 didn't jump to line 86 because the condition on line 79 was always true
80 self.cfg.n_key_value_heads = n_kv
82 # --- Weight processing conversions ---
83 # Standard GQA-aware Q/K/V/O rearrangements (same as Llama/Qwen2).
84 # n_kv is already set on self.cfg; _qkvo_weight_conversions reads it via
85 # getattr(self.cfg, "n_key_value_heads", None) when called with no args.
86 self.weight_processing_conversions = {
87 **self._qkvo_weight_conversions(),
88 }
90 # --- Logit scale ---
91 # CohereConfig.logit_scale is typed float | None; apply explicit None-check
92 # so cfg.logit_scale is always a plain float (never None).
93 # logit_scale is not a declared field on TransformerBridgeConfig; it is a
94 # Cohere-specific dynamic attribute accessed later in preprocess_weights.
95 _ls = getattr(cfg, "logit_scale", None)
96 self.cfg.logit_scale = float(_ls) if _ls is not None else 0.0625 # type: ignore[attr-defined]
98 # --- RoPE theta (informational metadata) ---
99 # CohereRotaryEmbedding reads config.rope_parameters["rope_theta"] directly;
100 # store it in cfg.rotary_base so TL config accurately reflects the model.
101 # TransformerBridgeConfig stores rotary_base as int, matching its declared type.
102 _rope_params = getattr(cfg, "rope_parameters", None) or {}
103 if isinstance(_rope_params, dict): 103 ↛ 106line 103 didn't jump to line 106 because the condition on line 103 was always true
104 _theta = _rope_params.get("rope_theta", getattr(cfg, "default_theta", 10000.0))
105 else:
106 _theta = getattr(cfg, "default_theta", 10000.0)
107 self.cfg.rotary_base = int(_theta)
109 # --- Component mapping ---
110 # Block structure follows Falcon's parallel_attn=True, num_ln_in_parallel_attn=1
111 # mode: single ln1 feeds both attn and MLP; NO ln2.
112 # Submodule shapes follow Llama: separate q/k/v/o projections and SwiGLU MLP.
113 # Rotary and attention both delegate to HF modules, preserving Cohere's
114 # repeat_interleave RoPE convention without re-implementing it in TL.
115 self.component_mapping = {
116 # Embedding: model.embed_tokens (same root as Llama, not transformer.* like Falcon)
117 "embed": EmbeddingBridge(name="model.embed_tokens"),
118 # Rotary embedding: top-level, delegates to CohereRotaryEmbedding.
119 # Pattern matches llama.py:75 and falcon.py:154 — NOT inside blocks.
120 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
121 "blocks": ParallelBlockBridge(
122 name="model.layers",
123 submodules={
124 # Single pre-norm only — Cohere has no post_attention_layernorm.
125 # NormalizationBridge handles weight-only CohereLayerNorm correctly:
126 # it checks `hasattr(original_component, "bias") and bias is not None`
127 # before adding bias, so the missing bias attribute is silently skipped.
128 "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg),
129 # No "ln2" — parallel block, same normed input goes to attn AND mlp.
130 "attn": PositionEmbeddingsAttentionBridge(
131 name="self_attn",
132 config=self.cfg,
133 submodules={
134 "q": LinearBridge(name="q_proj"),
135 "k": LinearBridge(name="k_proj"),
136 "v": LinearBridge(name="v_proj"),
137 "o": LinearBridge(name="o_proj"),
138 },
139 requires_attention_mask=True,
140 requires_position_embeddings=True,
141 ),
142 # GatedMLPBridge: gate/in/out matches Llama's gate_proj/up_proj/down_proj.
143 # Optional use_qk_norm is handled transparently by HF's
144 # CohereAttention.forward delegation (no extra submodules needed).
145 "mlp": self._gated_mlp(),
146 },
147 ),
148 # Final LayerNorm (CohereLayerNorm, weight-only) at model.norm
149 "ln_final": NormalizationBridge(name="model.norm", config=self.cfg),
150 # Unembed: lm_head. logit_scale is folded into weight in preprocess_weights.
151 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
152 }
154 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
155 """Fold logit_scale into unembed weights before ProcessWeights runs.
157 bridge.py lines 726-732 clone unembed.weight before calling this, so
158 scaling does not affect the tied embed.weight.
159 logit_scale=1.0 is a no-op (skipped for efficiency).
160 """
161 scale: float = getattr(self.cfg, "logit_scale") # always set by __init__
162 if scale != 1.0:
163 for key in ("unembed.weight", "unembed.bias"):
164 if key in state_dict:
165 orig_dtype = state_dict[key].dtype
166 state_dict[key] = (state_dict[key].float() * scale).to(orig_dtype)
167 return state_dict
169 def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
170 """Match Cohere's ``lm_head -> logit_scale -> optional softcap`` path."""
171 scale: float = getattr(self.cfg, "logit_scale")
172 return super().apply_output_logits_transform(logits * scale)
175class _Cohere2AttentionBridge(PositionEmbeddingsAttentionBridge):
176 """Attention bridge that honours Cohere2's RoPE/NoPE interleaving.
178 Cohere2 applies RoPE only on sliding-window layers. Full-attention global
179 layers receive the same position_embeddings tuple from the model loop but
180 intentionally skip apply_rotary_pos_emb inside HF's Cohere2Attention by
181 checking whether ``self.sliding_window`` is set.
183 The base bridge rotates whenever position_embeddings is present, so full
184 layers must suppress that argument before delegating to the shared attention
185 reconstruction path.
186 """
188 # Nulls position_embeddings on NoPE layers by design.
189 rope_optional = True
191 def forward(self, *args: Any, **kwargs: Any) -> Any:
192 """Drop position_embeddings on Cohere2 full-attention NoPE layers."""
193 if self._is_nope_layer():
194 kwargs["position_embeddings"] = None
195 if len(args) >= 2 and not isinstance(args[1], torch.Tensor):
196 args = (args[0], None) + args[2:]
197 return super().forward(*args, **kwargs)
199 def _is_nope_layer(self) -> bool:
200 """Return True when the wrapped Cohere2 attention is a full-attention layer."""
201 hf_attn = self.original_component
202 if hf_attn is None: 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 return False
205 if hasattr(hf_attn, "sliding_window"):
206 return getattr(hf_attn, "sliding_window") is None
208 layer_idx = getattr(hf_attn, "layer_idx", None)
209 layer_types = getattr(self.config, "layer_types", None)
210 if layer_idx is None or layer_types is None: 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true
211 return False
212 return layer_types[layer_idx] == "full_attention"
215def _cohere2_layer_types(cfg: Any) -> list[str]:
216 """Resolve Cohere2 layer types from explicit config or sliding-window pattern."""
217 n_layers = getattr(cfg, "n_layers", None)
218 if n_layers is None: 218 ↛ 219line 218 didn't jump to line 219 because the condition on line 218 was never true
219 n_layers = getattr(cfg, "num_hidden_layers")
220 n_layers = int(n_layers)
221 layer_types = getattr(cfg, "layer_types", None)
222 if layer_types is not None:
223 resolved = list(layer_types)
224 if len(resolved) != n_layers:
225 raise ValueError(
226 f"Cohere2 layer_types length ({len(resolved)}) must match n_layers ({n_layers})."
227 )
228 return resolved
230 pattern = getattr(cfg, "sliding_window_pattern", None)
231 if pattern is None:
232 pattern = getattr(cfg, "_sliding_window_pattern", None)
233 if pattern is None:
234 pattern = 4
235 pattern = int(pattern)
236 if pattern <= 0:
237 raise ValueError(f"Cohere2 sliding_window_pattern must be positive, got {pattern}.")
239 return [
240 "sliding_attention" if (layer_idx + 1) % pattern else "full_attention"
241 for layer_idx in range(n_layers)
242 ]
245class Cohere2ArchitectureAdapter(CohereArchitectureAdapter):
246 """Architecture adapter for Cohere2 / Command-A models.
248 Cohere2 keeps Cohere v1's parallel block, LayerNorm, GQA, gated MLP and
249 logit_scale behaviour, but interleaves sliding-window RoPE layers with
250 full-attention NoPE layers. HF represents that either as an explicit
251 ``layer_types`` list or as a legacy ``sliding_window_pattern`` integer.
252 """
254 def __init__(self, cfg: Any) -> None:
255 """Initialize the Cohere2 architecture adapter."""
256 super().__init__(cfg)
258 setattr(self.cfg, "layer_types", _cohere2_layer_types(cfg))
259 blocks = self.components["blocks"]
260 assert blocks.submodules is not None
261 blocks.submodules["attn"] = _Cohere2AttentionBridge(
262 name="self_attn",
263 config=self.cfg,
264 submodules={
265 "q": LinearBridge(name="q_proj"),
266 "k": LinearBridge(name="k_proj"),
267 "v": LinearBridge(name="v_proj"),
268 "o": LinearBridge(name="o_proj"),
269 },
270 requires_attention_mask=True,
271 requires_position_embeddings=True,
272 )