Coverage for transformer_lens/model_bridge/sources/native/model.py: 93%
259 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""TL-native transformer for TransformerBridge — minimal, no HF/HT dependency.
3Cfg-driven features: ``normalization_type`` (LN / RMS / RMSPre), ``final_rms``,
4``gated_mlp``, ``attn_only``, ``n_key_value_heads`` (GQA), ``attn_scores_soft_cap``,
5``output_logits_soft_cap``, ``positional_embedding_type`` (standard / rotary),
6``rotary_dim`` / ``rotary_base`` / ``rope_scaling`` (linear PI, dynamic/NTK,
7llama3 by-parts).
8"""
10from __future__ import annotations
12import math
13from typing import Callable, Optional, cast
15import torch
16import torch.nn as nn
17import torch.nn.functional as F
19from transformer_lens.config import TransformerBridgeConfig
20from transformer_lens.utilities import TypedModuleList
21from transformer_lens.utilities.activation_functions import apply_softcap
23# gelu_new = the tanh-approximation HF GPT-2 / HT use; F.gelu(approximate="tanh")
24# is the exact same formula.
25_Activation = Callable[[torch.Tensor], torch.Tensor]
26_ACTIVATIONS: dict[str, _Activation] = {
27 "gelu": F.gelu,
28 "gelu_new": lambda x: F.gelu(x, approximate="tanh"),
29 "relu": F.relu,
30 "silu": F.silu,
31 "swish": F.silu,
32}
35def _normalization_type(cfg: TransformerBridgeConfig) -> str | None:
36 normalization_type = cfg.normalization_type
37 return None if normalization_type is None else normalization_type.upper()
40def _uses_rms_norm(cfg: TransformerBridgeConfig) -> bool:
41 return _normalization_type(cfg) in ("RMS", "RMSPRE")
44def _uses_no_norm(cfg: TransformerBridgeConfig) -> bool:
45 return _normalization_type(cfg) is None
48def _positional_kind(cfg: TransformerBridgeConfig) -> str:
49 return (getattr(cfg, "positional_embedding_type", None) or "standard").lower()
52class NativeRMSNorm(nn.Module):
53 """Llama-style RMSNorm. Variance in fp32 regardless of input dtype, then
54 cast back before the per-channel scale (matches HF LlamaRMSNorm)."""
56 def __init__(self, d_model: int, eps: float = 1e-5):
57 super().__init__()
58 self.weight = nn.Parameter(torch.ones(d_model))
59 self.eps = eps
61 def forward(self, x: torch.Tensor) -> torch.Tensor:
62 input_dtype = x.dtype
63 x_fp32 = x.to(torch.float32)
64 rms_inv = torch.rsqrt(x_fp32.pow(2).mean(dim=-1, keepdim=True) + self.eps)
65 normalized = (x_fp32 * rms_inv).to(input_dtype)
66 return self.weight * normalized
69def _make_norm(cfg: TransformerBridgeConfig, *, force_rms: bool = False) -> nn.Module:
70 if force_rms or _uses_rms_norm(cfg):
71 return NativeRMSNorm(cfg.d_model, eps=cfg.eps)
72 if _uses_no_norm(cfg):
73 return nn.Identity()
74 return nn.LayerNorm(cfg.d_model, eps=cfg.eps)
77def _uses_causal_attention(cfg: TransformerBridgeConfig) -> bool:
78 return cfg.attention_dir == "causal"
81def _resolve_rope_scaling(
82 cfg: TransformerBridgeConfig, rotary_dim: int
83) -> tuple[float, float, torch.Tensor]:
84 """Returns (effective_base, position_scale, inv_freq) per cfg.rope_scaling."""
85 base = float(cfg.rotary_base)
86 rope_scaling = getattr(cfg, "rope_scaling", None)
87 inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim))
89 if not isinstance(rope_scaling, dict):
90 return base, 1.0, inv_freq
92 # Newer HF configs key on "rope_type"; older ones on "type".
93 scale_type = str(rope_scaling.get("rope_type") or rope_scaling.get("type") or "").lower()
94 factor = float(rope_scaling.get("factor", 1.0))
96 if scale_type in ("", "default") or factor <= 1.0:
97 return base, 1.0, inv_freq
99 if scale_type == "linear":
100 return base, factor, inv_freq
102 if scale_type in ("dynamic", "ntk"):
103 scaled_base = base * (factor ** (rotary_dim / (rotary_dim - 2)))
104 new_inv_freq = 1.0 / (scaled_base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim))
105 return scaled_base, 1.0, new_inv_freq
107 if scale_type == "llama3":
108 low_freq_factor = float(rope_scaling.get("low_freq_factor", 1.0))
109 high_freq_factor = float(rope_scaling.get("high_freq_factor", 4.0))
110 original_ctx = float(
111 rope_scaling.get("original_max_position_embeddings")
112 or rope_scaling.get("original_context_length")
113 or 8192
114 )
115 low_wavelen = original_ctx / low_freq_factor
116 high_wavelen = original_ctx / high_freq_factor
117 wavelens = 2 * math.pi / inv_freq
118 # Three regimes: low-freq → divide by factor; high-freq → unchanged;
119 # in-between → smooth linear interpolation between the two.
120 smooth = (original_ctx / wavelens - low_freq_factor) / (high_freq_factor - low_freq_factor)
121 new_inv_freq = torch.where(
122 wavelens > low_wavelen,
123 inv_freq / factor,
124 torch.where(
125 wavelens < high_wavelen,
126 inv_freq,
127 (1 - smooth) * inv_freq / factor + smooth * inv_freq,
128 ),
129 )
130 return base, 1.0, new_inv_freq
132 raise NotImplementedError(
133 f"rope_scaling type {scale_type!r} is not supported. "
134 f"Supported: 'linear', 'dynamic'/'ntk', 'llama3'."
135 )
138class NativeRotary(nn.Module):
139 """Shared cos/sin tables for RoPE. Honors ``cfg.rope_scaling``."""
141 # Declared so mypy sees the buffer dtype; register_buffer alone reports Module|Tensor.
142 cos_cached: torch.Tensor
143 sin_cached: torch.Tensor
145 def __init__(self, cfg: TransformerBridgeConfig):
146 super().__init__()
147 rotary_dim = cfg.rotary_dim if cfg.rotary_dim is not None else cfg.d_head
148 if rotary_dim <= 0 or rotary_dim % 2 != 0: 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true
149 raise ValueError(f"rotary_dim must be a positive even integer, got {rotary_dim!r}")
150 self.rotary_dim = rotary_dim
152 base, position_scale, inv_freq = _resolve_rope_scaling(cfg, rotary_dim)
154 positions = torch.arange(cfg.n_ctx).float() / position_scale
155 freqs = torch.outer(positions, inv_freq)
156 # Llama/HF adjacent-pair format: each (2i, 2i+1) pair rotates together.
157 cos = freqs.cos().repeat_interleave(2, dim=-1)
158 sin = freqs.sin().repeat_interleave(2, dim=-1)
159 self.register_buffer("cos_cached", cos, persistent=False)
160 self.register_buffer("sin_cached", sin, persistent=False)
161 self.effective_base = base
162 self.position_scale = position_scale
164 @staticmethod
165 def _rotate_half(x: torch.Tensor) -> torch.Tensor:
166 # Llama-style adjacent-pair rotation: (x0, x1) -> (-x1, x0).
167 x1 = x[..., 0::2]
168 x2 = x[..., 1::2]
169 rot = torch.stack((-x2, x1), dim=-1)
170 return rot.flatten(-2)
172 def apply_rope(
173 self,
174 q: torch.Tensor,
175 k: torch.Tensor,
176 *,
177 position_ids: Optional[torch.Tensor] = None,
178 ) -> tuple[torch.Tensor, torch.Tensor]:
179 """Apply RoPE to Q/K of shape [batch, heads, seq, d_head].
181 Named ``apply_rope`` rather than ``apply`` so ``nn.Module.apply(fn)``
182 — PyTorch's recursive function-application utility used by
183 ``bridge.apply(init_fn)`` — isn't shadowed.
184 """
185 seq = q.shape[-2]
186 rd = self.rotary_dim
187 if position_ids is None: 187 ↛ 188line 187 didn't jump to line 188 because the condition on line 187 was never true
188 cos = self.cos_cached[:seq].to(q.dtype)
189 sin = self.sin_cached[:seq].to(q.dtype)
190 else:
191 # [batch, seq] -> [batch, 1, seq, rd] (head dim for broadcast).
192 cos = self.cos_cached[position_ids].to(q.dtype).unsqueeze(1)
193 sin = self.sin_cached[position_ids].to(q.dtype).unsqueeze(1)
195 def _rope(x: torch.Tensor) -> torch.Tensor:
196 x_rot, x_pass = x[..., :rd], x[..., rd:]
197 x_rot = x_rot * cos + self._rotate_half(x_rot) * sin
198 return torch.cat([x_rot, x_pass], dim=-1) if x_pass.shape[-1] else x_rot
200 return _rope(q), _rope(k)
203class NativeAttention(nn.Module):
204 """Split-QKV causal self-attention. Returns (out, pattern); AttentionBridge
205 fires ``hook_pattern`` off the second element."""
207 causal_mask: torch.Tensor
209 def __init__(self, cfg: TransformerBridgeConfig, rotary: Optional[NativeRotary] = None):
210 super().__init__()
211 self.cfg = cfg
212 self.n_heads = cfg.n_heads
213 self.d_head = cfg.d_head
214 self.d_model = cfg.d_model
215 self.n_kv_heads = cfg.n_key_value_heads or cfg.n_heads
216 if self.n_heads % self.n_kv_heads != 0: 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 raise ValueError(
218 f"n_heads ({self.n_heads}) must be divisible by n_key_value_heads "
219 f"({self.n_kv_heads}) for GQA."
220 )
221 self.kv_repeats = self.n_heads // self.n_kv_heads
223 q_dim = self.n_heads * self.d_head
224 kv_dim = self.n_kv_heads * self.d_head
225 self.q = nn.Linear(cfg.d_model, q_dim, bias=True)
226 self.k = nn.Linear(cfg.d_model, kv_dim, bias=True)
227 self.v = nn.Linear(cfg.d_model, kv_dim, bias=True)
228 self.o = nn.Linear(q_dim, cfg.d_model, bias=True)
230 mask = torch.triu(torch.ones(cfg.n_ctx, cfg.n_ctx, dtype=torch.bool), diagonal=1)
231 self.register_buffer("causal_mask", mask, persistent=False)
233 # attn_scale=1.0 reads like "standard scaling" but is "divide by 1" —
234 # i.e. unscaled scores, which saturate softmax for d_head>1.
235 if cfg.use_attn_scale and cfg.attn_scale > 0:
236 if self.d_head > 1 and math.isclose(cfg.attn_scale, 1.0, abs_tol=1e-9):
237 raise ValueError(
238 f"attn_scale=1.0 with d_head={self.d_head} (>1) is unscaled "
239 f"attention; softmax will saturate. For standard scaling "
240 f"leave attn_scale at -1 (sentinel for sqrt(d_head))."
241 )
242 scale = cfg.attn_scale
243 else:
244 scale = math.sqrt(cfg.d_head)
245 self.scale = scale
246 self.rotary = rotary
247 self.attn_scores_soft_cap = float(cfg.attn_scores_soft_cap)
248 self.causal = _uses_causal_attention(cfg)
250 def forward(
251 self,
252 hidden_states: torch.Tensor,
253 attention_mask: Optional[torch.Tensor] = None,
254 position_ids: Optional[torch.Tensor] = None,
255 **kwargs,
256 ) -> tuple[torch.Tensor, torch.Tensor]:
257 batch, seq, _ = hidden_states.shape
259 q = self.q(hidden_states).view(batch, seq, self.n_heads, self.d_head).transpose(1, 2)
260 k = self.k(hidden_states).view(batch, seq, self.n_kv_heads, self.d_head).transpose(1, 2)
261 v = self.v(hidden_states).view(batch, seq, self.n_kv_heads, self.d_head).transpose(1, 2)
263 if self.rotary is not None:
264 q, k = self.rotary.apply_rope(q, k, position_ids=position_ids)
266 # GQA: repeat_interleave matches HF Llama's repeat_kv group ordering.
267 if self.kv_repeats > 1:
268 k = k.repeat_interleave(self.kv_repeats, dim=1)
269 v = v.repeat_interleave(self.kv_repeats, dim=1)
271 scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale
272 # Gemma2 soft-cap before the causal mask so masked positions stay -inf.
273 scores = apply_softcap(scores, self.attn_scores_soft_cap)
275 if self.causal:
276 block_mask = self.causal_mask[:seq, :seq]
277 else:
278 block_mask = torch.zeros(seq, seq, dtype=torch.bool, device=scores.device)
279 if attention_mask is not None:
280 block_mask = self._combine_attention_mask(block_mask, attention_mask, batch=batch)
281 scores = scores.masked_fill(block_mask, float("-inf"))
283 pattern = F.softmax(scores, dim=-1)
284 # Fully masked padding queries softmax to NaN; overwrite masked entries
285 # so those rows contribute a zero attention update instead of poisoning later layers.
286 pattern = pattern.masked_fill(block_mask, 0.0)
288 attn = torch.matmul(pattern, v).transpose(1, 2).contiguous().view(batch, seq, -1)
289 out = self.o(attn)
290 return out, pattern
292 @staticmethod
293 def _combine_attention_mask(
294 block_mask: torch.Tensor, attention_mask: torch.Tensor, *, batch: int
295 ) -> torch.Tensor:
296 """Combine an external attention_mask with the causal mask.
298 Accepts 2D HF padding mask ``[batch, seq]`` (1=keep, 0=mask), 4D bool
299 mask (True=mask), or 4D additive float mask (HF generation style; values
300 below -1 treated as masked).
301 """
302 if attention_mask.dim() == 2:
303 pad_mask = ~attention_mask.bool()
304 return block_mask | pad_mask[:, None, None, :]
305 if attention_mask.dim() == 4:
306 if attention_mask.dtype is torch.bool:
307 return block_mask | attention_mask
308 # HF additive masks use -inf or large negatives; benign biases bounded.
309 extra = attention_mask < -1.0
310 return block_mask | extra
311 raise ValueError(
312 f"attention_mask must be 2D [batch, seq] or 4D [batch, *, seq, seq], "
313 f"got shape {tuple(attention_mask.shape)}."
314 )
317class NativeMLP(nn.Module):
318 """Two-layer MLP with configurable activation."""
320 act: Callable[[torch.Tensor], torch.Tensor]
322 def __init__(self, cfg: TransformerBridgeConfig):
323 super().__init__()
324 assert cfg.d_mlp is not None, "NativeModel resolves d_mlp before instantiating MLPs"
325 d_mlp: int = cfg.d_mlp
326 self.fc_in = nn.Linear(cfg.d_model, d_mlp, bias=True)
327 self.fc_out = nn.Linear(d_mlp, cfg.d_model, bias=True)
328 act_name = (cfg.act_fn or "gelu").lower()
329 if act_name not in _ACTIVATIONS: 329 ↛ 330line 329 didn't jump to line 330 because the condition on line 329 was never true
330 raise ValueError(f"Unsupported act_fn={act_name!r}. Supported: {sorted(_ACTIVATIONS)}")
331 self.act = _ACTIVATIONS[act_name]
333 def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
334 return self.fc_out(self.act(self.fc_in(hidden_states)))
337class NativeGatedMLP(nn.Module):
338 """SwiGLU / ReGLU / GeGLU gated MLP (variant picked by ``cfg.act_fn``).
340 Submodules ``gate`` / ``in`` / ``out`` match GatedMLPBridge's expected slots.
341 """
343 act: Callable[[torch.Tensor], torch.Tensor]
345 def __init__(self, cfg: TransformerBridgeConfig):
346 super().__init__()
347 assert cfg.d_mlp is not None, "NativeModel resolves d_mlp before instantiating MLPs"
348 d_mlp: int = cfg.d_mlp
349 # Llama convention: no biases on gated MLP projections.
350 self.gate = nn.Linear(cfg.d_model, d_mlp, bias=False)
351 # ``in`` is a Python keyword; add_module + getattr(self, "in") works
352 # because the bridge resolves LinearBridge(name="in") the same way.
353 self.add_module("in", nn.Linear(cfg.d_model, d_mlp, bias=False))
354 self.out = nn.Linear(d_mlp, cfg.d_model, bias=False)
355 # Default to SwiGLU; mirror NativeMLP's dispatch so a typo'd act_fn
356 # raises instead of silently changing the model.
357 act_name = (cfg.act_fn or "silu").lower()
358 if act_name not in _ACTIVATIONS:
359 raise ValueError(f"Unsupported act_fn={act_name!r}. Supported: {sorted(_ACTIVATIONS)}")
360 self.act = _ACTIVATIONS[act_name]
362 def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
363 gate_out = self.act(self.gate(hidden_states))
364 in_proj = cast(nn.Linear, getattr(self, "in"))
365 up_out = in_proj(hidden_states)
366 return self.out(gate_out * up_out)
369class NativeBlock(nn.Module):
370 """Pre-LN transformer block. Layout adapts to ``cfg.attn_only`` and
371 ``cfg.gated_mlp``."""
373 def __init__(self, cfg: TransformerBridgeConfig, rotary: Optional[NativeRotary] = None):
374 super().__init__()
375 self.cfg = cfg
376 self.ln1 = _make_norm(cfg)
377 self.attn = NativeAttention(cfg, rotary=rotary)
378 if not cfg.attn_only:
379 self.ln2 = _make_norm(cfg)
380 self.mlp = NativeGatedMLP(cfg) if cfg.gated_mlp else NativeMLP(cfg)
382 def forward(
383 self,
384 hidden_states: torch.Tensor,
385 attention_mask: Optional[torch.Tensor] = None,
386 position_ids: Optional[torch.Tensor] = None,
387 **kwargs,
388 ) -> tuple[torch.Tensor]:
389 attn_out, _pattern = self.attn(
390 self.ln1(hidden_states),
391 attention_mask=attention_mask,
392 position_ids=position_ids,
393 )
394 hidden_states = hidden_states + attn_out
395 if not self.cfg.attn_only:
396 hidden_states = hidden_states + self.mlp(self.ln2(hidden_states))
397 # Tuple return matches HF block convention; BlockBridge's parser expects it.
398 return (hidden_states,)
401class NativeModel(nn.Module):
402 """TL-native transformer. See module docstring for the supported feature set."""
404 pos: Optional[nn.Embedding]
405 rotary: Optional[NativeRotary]
407 def __init__(self, cfg: TransformerBridgeConfig):
408 super().__init__()
409 # Write the resolved d_mlp back so downstream consumers see the real
410 # value, not None. Mutates cfg; isolating callers should deep-copy first.
411 if not getattr(cfg, "d_mlp", None): 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true
412 cfg.d_mlp = 4 * cfg.d_model
413 self.cfg = cfg
415 self.tok_embed = nn.Embedding(cfg.d_vocab, cfg.d_model)
417 kind = _positional_kind(cfg)
418 if kind == "standard":
419 self.pos = nn.Embedding(cfg.n_ctx, cfg.d_model)
420 self.rotary = None
421 elif kind == "rotary": 421 ↛ 425line 421 didn't jump to line 425 because the condition on line 421 was always true
422 self.pos = None
423 self.rotary = NativeRotary(cfg)
424 else:
425 raise ValueError(
426 f"Unsupported positional_embedding_type={kind!r}. "
427 f"NativeModel supports 'standard' and 'rotary'."
428 )
430 self.layers = TypedModuleList(
431 [NativeBlock(cfg, rotary=self.rotary) for _ in range(cfg.n_layers)]
432 )
433 # final_rms forces RMS on the final norm regardless of block-norm choice
434 # — matches the TL config semantic Llama uses.
435 self.ln_out = _make_norm(cfg, force_rms=cfg.final_rms)
436 d_vocab_out = cfg.d_vocab_out if cfg.d_vocab_out > 0 else cfg.d_vocab
437 self.head = nn.Linear(cfg.d_model, d_vocab_out, bias=False)
438 self.output_logits_soft_cap = float(cfg.output_logits_soft_cap)
440 def forward(
441 self,
442 input_ids: Optional[torch.Tensor] = None,
443 attention_mask: Optional[torch.Tensor] = None,
444 position_ids: Optional[torch.Tensor] = None,
445 inputs_embeds: Optional[torch.Tensor] = None,
446 **kwargs,
447 ) -> torch.Tensor:
448 """Returns logits directly."""
449 if input_ids is not None and inputs_embeds is not None: 449 ↛ 450line 449 didn't jump to line 450 because the condition on line 449 was never true
450 raise ValueError("Exactly one of input_ids or inputs_embeds must be provided.")
451 if input_ids is not None: 451 ↛ 454line 451 didn't jump to line 454 because the condition on line 451 was always true
452 model_input = input_ids
453 hidden_states = self.tok_embed(input_ids)
454 elif inputs_embeds is not None:
455 model_input = inputs_embeds
456 hidden_states = inputs_embeds
457 else:
458 raise ValueError("Exactly one of input_ids or inputs_embeds must be provided.")
460 # Bounds check up front so both absolute and rotary paths produce a
461 # self-explanatory error rather than IndexError / shape mismatch.
462 seq_len = model_input.shape[1]
463 if seq_len > self.cfg.n_ctx:
464 raise ValueError(
465 f"input length {seq_len} exceeds n_ctx={self.cfg.n_ctx}; "
466 f"position embeddings and rotary tables are pre-baked at n_ctx."
467 )
469 # Resolve position_ids before the block loop so rotary sees the caller's
470 # positions, not the dense default.
471 batch, seq = model_input.shape[:2]
472 if position_ids is None:
473 position_ids = (
474 torch.arange(seq, device=model_input.device).unsqueeze(0).expand(batch, -1)
475 )
477 if self.pos is not None:
478 hidden_states = hidden_states + self.pos(position_ids)
480 for block in self.layers:
481 (hidden_states,) = block(
482 hidden_states, attention_mask=attention_mask, position_ids=position_ids
483 )
484 hidden_states = self.ln_out(hidden_states)
485 logits = self.head(hidden_states)
486 logits = apply_softcap(logits, self.output_logits_soft_cap)
487 return logits