Coverage for transformer_lens/model_bridge/supported_architectures/cohere.py: 93%
99 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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]
97 self._logit_scale_fold_pending = False
98 # Separate from cfg.logit_scale, which callers (apply_output_logits_transform,
99 # tests) read as Cohere's declared constant and must never see mutated.
100 self._logit_scale_already_folded = False
102 # --- RoPE theta (informational metadata) ---
103 # CohereRotaryEmbedding reads config.rope_parameters["rope_theta"] directly;
104 # store it in cfg.rotary_base so TL config accurately reflects the model.
105 # TransformerBridgeConfig stores rotary_base as int, matching its declared type.
106 _rope_params = getattr(cfg, "rope_parameters", None) or {}
107 if isinstance(_rope_params, dict): 107 ↛ 110line 107 didn't jump to line 110 because the condition on line 107 was always true
108 _theta = _rope_params.get("rope_theta", getattr(cfg, "default_theta", 10000.0))
109 else:
110 _theta = getattr(cfg, "default_theta", 10000.0)
111 self.cfg.rotary_base = int(_theta)
113 # --- Component mapping ---
114 # Block structure follows Falcon's parallel_attn=True, num_ln_in_parallel_attn=1
115 # mode: single ln1 feeds both attn and MLP; NO ln2.
116 # Submodule shapes follow Llama: separate q/k/v/o projections and SwiGLU MLP.
117 # Rotary and attention both delegate to HF modules, preserving Cohere's
118 # repeat_interleave RoPE convention without re-implementing it in TL.
119 self.component_mapping = {
120 # Embedding: model.embed_tokens (same root as Llama, not transformer.* like Falcon)
121 "embed": EmbeddingBridge(name="model.embed_tokens"),
122 # Rotary embedding: top-level, delegates to CohereRotaryEmbedding.
123 # Pattern matches llama.py:75 and falcon.py:154 — NOT inside blocks.
124 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
125 "blocks": ParallelBlockBridge(
126 name="model.layers",
127 submodules={
128 # Single pre-norm only — Cohere has no post_attention_layernorm.
129 # NormalizationBridge handles weight-only CohereLayerNorm correctly:
130 # it checks `hasattr(original_component, "bias") and bias is not None`
131 # before adding bias, so the missing bias attribute is silently skipped.
132 "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg),
133 # No "ln2" — parallel block, same normed input goes to attn AND mlp.
134 "attn": PositionEmbeddingsAttentionBridge(
135 name="self_attn",
136 config=self.cfg,
137 submodules={
138 "q": LinearBridge(name="q_proj"),
139 "k": LinearBridge(name="k_proj"),
140 "v": LinearBridge(name="v_proj"),
141 "o": LinearBridge(name="o_proj"),
142 },
143 requires_attention_mask=True,
144 requires_position_embeddings=True,
145 ),
146 # GatedMLPBridge: gate/in/out matches Llama's gate_proj/up_proj/down_proj.
147 # Optional use_qk_norm is handled transparently by HF's
148 # CohereAttention.forward delegation (no extra submodules needed).
149 "mlp": self._gated_mlp(),
150 },
151 ),
152 # Final LayerNorm (CohereLayerNorm, weight-only) at model.norm
153 "ln_final": NormalizationBridge(name="model.norm", config=self.cfg),
154 # Unembed: lm_head. logit_scale is folded into weight in preprocess_weights.
155 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
156 }
158 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
159 """Fold logit_scale into unembed weights before ProcessWeights runs.
161 bridge.py clones unembed.weight before calling this, so
162 scaling does not affect the tied embed.weight.
163 logit_scale=1.0 is a no-op (skipped for efficiency).
164 """
165 self._logit_scale_fold_pending = False
166 if self._logit_scale_already_folded:
167 return state_dict
168 scale: float = getattr(self.cfg, "logit_scale") # always set by __init__
169 if scale != 1.0:
170 for key in ("unembed.weight", "unembed.bias"):
171 if key in state_dict:
172 orig_dtype = state_dict[key].dtype
173 state_dict[key] = (state_dict[key].float() * scale).to(orig_dtype)
174 self._logit_scale_fold_pending = True
175 return state_dict
177 def postprocess_weights(self, bridge: Any) -> None:
178 """Disable HF's outer scale after folding it into the live unembed weights."""
179 if not self._logit_scale_fold_pending:
180 return
182 model = getattr(bridge, "original_model", None)
183 if model is not None and hasattr(model, "logit_scale"):
184 model.logit_scale = 1.0
185 # Remember the fold happened without touching cfg.logit_scale itself: a repeat
186 # preprocess_weights() call (direct, or via a second process_weights()) must not
187 # re-fold the already-scaled unembed, but cfg.logit_scale is the model's declared
188 # constant and other readers (apply_output_logits_transform, tests) need it intact.
189 self._logit_scale_already_folded = True
190 self._logit_scale_fold_pending = False
192 def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
193 """Match Cohere's ``lm_head -> logit_scale -> optional softcap`` path."""
194 scale: float = getattr(self.cfg, "logit_scale")
195 return super().apply_output_logits_transform(logits * scale)
198class _Cohere2AttentionBridge(PositionEmbeddingsAttentionBridge):
199 """Attention bridge that honours Cohere2's RoPE/NoPE interleaving.
201 Cohere2 applies RoPE only on sliding-window layers. Full-attention global
202 layers receive the same position_embeddings tuple from the model loop but
203 intentionally skip apply_rotary_pos_emb inside HF's Cohere2Attention by
204 checking whether ``self.sliding_window`` is set.
206 The base bridge rotates whenever position_embeddings is present, so full
207 layers must suppress that argument before delegating to the shared attention
208 reconstruction path.
209 """
211 # Nulls position_embeddings on NoPE layers by design.
212 rope_optional = True
214 def forward(self, *args: Any, **kwargs: Any) -> Any:
215 """Drop position_embeddings on Cohere2 full-attention NoPE layers."""
216 if self._is_nope_layer():
217 kwargs["position_embeddings"] = None
218 if len(args) >= 2 and not isinstance(args[1], torch.Tensor):
219 args = (args[0], None) + args[2:]
220 return super().forward(*args, **kwargs)
222 def _is_nope_layer(self) -> bool:
223 """Return True when the wrapped Cohere2 attention is a full-attention layer."""
224 hf_attn = self.original_component
225 if hf_attn is None: 225 ↛ 226line 225 didn't jump to line 226 because the condition on line 225 was never true
226 return False
228 if hasattr(hf_attn, "sliding_window"):
229 return getattr(hf_attn, "sliding_window") is None
231 layer_idx = getattr(hf_attn, "layer_idx", None)
232 layer_types = getattr(self.config, "layer_types", None)
233 if layer_idx is None or layer_types is None: 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true
234 return False
235 return layer_types[layer_idx] == "full_attention"
238def _cohere2_layer_types(cfg: Any) -> list[str]:
239 """Resolve Cohere2 layer types from explicit config or sliding-window pattern."""
240 n_layers = getattr(cfg, "n_layers", None)
241 if n_layers is None: 241 ↛ 242line 241 didn't jump to line 242 because the condition on line 241 was never true
242 n_layers = getattr(cfg, "num_hidden_layers")
243 n_layers = int(n_layers)
244 layer_types = getattr(cfg, "layer_types", None)
245 if layer_types is not None:
246 resolved = list(layer_types)
247 if len(resolved) != n_layers:
248 raise ValueError(
249 f"Cohere2 layer_types length ({len(resolved)}) must match n_layers ({n_layers})."
250 )
251 return resolved
253 pattern = getattr(cfg, "sliding_window_pattern", None)
254 if pattern is None:
255 pattern = getattr(cfg, "_sliding_window_pattern", None)
256 if pattern is None:
257 pattern = 4
258 pattern = int(pattern)
259 if pattern <= 0:
260 raise ValueError(f"Cohere2 sliding_window_pattern must be positive, got {pattern}.")
262 return [
263 "sliding_attention" if (layer_idx + 1) % pattern else "full_attention"
264 for layer_idx in range(n_layers)
265 ]
268class Cohere2ArchitectureAdapter(CohereArchitectureAdapter):
269 """Architecture adapter for Cohere2 / Command-A models.
271 Cohere2 keeps Cohere v1's parallel block, LayerNorm, GQA, gated MLP and
272 logit_scale behaviour, but interleaves sliding-window RoPE layers with
273 full-attention NoPE layers. HF represents that either as an explicit
274 ``layer_types`` list or as a legacy ``sliding_window_pattern`` integer.
275 """
277 def __init__(self, cfg: Any) -> None:
278 """Initialize the Cohere2 architecture adapter."""
279 super().__init__(cfg)
281 setattr(self.cfg, "layer_types", _cohere2_layer_types(cfg))
282 blocks = self.components["blocks"]
283 assert blocks.submodules is not None
284 blocks.submodules["attn"] = _Cohere2AttentionBridge(
285 name="self_attn",
286 config=self.cfg,
287 submodules={
288 "q": LinearBridge(name="q_proj"),
289 "k": LinearBridge(name="k_proj"),
290 "v": LinearBridge(name="v_proj"),
291 "o": LinearBridge(name="o_proj"),
292 },
293 requires_attention_mask=True,
294 requires_position_embeddings=True,
295 )