Coverage for transformer_lens/model_bridge/supported_architectures/cohere.py: 93%
100 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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: cos/sin come from HF's CohereRotaryEmbedding, but the
11 bridge rotates Q/K itself (cfg.rotary_interleaved_cos_sin=True)
12"""
14from typing import Any
16import torch
18from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
19from transformer_lens.model_bridge.generalized_components import (
20 EmbeddingBridge,
21 LinearBridge,
22 NormalizationBridge,
23 ParallelBlockBridge,
24 PositionEmbeddingsAttentionBridge,
25 RotaryEmbeddingBridge,
26 UnembeddingBridge,
27)
30class CohereArchitectureAdapter(ArchitectureAdapter):
31 """Architecture adapter for Cohere models (CohereForCausalLM).
33 Architectural quirks vs. standard decoder-only models:
34 - Single input_layernorm per block; NO post_attention_layernorm.
35 Attention and MLP both read the SAME normed hidden states (parallel).
36 - CohereLayerNorm is true LayerNorm (mean-subtracting), NOT RMSNorm.
37 It has a weight parameter but NO bias parameter.
38 - Logit scale: CohereForCausalLM.forward multiplies logits by logit_scale
39 (default 0.0625 = 1/16). Folded into unembed.weight via preprocess_weights.
40 - Rotary embeddings use repeat_interleave instead of cat-split; the bridge
41 applies that rotation itself in its reconstructed attention.
43 Optional parameters (absent from state_dict by default):
44 - blocks.{i}.attn.b_Q/b_K/b_V/b_O — no bias on projections (attention_bias=False)
45 - blocks.{i}.mlp.b_gate/b_in/b_out — no bias on MLP projections
46 - blocks.{i}.ln1.b — CohereLayerNorm has no bias
47 - ln_final.b — CohereLayerNorm has no bias
48 """
50 _testing_eager = None
52 def __init__(self, cfg: Any) -> None:
53 """Initialize the Cohere architecture adapter."""
54 super().__init__(cfg)
56 # --- Normalization ---
57 # CohereLayerNorm is true LayerNorm (subtracts mean), NOT RMSNorm.
58 # uses_rms_norm=False tells NormalizationBridge to subtract the mean.
59 self.cfg.normalization_type = "LN"
60 self.cfg.uses_rms_norm = False
61 self.cfg.final_rms = False
63 # --- Position embeddings and MLP ---
64 self.cfg.positional_embedding_type = "rotary"
65 self.cfg.gated_mlp = True
66 self.cfg.attn_only = False
68 # --- Parallel block: single norm, no post_attention_layernorm ---
69 self.cfg.parallel_attn_mlp = True
71 # CohereRotaryEmbedding repeat_interleaves cos/sin and rotate_half pairs adjacent elements.
72 self.cfg.rotary_interleaved_cos_sin = True
74 # --- Tokenizer: BOS is prepended by default ---
75 # CohereTokenizerFast has add_bos_token=False but HF's __call__ with
76 # add_special_tokens=True (the default) prepends BOS. Verified against
77 # trl-internal-testing/tiny-CohereForCausalLM.
78 self.cfg.default_prepend_bos = True
80 # --- GQA: n_key_value_heads ---
81 # sources/transformers.py copies num_key_value_heads generically.
82 # Re-read here to ensure it's set on cfg for _qkvo_weight_conversions.
83 n_kv = getattr(cfg, "n_key_value_heads", None)
84 if n_kv is not None: 84 ↛ 91line 84 didn't jump to line 91 because the condition on line 84 was always true
85 self.cfg.n_key_value_heads = n_kv
87 # --- Weight processing conversions ---
88 # Standard GQA-aware Q/K/V/O rearrangements (same as Llama/Qwen2).
89 # n_kv is already set on self.cfg; _qkvo_weight_conversions reads it via
90 # getattr(self.cfg, "n_key_value_heads", None) when called with no args.
91 self.weight_processing_conversions = {
92 **self._qkvo_weight_conversions(),
93 }
95 # --- Logit scale ---
96 # CohereConfig.logit_scale is typed float | None; apply explicit None-check
97 # so cfg.logit_scale is always a plain float (never None).
98 # logit_scale is not a declared field on TransformerBridgeConfig; it is a
99 # Cohere-specific dynamic attribute accessed later in preprocess_weights.
100 _ls = getattr(cfg, "logit_scale", None)
101 self.cfg.logit_scale = float(_ls) if _ls is not None else 0.0625 # type: ignore[attr-defined]
102 self._logit_scale_fold_pending = False
103 # Separate from cfg.logit_scale, which callers (apply_output_logits_transform,
104 # tests) read as Cohere's declared constant and must never see mutated.
105 self._logit_scale_already_folded = False
107 # --- RoPE theta (informational metadata) ---
108 # CohereRotaryEmbedding reads config.rope_parameters["rope_theta"] directly;
109 # store it in cfg.rotary_base so TL config accurately reflects the model.
110 # TransformerBridgeConfig stores rotary_base as int, matching its declared type.
111 _rope_params = getattr(cfg, "rope_parameters", None) or {}
112 if isinstance(_rope_params, dict): 112 ↛ 115line 112 didn't jump to line 115 because the condition on line 112 was always true
113 _theta = _rope_params.get("rope_theta", getattr(cfg, "default_theta", 10000.0))
114 else:
115 _theta = getattr(cfg, "default_theta", 10000.0)
116 self.cfg.rotary_base = int(_theta)
118 # --- Component mapping ---
119 # Block structure follows Falcon's parallel_attn=True, num_ln_in_parallel_attn=1
120 # mode: single ln1 feeds both attn and MLP; NO ln2.
121 # Submodule shapes follow Llama: separate q/k/v/o projections and SwiGLU MLP.
122 # Attention is reconstructed in the bridge (so hooks fire mid-computation);
123 # cfg.rotary_interleaved_cos_sin selects Cohere's repeat_interleave RoPE
124 # convention there, while cos/sin generation stays with the HF rotary module.
125 self.component_mapping = {
126 # Embedding: model.embed_tokens (same root as Llama, not transformer.* like Falcon)
127 "embed": EmbeddingBridge(name="model.embed_tokens"),
128 # Rotary embedding: top-level, delegates to CohereRotaryEmbedding.
129 # Pattern matches llama.py:75 and falcon.py:154 — NOT inside blocks.
130 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
131 "blocks": ParallelBlockBridge(
132 name="model.layers",
133 submodules={
134 # Single pre-norm only — Cohere has no post_attention_layernorm.
135 # NormalizationBridge handles weight-only CohereLayerNorm correctly:
136 # it checks `hasattr(original_component, "bias") and bias is not None`
137 # before adding bias, so the missing bias attribute is silently skipped.
138 "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg),
139 # No "ln2" — parallel block, same normed input goes to attn AND mlp.
140 "attn": PositionEmbeddingsAttentionBridge(
141 name="self_attn",
142 config=self.cfg,
143 submodules={
144 "q": LinearBridge(name="q_proj"),
145 "k": LinearBridge(name="k_proj"),
146 "v": LinearBridge(name="v_proj"),
147 "o": LinearBridge(name="o_proj"),
148 # Present only when use_qk_norm=True (config default False).
149 # Per-head CohereLayerNorm with weight (n_heads, d_head);
150 # the reconstruction applies it on the post-reshape path.
151 "q_norm": NormalizationBridge(
152 name="q_norm", config=self.cfg, optional=True
153 ),
154 "k_norm": NormalizationBridge(
155 name="k_norm", config=self.cfg, optional=True
156 ),
157 },
158 requires_attention_mask=True,
159 requires_position_embeddings=True,
160 ),
161 # GatedMLPBridge: gate/in/out matches Llama's gate_proj/up_proj/down_proj.
162 "mlp": self._gated_mlp(),
163 },
164 ),
165 # Final LayerNorm (CohereLayerNorm, weight-only) at model.norm
166 "ln_final": NormalizationBridge(name="model.norm", config=self.cfg),
167 # Unembed: lm_head. logit_scale is folded into weight in preprocess_weights.
168 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
169 }
171 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
172 """Fold logit_scale into unembed weights before ProcessWeights runs.
174 bridge.py clones unembed.weight before calling this, so
175 scaling does not affect the tied embed.weight.
176 logit_scale=1.0 is a no-op (skipped for efficiency).
177 """
178 self._logit_scale_fold_pending = False
179 if self._logit_scale_already_folded:
180 return state_dict
181 scale: float = getattr(self.cfg, "logit_scale") # always set by __init__
182 if scale != 1.0:
183 for key in ("unembed.weight", "unembed.bias"):
184 if key in state_dict:
185 orig_dtype = state_dict[key].dtype
186 state_dict[key] = (state_dict[key].float() * scale).to(orig_dtype)
187 self._logit_scale_fold_pending = True
188 return state_dict
190 def postprocess_weights(self, bridge: Any) -> None:
191 """Disable HF's outer scale after folding it into the live unembed weights."""
192 if not self._logit_scale_fold_pending:
193 return
195 model = getattr(bridge, "original_model", None)
196 if model is not None and hasattr(model, "logit_scale"):
197 model.logit_scale = 1.0
198 # Remember the fold happened without touching cfg.logit_scale itself: a repeat
199 # preprocess_weights() call (direct, or via a second process_weights()) must not
200 # re-fold the already-scaled unembed, but cfg.logit_scale is the model's declared
201 # constant and other readers (apply_output_logits_transform, tests) need it intact.
202 self._logit_scale_already_folded = True
203 self._logit_scale_fold_pending = False
205 def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
206 """Match Cohere's ``lm_head -> logit_scale -> optional softcap`` path.
208 Once preprocess_weights has folded the scale into the live unembed, the
209 head output already carries it, so applying cfg.logit_scale again would
210 double-scale every reader of this transform.
211 """
212 scale: float = 1.0 if self._logit_scale_already_folded else getattr(self.cfg, "logit_scale")
213 return super().apply_output_logits_transform(logits * scale)
216class _Cohere2AttentionBridge(PositionEmbeddingsAttentionBridge):
217 """Attention bridge that honours Cohere2's RoPE/NoPE interleaving.
219 Cohere2 applies RoPE only on sliding-window layers. Full-attention global
220 layers receive the same position_embeddings tuple from the model loop but
221 intentionally skip apply_rotary_pos_emb inside HF's Cohere2Attention by
222 checking whether ``self.sliding_window`` is set.
224 The base bridge rotates whenever position_embeddings is present, so full
225 layers must suppress that argument before delegating to the shared attention
226 reconstruction path.
227 """
229 # Nulls position_embeddings on NoPE layers by design.
230 rope_optional = True
232 def forward(self, *args: Any, **kwargs: Any) -> Any:
233 """Drop position_embeddings on Cohere2 full-attention NoPE layers."""
234 if self._is_nope_layer():
235 kwargs["position_embeddings"] = None
236 if len(args) >= 2 and not isinstance(args[1], torch.Tensor):
237 args = (args[0], None) + args[2:]
238 return super().forward(*args, **kwargs)
240 def _is_nope_layer(self) -> bool:
241 """Return True when the wrapped Cohere2 attention is a full-attention layer."""
242 hf_attn = self.original_component
243 if hf_attn is None: 243 ↛ 244line 243 didn't jump to line 244 because the condition on line 243 was never true
244 return False
246 if hasattr(hf_attn, "sliding_window"):
247 return getattr(hf_attn, "sliding_window") is None
249 layer_idx = getattr(hf_attn, "layer_idx", None)
250 layer_types = getattr(self.config, "layer_types", None)
251 if layer_idx is None or layer_types is None: 251 ↛ 252line 251 didn't jump to line 252 because the condition on line 251 was never true
252 return False
253 return layer_types[layer_idx] == "full_attention"
256def _cohere2_layer_types(cfg: Any) -> list[str]:
257 """Resolve Cohere2 layer types from explicit config or sliding-window pattern."""
258 n_layers = getattr(cfg, "n_layers", None)
259 if n_layers is None: 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true
260 n_layers = getattr(cfg, "num_hidden_layers")
261 n_layers = int(n_layers)
262 layer_types = getattr(cfg, "layer_types", None)
263 if layer_types is not None:
264 resolved = list(layer_types)
265 if len(resolved) != n_layers:
266 raise ValueError(
267 f"Cohere2 layer_types length ({len(resolved)}) must match n_layers ({n_layers})."
268 )
269 return resolved
271 pattern = getattr(cfg, "sliding_window_pattern", None)
272 if pattern is None:
273 pattern = getattr(cfg, "_sliding_window_pattern", None)
274 if pattern is None:
275 pattern = 4
276 pattern = int(pattern)
277 if pattern <= 0:
278 raise ValueError(f"Cohere2 sliding_window_pattern must be positive, got {pattern}.")
280 return [
281 "sliding_attention" if (layer_idx + 1) % pattern else "full_attention"
282 for layer_idx in range(n_layers)
283 ]
286class Cohere2ArchitectureAdapter(CohereArchitectureAdapter):
287 """Architecture adapter for Cohere2 / Command-A models.
289 Cohere2 keeps Cohere v1's parallel block, LayerNorm, GQA, gated MLP and
290 logit_scale behaviour, but interleaves sliding-window RoPE layers with
291 full-attention NoPE layers. HF represents that either as an explicit
292 ``layer_types`` list or as a legacy ``sliding_window_pattern`` integer.
293 """
295 def __init__(self, cfg: Any) -> None:
296 """Initialize the Cohere2 architecture adapter."""
297 super().__init__(cfg)
299 setattr(self.cfg, "layer_types", _cohere2_layer_types(cfg))
300 blocks = self.components["blocks"]
301 assert blocks.submodules is not None
302 blocks.submodules["attn"] = _Cohere2AttentionBridge(
303 name="self_attn",
304 config=self.cfg,
305 submodules={
306 "q": LinearBridge(name="q_proj"),
307 "k": LinearBridge(name="k_proj"),
308 "v": LinearBridge(name="v_proj"),
309 "o": LinearBridge(name="o_proj"),
310 },
311 requires_attention_mask=True,
312 requires_position_embeddings=True,
313 )