Coverage for transformer_lens/model_bridge/sources/native/model.py: 95%
292 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""TL-native transformer for TransformerBridge — minimal, no HF/HT dependency.
3Cfg-driven features: ``normalization_type`` (LN / RMS / LNPre / RMSPre —
4the ``Pre`` variants are param-free), ``final_rms``, ``gated_mlp``,
5``attn_only``, ``n_key_value_heads`` (GQA), ``attn_scores_soft_cap``,
6``output_logits_soft_cap``, ``positional_embedding_type`` (standard / rotary),
7``rotary_dim`` / ``rotary_base`` / ``rope_scaling`` (linear PI, dynamic/NTK,
8llama3 by-parts).
9"""
11from __future__ import annotations
13import math
14from typing import Callable, Optional, cast
16import torch
17import torch.nn as nn
18import torch.nn.functional as F
20from transformer_lens.config import TransformerBridgeConfig
21from transformer_lens.utilities import TypedModuleList
22from transformer_lens.utilities.activation_functions import apply_softcap
24# gelu_new = the tanh-approximation HF GPT-2 / HT use; F.gelu(approximate="tanh")
25# is the exact same formula.
26_Activation = Callable[[torch.Tensor], torch.Tensor]
27_ACTIVATIONS: dict[str, _Activation] = {
28 "gelu": F.gelu,
29 "gelu_new": lambda x: F.gelu(x, approximate="tanh"),
30 "relu": F.relu,
31 "silu": F.silu,
32 "swish": F.silu,
33 # SoLU (https://transformer-circuits.pub/2022/solu/index.html): x*softmax(x).
34 # "solu_ln" is the same activation; the mid-MLP LayerNorm that follows it is
35 # a NativeMLP submodule, not part of the pointwise function.
36 "solu": lambda x: x * F.softmax(x, dim=-1),
37 "solu_ln": lambda x: x * F.softmax(x, dim=-1),
38}
41def _normalization_type(cfg: TransformerBridgeConfig) -> str | None:
42 normalization_type = cfg.normalization_type
43 return None if normalization_type is None else normalization_type.upper()
46def _uses_rms_norm(cfg: TransformerBridgeConfig) -> bool:
47 return _normalization_type(cfg) in ("RMS", "RMSPRE")
50def _uses_no_norm(cfg: TransformerBridgeConfig) -> bool:
51 return _normalization_type(cfg) is None
54def _positional_kind(cfg: TransformerBridgeConfig) -> str:
55 return (getattr(cfg, "positional_embedding_type", None) or "standard").lower()
58class NativeRMSNorm(nn.Module):
59 """Llama-style RMSNorm. Variance in fp32 regardless of input dtype, then
60 cast back before the per-channel scale (matches HF LlamaRMSNorm)."""
62 def __init__(self, d_model: int, eps: float = 1e-5):
63 super().__init__()
64 self.weight = nn.Parameter(torch.ones(d_model))
65 self.eps = eps
67 def forward(self, x: torch.Tensor) -> torch.Tensor:
68 input_dtype = x.dtype
69 x_fp32 = x.to(torch.float32)
70 rms_inv = torch.rsqrt(x_fp32.pow(2).mean(dim=-1, keepdim=True) + self.eps)
71 normalized = (x_fp32 * rms_inv).to(input_dtype)
72 return self.weight * normalized
75class NativeRMSNormPre(nn.Module):
76 """Param-free RMSNorm — normalization only, no learnable scale."""
78 def __init__(self, eps: float = 1e-5):
79 super().__init__()
80 self.eps = eps
82 def forward(self, x: torch.Tensor) -> torch.Tensor:
83 input_dtype = x.dtype
84 x_fp32 = x.to(torch.float32)
85 rms_inv = torch.rsqrt(x_fp32.pow(2).mean(dim=-1, keepdim=True) + self.eps)
86 return (x_fp32 * rms_inv).to(input_dtype)
89class NativeLayerNormPre(nn.Module):
90 """Param-free LayerNorm — center + normalize only, no learnable scale/bias."""
92 def __init__(self, eps: float = 1e-5):
93 super().__init__()
94 self.eps = eps
96 def forward(self, x: torch.Tensor) -> torch.Tensor:
97 input_dtype = x.dtype
98 x_fp32 = x.to(torch.float32)
99 x_fp32 = x_fp32 - x_fp32.mean(dim=-1, keepdim=True)
100 scale = (x_fp32.pow(2).mean(dim=-1, keepdim=True) + self.eps).sqrt()
101 return (x_fp32 / scale).to(input_dtype)
104def _uses_param_free_norm(cfg: TransformerBridgeConfig) -> bool:
105 return _normalization_type(cfg) in ("RMSPRE", "LNPRE")
108def _make_norm(cfg: TransformerBridgeConfig, *, force_rms: bool = False) -> nn.Module:
109 param_free = _uses_param_free_norm(cfg)
110 if force_rms or _uses_rms_norm(cfg):
111 # final_rms swaps the norm family but must not reintroduce a scale the
112 # checkpoint doesn't carry.
113 if param_free:
114 return NativeRMSNormPre(eps=cfg.eps)
115 return NativeRMSNorm(cfg.d_model, eps=cfg.eps)
116 if _normalization_type(cfg) == "LNPRE":
117 return NativeLayerNormPre(eps=cfg.eps)
118 if _uses_no_norm(cfg):
119 return nn.Identity()
121 return nn.LayerNorm(cfg.d_model, eps=cfg.eps)
124def _uses_causal_attention(cfg: TransformerBridgeConfig) -> bool:
125 return cfg.attention_dir == "causal"
128def _resolve_rope_scaling(
129 cfg: TransformerBridgeConfig, rotary_dim: int
130) -> tuple[float, float, torch.Tensor]:
131 """Returns (effective_base, position_scale, inv_freq) per cfg.rope_scaling."""
132 base = float(cfg.rotary_base)
133 rope_scaling = getattr(cfg, "rope_scaling", None)
134 inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim))
136 if not isinstance(rope_scaling, dict):
137 return base, 1.0, inv_freq
139 # Newer HF configs key on "rope_type"; older ones on "type".
140 scale_type = str(rope_scaling.get("rope_type") or rope_scaling.get("type") or "").lower()
141 factor = float(rope_scaling.get("factor", 1.0))
143 if scale_type in ("", "default") or factor <= 1.0:
144 return base, 1.0, inv_freq
146 if scale_type == "linear":
147 return base, factor, inv_freq
149 if scale_type in ("dynamic", "ntk"):
150 scaled_base = base * (factor ** (rotary_dim / (rotary_dim - 2)))
151 new_inv_freq = 1.0 / (scaled_base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim))
152 return scaled_base, 1.0, new_inv_freq
154 if scale_type == "llama3":
155 low_freq_factor = float(rope_scaling.get("low_freq_factor", 1.0))
156 high_freq_factor = float(rope_scaling.get("high_freq_factor", 4.0))
157 original_ctx = float(
158 rope_scaling.get("original_max_position_embeddings")
159 or rope_scaling.get("original_context_length")
160 or 8192
161 )
162 low_wavelen = original_ctx / low_freq_factor
163 high_wavelen = original_ctx / high_freq_factor
164 wavelens = 2 * math.pi / inv_freq
165 # Three regimes: low-freq → divide by factor; high-freq → unchanged;
166 # in-between → smooth linear interpolation between the two.
167 smooth = (original_ctx / wavelens - low_freq_factor) / (high_freq_factor - low_freq_factor)
168 new_inv_freq = torch.where(
169 wavelens > low_wavelen,
170 inv_freq / factor,
171 torch.where(
172 wavelens < high_wavelen,
173 inv_freq,
174 (1 - smooth) * inv_freq / factor + smooth * inv_freq,
175 ),
176 )
177 return base, 1.0, new_inv_freq
179 raise NotImplementedError(
180 f"rope_scaling type {scale_type!r} is not supported. "
181 f"Supported: 'linear', 'dynamic'/'ntk', 'llama3'."
182 )
185class NativeRotary(nn.Module):
186 """Shared cos/sin tables for RoPE. Honors ``cfg.rope_scaling``."""
188 # Declared so mypy sees the buffer dtype; register_buffer alone reports Module|Tensor.
189 cos_cached: torch.Tensor
190 sin_cached: torch.Tensor
192 def __init__(self, cfg: TransformerBridgeConfig):
193 super().__init__()
194 rotary_dim = cfg.rotary_dim if cfg.rotary_dim is not None else cfg.d_head
195 if rotary_dim <= 0 or rotary_dim % 2 != 0: 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 raise ValueError(f"rotary_dim must be a positive even integer, got {rotary_dim!r}")
197 self.rotary_dim = rotary_dim
199 base, position_scale, inv_freq = _resolve_rope_scaling(cfg, rotary_dim)
201 positions = torch.arange(cfg.n_ctx).float() / position_scale
202 freqs = torch.outer(positions, inv_freq)
203 # Llama/HF adjacent-pair format: each (2i, 2i+1) pair rotates together.
204 cos = freqs.cos().repeat_interleave(2, dim=-1)
205 sin = freqs.sin().repeat_interleave(2, dim=-1)
206 self.register_buffer("cos_cached", cos, persistent=False)
207 self.register_buffer("sin_cached", sin, persistent=False)
208 self.effective_base = base
209 self.position_scale = position_scale
211 @staticmethod
212 def _rotate_half(x: torch.Tensor) -> torch.Tensor:
213 # Llama-style adjacent-pair rotation: (x0, x1) -> (-x1, x0).
214 x1 = x[..., 0::2]
215 x2 = x[..., 1::2]
216 rot = torch.stack((-x2, x1), dim=-1)
217 return rot.flatten(-2)
219 def apply_rope(
220 self,
221 q: torch.Tensor,
222 k: torch.Tensor,
223 *,
224 position_ids: Optional[torch.Tensor] = None,
225 ) -> tuple[torch.Tensor, torch.Tensor]:
226 """Apply RoPE to Q/K of shape [batch, heads, seq, d_head].
228 Named ``apply_rope`` rather than ``apply`` so ``nn.Module.apply(fn)``
229 — PyTorch's recursive function-application utility used by
230 ``bridge.apply(init_fn)`` — isn't shadowed.
231 """
232 seq = q.shape[-2]
233 rd = self.rotary_dim
234 if position_ids is None: 234 ↛ 235line 234 didn't jump to line 235 because the condition on line 234 was never true
235 cos = self.cos_cached[:seq].to(q.dtype)
236 sin = self.sin_cached[:seq].to(q.dtype)
237 else:
238 # [batch, seq] -> [batch, 1, seq, rd] (head dim for broadcast).
239 cos = self.cos_cached[position_ids].to(q.dtype).unsqueeze(1)
240 sin = self.sin_cached[position_ids].to(q.dtype).unsqueeze(1)
242 def _rope(x: torch.Tensor) -> torch.Tensor:
243 x_rot, x_pass = x[..., :rd], x[..., rd:]
244 x_rot = x_rot * cos + self._rotate_half(x_rot) * sin
245 return torch.cat([x_rot, x_pass], dim=-1) if x_pass.shape[-1] else x_rot
247 return _rope(q), _rope(k)
250class NativeAttention(nn.Module):
251 """Split-QKV causal self-attention. Returns (out, pattern).
253 ``accepts_pattern_fn``: AttentionBridge injects its ``hook_pattern`` as
254 ``pattern_fn``, applied BEFORE the value matmul — so hook edits genuinely
255 re-weight the attention output instead of only decorating the returned
256 tuple (which the wrapper cannot recompute from).
257 """
259 accepts_pattern_fn = True
261 causal_mask: torch.Tensor
263 def __init__(self, cfg: TransformerBridgeConfig, rotary: Optional[NativeRotary] = None):
264 super().__init__()
265 self.cfg = cfg
266 self.n_heads = cfg.n_heads
267 self.d_head = cfg.d_head
268 self.d_model = cfg.d_model
269 self.n_kv_heads = cfg.n_key_value_heads or cfg.n_heads
270 if self.n_heads % self.n_kv_heads != 0: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 raise ValueError(
272 f"n_heads ({self.n_heads}) must be divisible by n_key_value_heads "
273 f"({self.n_kv_heads}) for GQA."
274 )
275 self.kv_repeats = self.n_heads // self.n_kv_heads
277 q_dim = self.n_heads * self.d_head
278 kv_dim = self.n_kv_heads * self.d_head
279 self.q = nn.Linear(cfg.d_model, q_dim, bias=True)
280 self.k = nn.Linear(cfg.d_model, kv_dim, bias=True)
281 self.v = nn.Linear(cfg.d_model, kv_dim, bias=True)
282 self.o = nn.Linear(q_dim, cfg.d_model, bias=True)
284 mask = torch.triu(torch.ones(cfg.n_ctx, cfg.n_ctx, dtype=torch.bool), diagonal=1)
285 self.register_buffer("causal_mask", mask, persistent=False)
287 # attn_scale=1.0 reads like "standard scaling" but is "divide by 1" —
288 # i.e. unscaled scores, which saturate softmax for d_head>1.
289 if cfg.use_attn_scale and cfg.attn_scale > 0:
290 if self.d_head > 1 and math.isclose(cfg.attn_scale, 1.0, abs_tol=1e-9):
291 raise ValueError(
292 f"attn_scale=1.0 with d_head={self.d_head} (>1) is unscaled "
293 f"attention; softmax will saturate. For standard scaling "
294 f"leave attn_scale at -1 (sentinel for sqrt(d_head))."
295 )
296 scale = cfg.attn_scale
297 else:
298 scale = math.sqrt(cfg.d_head)
299 self.scale = scale
300 self.rotary = rotary
301 self.attn_scores_soft_cap = float(cfg.attn_scores_soft_cap)
302 self.causal = _uses_causal_attention(cfg)
304 def forward(
305 self,
306 hidden_states: torch.Tensor,
307 attention_mask: Optional[torch.Tensor] = None,
308 position_ids: Optional[torch.Tensor] = None,
309 pattern_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
310 **kwargs,
311 ) -> tuple[torch.Tensor, torch.Tensor]:
312 batch, seq, _ = hidden_states.shape
314 q = self.q(hidden_states).view(batch, seq, self.n_heads, self.d_head).transpose(1, 2)
315 k = self.k(hidden_states).view(batch, seq, self.n_kv_heads, self.d_head).transpose(1, 2)
316 v = self.v(hidden_states).view(batch, seq, self.n_kv_heads, self.d_head).transpose(1, 2)
318 if self.rotary is not None:
319 q, k = self.rotary.apply_rope(q, k, position_ids=position_ids)
321 # GQA: repeat_interleave matches HF Llama's repeat_kv group ordering.
322 if self.kv_repeats > 1:
323 k = k.repeat_interleave(self.kv_repeats, dim=1)
324 v = v.repeat_interleave(self.kv_repeats, dim=1)
326 scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale
327 # Gemma2 soft-cap before the causal mask so masked positions stay -inf.
328 scores = apply_softcap(scores, self.attn_scores_soft_cap)
330 if self.causal:
331 block_mask = self.causal_mask[:seq, :seq]
332 else:
333 block_mask = torch.zeros(seq, seq, dtype=torch.bool, device=scores.device)
334 if attention_mask is not None:
335 block_mask = self._combine_attention_mask(block_mask, attention_mask, batch=batch)
336 scores = scores.masked_fill(block_mask, float("-inf"))
338 pattern = F.softmax(scores, dim=-1)
339 # Fully masked padding queries softmax to NaN; overwrite masked entries
340 # so those rows contribute a zero attention update instead of poisoning later layers.
341 pattern = pattern.masked_fill(block_mask, 0.0)
342 if pattern_fn is not None:
343 pattern = pattern_fn(pattern)
345 attn = torch.matmul(pattern, v).transpose(1, 2).contiguous().view(batch, seq, -1)
346 out = self.o(attn)
347 return out, pattern
349 @staticmethod
350 def _combine_attention_mask(
351 block_mask: torch.Tensor, attention_mask: torch.Tensor, *, batch: int
352 ) -> torch.Tensor:
353 """Combine an external attention_mask with the causal mask.
355 Accepts 2D HF padding mask ``[batch, seq]`` (1=keep, 0=mask), 4D bool
356 mask (True=mask), or 4D additive float mask (HF generation style; values
357 below -1 treated as masked).
358 """
359 if attention_mask.dim() == 2:
360 pad_mask = ~attention_mask.bool()
361 return block_mask | pad_mask[:, None, None, :]
362 if attention_mask.dim() == 4:
363 if attention_mask.dtype is torch.bool:
364 return block_mask | attention_mask
365 # HF additive masks use -inf or large negatives; benign biases bounded.
366 extra = attention_mask < -1.0
367 return block_mask | extra
368 raise ValueError(
369 f"attention_mask must be 2D [batch, seq] or 4D [batch, *, seq, seq], "
370 f"got shape {tuple(attention_mask.shape)}."
371 )
374class NativeMLP(nn.Module):
375 """Two-layer MLP with configurable activation."""
377 act: Callable[[torch.Tensor], torch.Tensor]
379 def __init__(self, cfg: TransformerBridgeConfig):
380 super().__init__()
381 assert cfg.d_mlp is not None, "NativeModel resolves d_mlp before instantiating MLPs"
382 d_mlp: int = cfg.d_mlp
383 self.fc_in = nn.Linear(cfg.d_model, d_mlp, bias=True)
384 self.fc_out = nn.Linear(d_mlp, cfg.d_model, bias=True)
385 act_name = (cfg.act_fn or "gelu").lower()
386 if act_name not in _ACTIVATIONS: 386 ↛ 387line 386 didn't jump to line 387 because the condition on line 386 was never true
387 raise ValueError(f"Unsupported act_fn={act_name!r}. Supported: {sorted(_ACTIVATIONS)}")
388 self.act = _ACTIVATIONS[act_name]
389 # SoLU-LN models (NeelNanda's SoLU family) apply a LayerNorm between the
390 # activation and the out-projection; without it their checkpoints load
391 # but compute the wrong function. Named ``ln`` to match the legacy
392 # property-format key blocks.{i}.mlp.ln.{w,b}.
393 self.ln: Optional[nn.LayerNorm] = (
394 nn.LayerNorm(d_mlp, eps=cfg.eps) if act_name == "solu_ln" else None
395 )
397 def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
398 mid = self.act(self.fc_in(hidden_states))
399 if self.ln is not None:
400 mid = self.ln(mid)
401 return self.fc_out(mid)
404class NativeGatedMLP(nn.Module):
405 """SwiGLU / ReGLU / GeGLU gated MLP (variant picked by ``cfg.act_fn``).
407 Submodules ``gate`` / ``in`` / ``out`` match GatedMLPBridge's expected slots.
408 """
410 act: Callable[[torch.Tensor], torch.Tensor]
412 def __init__(self, cfg: TransformerBridgeConfig):
413 super().__init__()
414 assert cfg.d_mlp is not None, "NativeModel resolves d_mlp before instantiating MLPs"
415 d_mlp: int = cfg.d_mlp
416 # Llama convention: no biases on gated MLP projections.
417 self.gate = nn.Linear(cfg.d_model, d_mlp, bias=False)
418 # ``in`` is a Python keyword; add_module + getattr(self, "in") works
419 # because the bridge resolves LinearBridge(name="in") the same way.
420 self.add_module("in", nn.Linear(cfg.d_model, d_mlp, bias=False))
421 self.out = nn.Linear(d_mlp, cfg.d_model, bias=False)
422 # Default to SwiGLU; mirror NativeMLP's dispatch so a typo'd act_fn
423 # raises instead of silently changing the model.
424 act_name = (cfg.act_fn or "silu").lower()
425 if act_name not in _ACTIVATIONS:
426 raise ValueError(f"Unsupported act_fn={act_name!r}. Supported: {sorted(_ACTIVATIONS)}")
427 self.act = _ACTIVATIONS[act_name]
429 def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
430 gate_out = self.act(self.gate(hidden_states))
431 in_proj = cast(nn.Linear, getattr(self, "in"))
432 up_out = in_proj(hidden_states)
433 return self.out(gate_out * up_out)
436class NativeBlock(nn.Module):
437 """Pre-LN transformer block. Layout adapts to ``cfg.attn_only`` and
438 ``cfg.gated_mlp``."""
440 def __init__(self, cfg: TransformerBridgeConfig, rotary: Optional[NativeRotary] = None):
441 super().__init__()
442 self.cfg = cfg
443 self.ln1 = _make_norm(cfg)
444 self.attn = NativeAttention(cfg, rotary=rotary)
445 if not cfg.attn_only:
446 self.ln2 = _make_norm(cfg)
447 self.mlp = NativeGatedMLP(cfg) if cfg.gated_mlp else NativeMLP(cfg)
449 def forward(
450 self,
451 hidden_states: torch.Tensor,
452 attention_mask: Optional[torch.Tensor] = None,
453 position_ids: Optional[torch.Tensor] = None,
454 **kwargs,
455 ) -> tuple[torch.Tensor]:
456 attn_out, _pattern = self.attn(
457 self.ln1(hidden_states),
458 attention_mask=attention_mask,
459 position_ids=position_ids,
460 )
461 hidden_states = hidden_states + attn_out
462 if not self.cfg.attn_only:
463 hidden_states = hidden_states + self.mlp(self.ln2(hidden_states))
464 # Tuple return matches HF block convention; BlockBridge's parser expects it.
465 return (hidden_states,)
468class NativeModel(nn.Module):
469 """TL-native transformer. See module docstring for the supported feature set."""
471 pos: Optional[nn.Embedding]
472 rotary: Optional[NativeRotary]
474 def __init__(self, cfg: TransformerBridgeConfig):
475 super().__init__()
476 # Write the resolved d_mlp back so downstream consumers see the real
477 # value, not None. Mutates cfg; isolating callers should deep-copy first.
478 if not getattr(cfg, "d_mlp", None): 478 ↛ 479line 478 didn't jump to line 479 because the condition on line 478 was never true
479 cfg.d_mlp = 4 * cfg.d_model
480 self.cfg = cfg
482 self.tok_embed = nn.Embedding(cfg.d_vocab, cfg.d_model)
484 kind = _positional_kind(cfg)
485 if kind == "standard":
486 self.pos = nn.Embedding(cfg.n_ctx, cfg.d_model)
487 self.rotary = None
488 elif kind == "rotary": 488 ↛ 492line 488 didn't jump to line 492 because the condition on line 488 was always true
489 self.pos = None
490 self.rotary = NativeRotary(cfg)
491 else:
492 raise ValueError(
493 f"Unsupported positional_embedding_type={kind!r}. "
494 f"NativeModel supports 'standard' and 'rotary'."
495 )
497 self.layers = TypedModuleList(
498 [NativeBlock(cfg, rotary=self.rotary) for _ in range(cfg.n_layers)]
499 )
500 # final_rms forces RMS on the final norm regardless of block-norm choice
501 # — matches the TL config semantic Llama uses.
502 self.ln_out = _make_norm(cfg, force_rms=cfg.final_rms)
503 d_vocab_out = cfg.d_vocab_out if cfg.d_vocab_out > 0 else cfg.d_vocab
504 self.head = nn.Linear(cfg.d_model, d_vocab_out, bias=False)
505 self.output_logits_soft_cap = float(cfg.output_logits_soft_cap)
507 def forward(
508 self,
509 input_ids: Optional[torch.Tensor] = None,
510 attention_mask: Optional[torch.Tensor] = None,
511 position_ids: Optional[torch.Tensor] = None,
512 inputs_embeds: Optional[torch.Tensor] = None,
513 **kwargs,
514 ) -> torch.Tensor:
515 """Returns logits directly."""
516 if input_ids is not None and inputs_embeds is not None: 516 ↛ 517line 516 didn't jump to line 517 because the condition on line 516 was never true
517 raise ValueError("Exactly one of input_ids or inputs_embeds must be provided.")
518 if input_ids is not None:
519 model_input = input_ids
520 hidden_states = self.tok_embed(input_ids)
521 elif inputs_embeds is not None: 521 ↛ 525line 521 didn't jump to line 525 because the condition on line 521 was always true
522 model_input = inputs_embeds
523 hidden_states = inputs_embeds
524 else:
525 raise ValueError("Exactly one of input_ids or inputs_embeds must be provided.")
527 # Bounds check up front so both absolute and rotary paths produce a
528 # self-explanatory error rather than IndexError / shape mismatch.
529 seq_len = model_input.shape[1]
530 if seq_len > self.cfg.n_ctx:
531 raise ValueError(
532 f"input length {seq_len} exceeds n_ctx={self.cfg.n_ctx}; "
533 f"position embeddings and rotary tables are pre-baked at n_ctx."
534 )
536 # Resolve position_ids before the block loop so rotary sees the caller's
537 # positions, not the dense default.
538 batch, seq = model_input.shape[:2]
539 if position_ids is None:
540 position_ids = (
541 torch.arange(seq, device=model_input.device).unsqueeze(0).expand(batch, -1)
542 )
544 if self.pos is not None:
545 hidden_states = hidden_states + self.pos(position_ids)
547 for block in self.layers:
548 (hidden_states,) = block(
549 hidden_states, attention_mask=attention_mask, position_ids=position_ids
550 )
551 hidden_states = self.ln_out(hidden_states)
552 logits = self.head(hidden_states)
553 logits = apply_softcap(logits, self.output_logits_soft_cap)
554 return logits