Coverage for transformer_lens/components/abstract_attention.py: 85%
407 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1import math
2from abc import ABC
3from typing import Dict, Optional, Tuple, Union, cast
5import einops
6import torch
7import torch.nn as nn
8import torch.nn.functional as F
9from better_abc import abstract_attribute
10from jaxtyping import Float, Int
11from torch import Tensor
12from transformers.utils.import_utils import is_bitsandbytes_available
14from transformer_lens.cache.key_value_cache_entry import (
15 TransformerLensKeyValueCacheEntry,
16)
17from transformer_lens.components.rms_norm import RMSNorm
18from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig
19from transformer_lens.FactoredMatrix import FactoredMatrix
20from transformer_lens.hook_points import HookPoint
21from transformer_lens.utilities import get_offset_position_ids
22from transformer_lens.utilities.activation_functions import apply_softcap
23from transformer_lens.utilities.attention import complex_attn_linear, simple_attn_linear
25if is_bitsandbytes_available(): 25 ↛ 26line 25 didn't jump to line 26 because the condition on line 25 was never true
26 import bitsandbytes as bnb
27 from bitsandbytes.nn.modules import Params4bit
30class AbstractAttention(ABC, nn.Module):
31 ROTARY_INITIAL_CACHE_SIZE = 2048
33 alibi: Union[torch.Tensor, None]
34 q_norm: Optional[RMSNorm]
35 k_norm: Optional[RMSNorm]
36 mask: torch.Tensor
37 IGNORE: torch.Tensor
38 rotary_sin: torch.Tensor
39 rotary_cos: torch.Tensor
41 def __init__(
42 self,
43 cfg: Union[Dict, HookedTransformerConfig],
44 attn_type: str = "global",
45 layer_id: Optional[int] = None,
46 ):
47 """Abstract Base Class of Attention Blocks, featuring common functionality of both Attention and GroupedQueryAttention blocks.
49 Query and Output projections are defined in this class as they are the same for regular and grouped query attention.
50 Attributes related to Key and Value projections are abstract as their implementations may differ. For example, in GroupedQueryAttention there are less query and key heads than value heads.
51 To enforce implementation of W_K, W_V, b_K, and b_V by child classes, the better_abc.abstract_attribute class is used. See here for details: https://stackoverflow.com/questions/23831510/abstract-attribute-not-property.
53 Args:
54 cfg (Union[Dict, HookedTransformerConfig]): Config
55 attn_type (str, optional): "global" or "local", used by GPT-Neo. Local attention means the model can only attend back cfg.window_size tokens (here, 256). Not used by any other model at the moment. Defaults to "global".
56 layer_id (int, optional): The index of the current layer. Used by the Mistral models (labelled here as stanford-gpt2) to scale down attention scores pre softmax for numerical stability reasons by 1/(layer_id+1). Defaults to None.
57 """
58 super().__init__()
59 self.cfg = HookedTransformerConfig.unwrap(cfg)
61 if self.cfg.load_in_4bit: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true
62 nq = int((self.cfg.d_model * self.cfg.d_head * self.cfg.n_heads) / 2)
63 self.W_Q: Union[nn.Parameter, "Params4bit"] = Params4bit(
64 torch.empty(nq, 1, dtype=torch.uint8), requires_grad=False
65 )
66 self.W_O: Union[nn.Parameter, "Params4bit"] = Params4bit(
67 torch.empty(nq, 1, dtype=torch.uint8), requires_grad=False
68 )
69 else:
70 self.W_Q = nn.Parameter(
71 torch.empty(
72 self.cfg.n_heads,
73 self.cfg.d_model,
74 self.cfg.d_head,
75 dtype=self.cfg.dtype,
76 )
77 )
78 self.W_O = nn.Parameter(
79 torch.empty(
80 self.cfg.n_heads,
81 self.cfg.d_head,
82 self.cfg.d_model,
83 dtype=self.cfg.dtype,
84 )
85 )
86 self.W_K = abstract_attribute()
87 self.W_V = abstract_attribute()
89 self.b_Q = nn.Parameter(
90 torch.zeros(self.cfg.n_heads, self.cfg.d_head, dtype=self.cfg.dtype)
91 )
92 self.b_K: nn.Parameter = abstract_attribute()
93 self.b_V: nn.Parameter = abstract_attribute()
94 self.b_O = nn.Parameter(torch.zeros(self.cfg.d_model, dtype=self.cfg.dtype))
96 if self.cfg.use_qk_norm:
97 self.q_norm = RMSNorm(self.cfg, length=self.cfg.d_head)
98 self.k_norm = RMSNorm(self.cfg, length=self.cfg.d_head)
100 elif self.cfg.original_architecture in (
101 "OlmoeForCausalLM",
102 "Olmo2ForCausalLM",
103 "Olmo3ForCausalLM",
104 ):
105 # Q/K norms applied on full projected vectors (before head reshape).
106 # q_norm dim = n_heads * d_head = d_model
107 self.q_norm: Optional[RMSNorm] = RMSNorm(self.cfg, self.cfg.d_model)
108 # k_norm dim depends on whether GQA is used:
109 # OLMo 2 (MHA): n_kv_heads == n_heads, so d_model
110 # OLMo 3 / OLMoE (GQA): n_kv_heads * d_head
111 if self.cfg.n_key_value_heads is not None:
112 k_norm_dim = self.cfg.d_head * self.cfg.n_key_value_heads
113 else:
114 k_norm_dim = self.cfg.d_model
115 self.k_norm: Optional[RMSNorm] = RMSNorm(self.cfg, k_norm_dim)
116 else:
117 self.q_norm = None
118 self.k_norm = None
120 self.attn_type = attn_type
121 if self.attn_type == "local":
122 if not isinstance(self.cfg.window_size, int): 122 ↛ 123line 122 didn't jump to line 123 because the condition on line 122 was never true
123 raise ValueError("Window size must be an integer for local attention")
124 elif self.attn_type != "global": 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true
125 raise ValueError(f"Invalid attention type: {self.attn_type}")
127 # Encoder-decoder models share one cfg across stacks whose directions
128 # differ, so a stack that must mask sets this rather than cfg.
129 self._attention_dir_override: Optional[str] = None
130 self._ntk_alpha_cached: float = 1.0
132 # Kept as a tiny buffer for state-dict/device compatibility. The actual
133 # causal mask is built at forward time for the current sequence length.
134 self.register_buffer("mask", torch.empty((0, 0), dtype=torch.bool))
135 self.register_buffer("IGNORE", torch.tensor(-torch.inf))
137 self.layer_id = layer_id
139 # attn_scale is a constant that we divide the attention scores by pre-softmax. I'm not entirely sure why it matters, but it's probably a mix of softmax not being scale invariant and numerical stability?
140 if self.cfg.use_attn_scale:
141 self.attn_scale = self.cfg.attn_scale # Defaults to sqrt(d_head)
142 else:
143 self.attn_scale = 1.0
144 if self.cfg.scale_attn_by_inverse_layer_idx:
145 if self.layer_id is None: # keep mypy happy 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true
146 raise ValueError("Layer ID must be provided to scale attention scores")
147 self.attn_scale *= self.layer_id + 1
149 if self.cfg.use_attention_sinks:
150 # Learned per-head sink logit (GPT-OSS); joins the softmax as an
151 # extra key column and is dropped afterward.
152 self.sinks = nn.Parameter(torch.zeros(self.cfg.n_heads, dtype=self.cfg.dtype))
154 self.hook_k = HookPoint() # [batch, pos, head_index, d_head]
155 self.hook_q = HookPoint() # [batch, pos, head_index, d_head]
156 self.hook_v = HookPoint() # [batch, pos, head_index, d_head]
157 self.hook_z = HookPoint() # [batch, pos, head_index, d_head]
158 self.hook_attn_scores = HookPoint() # [batch, head_index, query_pos, key_pos]
159 self.hook_pattern = HookPoint() # [batch, head_index, query_pos, key_pos]
160 self.hook_result = HookPoint() # [batch, pos, head_index, d_model]
162 # See HookedTransformerConfig for more details.
163 if self.cfg.positional_embedding_type == "shortformer":
164 # This tracks the input to the keys and queries, which is resid_pre + pos_embeds
165 self.hook_attn_input = HookPoint() # [batch, pos, d_model]
166 elif self.cfg.positional_embedding_type == "rotary":
167 # Applies a rotation to each two-element chunk of keys and queries pre dot producting to bake in relative position. See HookedTransformerConfig for details
168 self.hook_rot_k = HookPoint()
169 self.hook_rot_q = HookPoint()
170 if self.cfg.rotary_dim is None: # keep mypy happy 170 ↛ 171line 170 didn't jump to line 171 because the condition on line 170 was never true
171 raise ValueError("Rotary dim must be provided for rotary positional embeddings")
172 rotary_cache_size = min(self.cfg.n_ctx, self.ROTARY_INITIAL_CACHE_SIZE)
173 sin, cos = self.calculate_sin_cos_rotary(
174 self.cfg.rotary_dim,
175 rotary_cache_size,
176 base=self._rotary_base(),
177 dtype=self.cfg.dtype,
178 )
179 self.register_buffer("rotary_sin", sin)
180 self.register_buffer("rotary_cos", cos)
181 elif self.cfg.positional_embedding_type == "alibi":
182 # ALiBi bias will be constructed on the first forward pass.
183 # Note: While computationally efficient, initializing an bias with max n_ctx (16, 1024, 1024) of float32 will occupy ~256MiB of contiguous GPU memory, which may not be optimal for memory usage.
184 self.alibi = None
186 elif self.cfg.positional_embedding_type == "relative_positional_bias":
187 # will be overwritten by the child T5Attention class
188 self.has_relative_attention_bias = False
190 @property
191 def OV(self) -> FactoredMatrix:
192 """
193 OV-Circuit, as defined in A Mathematical Framework. Because there's no non-linearity between the value vector and the output of the layer, the output is purely determined by the matrix W_OV = W_V @ W_O, and not W_V or W_O individually. (Mathematically, for a single head, output == pattern @ residual @ W_V @ W_O, see the glossary for more)
195 Done in the order W_V, W_O because the paper uses left-multiplying weight matrices, and TransformerLens uses right-multiplying, sorry!
197 Returns a FactoredMatrix, with left matrix W_V [head_index, d_model, d_head] and right matrix W_O [head_index, d_head, d_model] - this is a low rank factorisation of the underlying [head_index, d_model, d_model]. FactoredMatrix has helper functions to deal with these large matrices efficiently. To get the OV circuit of a head k, attn.OV[k] works.
198 """
199 return FactoredMatrix(self.W_V, self.W_O)
201 @property
202 def QK(self) -> FactoredMatrix:
203 """
204 QK-Circuit, as defined in A Mathematical Framework. Because there's no non-linearity in the key-query dot product, the output is purely determined by the matrix W_QK = W_Q.T @ W_K, and not W_Q or W_K individually. (Mathematically, for a single head, pattern = destination_residual.T @ W_Q.T @ W_K @ source-residual, see the glossary for more).
206 Done in the order Q on the left, K on the right, because the pattern has dimensions [destination_pos, source_pos]
208 Returns a FactoredMatrix, with left matrix W_Q [head_index, d_model, d_head] and right matrix W_K.T [head_index, d_head, d_model] - this is a low rank factorisation of the underlying [head_index, d_model, d_model] matrix. FactoredMatrix has helper functions to deal with these large matrices efficiently. To get the QK circuit of a head k, attn.QK[k] works.
209 """
210 W_K_transpose = einops.rearrange(
211 self.W_K, "head_index d_model d_head -> head_index d_head d_model"
212 )
213 return FactoredMatrix(self.W_Q, W_K_transpose)
215 def forward(
216 self,
217 query_input: Union[
218 Float[torch.Tensor, "batch pos d_model"],
219 Float[torch.Tensor, "batch pos head_index d_model"],
220 ],
221 key_input: Union[
222 Float[torch.Tensor, "batch kv_pos d_model"],
223 Float[torch.Tensor, "batch kv_pos head_index d_model"],
224 Float[torch.Tensor, "batch kv_pos kv_head_index d_model"],
225 ],
226 value_input: Union[
227 Float[torch.Tensor, "batch kv_pos d_model"],
228 Float[torch.Tensor, "batch kv_pos head_index d_model"],
229 Float[torch.Tensor, "batch kv_pos kv_head_index d_model"],
230 ],
231 past_kv_cache_entry: Optional[TransformerLensKeyValueCacheEntry] = None,
232 additive_attention_mask: Optional[Float[torch.Tensor, "batch 1 1 kv_pos"]] = None,
233 attention_mask: Optional[Int[torch.Tensor, "batch offset_pos"]] = None,
234 position_bias: Optional[Float[torch.Tensor, "1 head_index pos kv_pos"]] = None,
235 ) -> Float[torch.Tensor, "batch pos d_model"]:
236 """
237 shortformer_pos_embed is only used if self.cfg.positional_embedding_type == "shortformer", else defaults to None and is irrelevant. See HookedTransformerConfig for more details
238 past_kv_cache_entry is an optional entry of past keys and values for this layer, only relevant if generating text. Defaults to None
239 additive_attention_mask is an optional mask to add to the attention weights. Defaults to None.
240 attention_mask is the attention mask for padded tokens. Defaults to None.
241 """
243 q, k, v = self.calculate_qkv_matrices(query_input, key_input, value_input)
245 # OLMo-family QK-norm: applied on full projected vectors before head reshape.
246 if self.cfg.original_architecture in (
247 "OlmoeForCausalLM",
248 "Olmo2ForCausalLM",
249 "Olmo3ForCausalLM",
250 ):
251 assert self.q_norm is not None
252 assert self.k_norm is not None
253 q = einops.rearrange(
254 self.q_norm(
255 einops.rearrange(
256 q,
257 "batch pos head_index d_head -> batch pos (head_index d_head)",
258 )
259 ),
260 "batch kv_pos (head_index d_head) -> batch kv_pos head_index d_head",
261 head_index=q.shape[2],
262 )
263 k = einops.rearrange(
264 self.k_norm(
265 einops.rearrange(
266 k,
267 "batch pos head_index d_head -> batch pos (head_index d_head)",
268 )
269 ),
270 "batch kv_pos (head_index d_head) -> batch kv_pos head_index d_head",
271 head_index=k.shape[2],
272 )
274 # OLMo v1 / OLMoE clamp Q/K/V after projection (and any qk-norm),
275 # before RoPE and cache append. Out-of-place so hooks stay legal.
276 if self.cfg.clip_qkv is not None:
277 q = q.clamp(min=-self.cfg.clip_qkv, max=self.cfg.clip_qkv)
278 k = k.clamp(min=-self.cfg.clip_qkv, max=self.cfg.clip_qkv)
279 v = v.clamp(min=-self.cfg.clip_qkv, max=self.cfg.clip_qkv)
281 if past_kv_cache_entry is not None:
282 # Appends the new keys and values to the cached values, and automatically updates the cache
283 kv_cache_pos_offset = past_kv_cache_entry.past_keys.size(1)
284 k, v = past_kv_cache_entry.append(k, v)
285 else:
286 # Not using a cache
287 kv_cache_pos_offset = 0
289 if self.cfg.positional_embedding_type == "rotary":
290 q = self.hook_rot_q(self.apply_rotary(q, kv_cache_pos_offset, attention_mask))
291 k = self.hook_rot_k(
292 self.apply_rotary(k, 0, attention_mask)
293 ) # keys are cached so no offset
295 if self.cfg.use_logn_attn and not self.training:
296 q = self._apply_logn_scaling(q, kv_cache_pos_offset)
298 attn_scores = self.calculate_attention_scores(
299 q, k
300 ) # [batch, head_index, query_pos, key_pos]
302 if self.cfg.positional_embedding_type == "alibi":
303 query_ctx = attn_scores.size(-2)
304 # The key context length is the number of positions in the past - this includes all positions in the cache
305 key_ctx = attn_scores.size(-1)
307 # only recompute when necessary to increase efficiency.
308 if self.alibi is None or key_ctx > self.alibi.size(-1):
309 self.alibi = AbstractAttention.create_alibi_bias(
310 self.cfg.n_heads, key_ctx, self.cfg.device
311 )
313 # Take the last query_ctx positions so it also works with past_kv_cache
314 if isinstance(self.alibi, torch.Tensor): 314 ↛ 317line 314 didn't jump to line 317 because the condition on line 314 was always true
315 attn_scores += self.alibi[:, -query_ctx:, :key_ctx]
316 else:
317 raise TypeError(
318 f"Expected self.alibi to be a Tensor, but got {type(self.alibi)}"
319 ) # [batch, head_index, query_pos, key_pos]
320 elif self.cfg.positional_embedding_type == "relative_positional_bias":
321 if position_bias is None:
322 if self.has_relative_attention_bias: 322 ↛ 323line 322 didn't jump to line 323 because the condition on line 322 was never true
323 raise ValueError("Positional bias is required for relative_positional_bias")
324 else:
325 position_bias = torch.zeros(
326 1,
327 self.cfg.n_heads,
328 attn_scores.shape[2],
329 attn_scores.shape[3],
330 device=attn_scores.device,
331 )
333 if position_bias is not None: # Add None check 333 ↛ 335line 333 didn't jump to line 335 because the condition on line 333 was always true
334 attn_scores += position_bias
335 if (self._attention_dir_override or self.cfg.attention_dir) == "causal":
336 # If causal attention, we mask it to only attend backwards. If bidirectional, we don't mask.
337 attn_scores = self.apply_causal_mask(
338 attn_scores, kv_cache_pos_offset, attention_mask
339 ) # [batch, head_index, query_pos, key_pos]
340 if additive_attention_mask is not None:
341 attn_scores += additive_attention_mask
343 attn_scores = self.hook_attn_scores(attn_scores)
344 if self.cfg.use_attention_sinks:
345 # The sink column is appended after hook_attn_scores so hooks keep
346 # the [batch, head, query_pos, key_pos] shape, and dropped after
347 # the softmax — every real position's weight is scaled down by the
348 # sink's share, so pattern rows sum to less than 1.
349 sink = (
350 self.sinks.reshape(1, -1, 1, 1)
351 .expand(attn_scores.shape[0], -1, attn_scores.shape[-2], -1)
352 .to(attn_scores.dtype)
353 )
354 combined = torch.cat([attn_scores, sink], dim=-1)
355 combined = combined - combined.max(dim=-1, keepdim=True).values
356 pattern = F.softmax(combined, dim=-1)[..., :-1]
357 else:
358 pattern = F.softmax(attn_scores, dim=-1)
359 if not isinstance(pattern, torch.Tensor): 359 ↛ 360line 359 didn't jump to line 360 because the condition on line 359 was never true
360 raise TypeError(f"Expected 'pattern' to be a Tensor, got {type(pattern)}")
361 pattern = torch.where(torch.isnan(pattern), torch.zeros_like(pattern), pattern)
362 pattern = self.hook_pattern(pattern) # [batch, head_index, query_pos, key_pos]
363 pattern = pattern.to(device=v.device, dtype=v.dtype)
364 z = self.calculate_z_scores(v, pattern) # [batch, pos, head_index, d_head]
365 if not self.cfg.use_attn_result:
366 if self.cfg.load_in_4bit: 366 ↛ 368line 366 didn't jump to line 368 because the condition on line 366 was never true
367 # call bitsandbytes method to dequantize and multiply
368 W_O_4bit = cast(Params4bit, self.W_O)
369 out = (
370 bnb.matmul_4bit(
371 z.reshape(z.shape[0], z.shape[1], self.cfg.d_head * self.cfg.n_heads),
372 W_O_4bit.t(),
373 bias=None,
374 quant_state=W_O_4bit.quant_state,
375 )
376 + self.b_O
377 )
378 else:
379 w = einops.rearrange(
380 self.W_O, "head_index d_head d_model -> d_model (head_index d_head)"
381 ).contiguous()
383 # Move output projection weights and bias to the same device as z
384 # so that the final linear operation occurs on the device of the inputs
385 if w.device != z.device: 385 ↛ 386line 385 didn't jump to line 386 because the condition on line 385 was never true
386 w = w.to(z.device)
387 b_O: Tensor = self.b_O
388 if b_O.device != z.device: 388 ↛ 389line 388 didn't jump to line 389 because the condition on line 388 was never true
389 b_O = b_O.to(z.device)
390 # Ensure z has the same dtype as weights used in the output projection
391 if z.dtype != w.dtype: 391 ↛ 392line 391 didn't jump to line 392 because the condition on line 391 was never true
392 z = z.to(w.dtype)
394 z = z.reshape(z.shape[0], z.shape[1], self.cfg.d_head * self.cfg.n_heads)
396 # F.linear is a fused matmul+bias that matches HuggingFace exactly,
397 # but has a bug on MPS with PyTorch 2.8 (pytorch#161640).
398 # Fall back to manual matmul on MPS to work around it.
399 if z.device.type == "mps": 399 ↛ 400line 399 didn't jump to line 400 because the condition on line 399 was never true
400 out = torch.matmul(z, w.T) + b_O
401 else:
402 out = F.linear(z, w, b_O)
403 else:
404 # Explicitly calculate the attention result so it can be accessed by a hook
405 # This is off by default because it can easily eat through your GPU memory.
406 if self.cfg.load_in_4bit: 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true
407 W_O_4bit = cast(Params4bit, self.W_O)
408 result = self.hook_result(
409 bnb.matmul_4bit(
410 z.reshape(z.shape[0], z.shape[1], self.cfg.d_head * self.cfg.n_heads),
411 W_O_4bit.t(),
412 bias=None,
413 quant_state=W_O_4bit.quant_state,
414 )
415 )
416 else:
417 # Add singleton dimensions to make shapes compatible for broadcasting:
418 w = einops.rearrange(
419 self.W_O,
420 "head_index d_head d_model -> 1 1 head_index d_head d_model",
421 )
422 if w.device != z.device: 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true
423 w = w.to(z.device)
424 # Ensure z has the same dtype as w before multiplication
425 if z.dtype != w.dtype: 425 ↛ 426line 425 didn't jump to line 426 because the condition on line 425 was never true
426 z = z.to(w.dtype)
427 z = einops.rearrange(
428 z, "batch pos head_index d_head -> batch pos head_index d_head 1"
429 )
431 unhooked_result = (z * w).sum(-2)
433 result = self.hook_result(unhooked_result) # [batch, pos, head_index, d_model]
434 out = (
435 einops.reduce(result, "batch position index model->batch position model", "sum")
436 + self.b_O
437 ) # [batch, pos, d_model]
438 return out
440 def _apply_qk_norm(
441 self, x: Float[torch.Tensor, "batch pos head_index d_head"], norm_module: RMSNorm
442 ) -> Float[torch.Tensor, "batch pos head_index d_head"]:
443 """Apply QK normalization with proper reshaping.
445 Args:
446 x: Input tensor with shape [batch, pos, head_index, d_head]
447 norm_module: RMSNorm module to apply
449 Returns:
450 Normalized tensor with same shape as input
451 """
452 # Reshape from [batch, pos, head_index, d_head] to [batch * pos * head_index, d_head]
453 d_head = x.shape[-1]
454 x_normed = norm_module(x.reshape(-1, d_head))
455 return x_normed.reshape(x.shape)
457 def calculate_qkv_matrices(
458 self,
459 query_input: Union[
460 Float[torch.Tensor, "batch pos d_model"],
461 Float[torch.Tensor, "batch pos head_index d_model"],
462 ],
463 key_input: Union[
464 Float[torch.Tensor, "batch kv_pos d_model"],
465 Float[torch.Tensor, "batch kv_pos head_index d_model"],
466 ],
467 value_input: Union[
468 Float[torch.Tensor, "batch kv_pos d_model"],
469 Float[torch.Tensor, "batch kv_pos head_index d_model"],
470 ],
471 ) -> Tuple[
472 Float[torch.Tensor, "batch pos head_index d_head"],
473 Float[torch.Tensor, "batch kv_pos head_index d_head"],
474 Float[torch.Tensor, "batch kv_pos head_index d_head"],
475 ]:
476 attn_fn = (
477 complex_attn_linear
478 if self.cfg.use_split_qkv_input or self.cfg.use_attn_in
479 else simple_attn_linear
480 )
481 if self.cfg.load_in_4bit:
482 q = self.hook_q(self._project_4bit_qkv(query_input, self.W_Q, self.b_Q))
483 else:
484 q = self.hook_q(attn_fn(query_input, self.W_Q, self.b_Q))
485 if self.cfg.load_in_4bit:
486 k = self.hook_k(self._project_4bit_qkv(key_input, self.W_K, self.b_K))
487 else:
488 k = self.hook_k(attn_fn(key_input, self.W_K, self.b_K))
490 if self.cfg.load_in_4bit:
491 v = self.hook_v(self._project_4bit_qkv(value_input, self.W_V, self.b_V))
492 else:
493 v = self.hook_v(attn_fn(value_input, self.W_V, self.b_V))
495 if self.cfg.use_qk_norm: 495 ↛ 496line 495 didn't jump to line 496 because the condition on line 495 was never true
496 assert self.q_norm is not None
497 assert self.k_norm is not None
498 q = self._apply_qk_norm(q, self.q_norm)
499 k = self._apply_qk_norm(k, self.k_norm)
501 return q, k, v
503 def _project_4bit_qkv(
504 self,
505 input: Union[
506 Float[torch.Tensor, "batch pos d_model"],
507 Float[torch.Tensor, "batch pos head_index d_model"],
508 ],
509 weight: Union[nn.Parameter, "Params4bit"],
510 bias: Float[torch.Tensor, "head_index d_head"],
511 ) -> Float[torch.Tensor, "batch pos head_index d_head"]:
512 """Project Q/K/V inputs with a 4-bit weight.
514 Split inputs dequantize once and reuse the head-wise projection path.
515 """
516 if not isinstance(weight, Params4bit): 516 ↛ 517line 516 didn't jump to line 517 because the condition on line 516 was never true
517 raise ValueError("QKV weights must be Params4bit objects if load_in_4bit is True")
519 n_heads, d_head = bias.shape
521 if input.ndim == 3:
522 projected = bnb.matmul_4bit(
523 input,
524 weight.t(),
525 bias=None,
526 quant_state=weight.quant_state,
527 )
528 return projected.reshape(input.shape[0], input.shape[1], n_heads, d_head) + bias
530 if input.ndim == 4: 530 ↛ 547line 530 didn't jump to line 547 because the condition on line 530 was always true
531 if input.shape[2] != n_heads: 531 ↛ 532line 531 didn't jump to line 532 because the condition on line 531 was never true
532 raise ValueError(
533 "4-bit split QKV inputs must have one input slice per attention head; "
534 f"got {input.shape[2]} input heads for {n_heads} projection heads."
535 )
536 dequantized_weight = bnb.functional.dequantize_4bit(
537 weight.data, quant_state=weight.quant_state
538 )
539 split_weight = einops.rearrange(
540 dequantized_weight,
541 "(head_index d_head) d_model -> head_index d_model d_head",
542 head_index=n_heads,
543 d_head=d_head,
544 )
545 return complex_attn_linear(input, split_weight, bias)
547 raise ValueError(
548 "4-bit QKV projection input must have shape [batch, pos, d_model] or "
549 "[batch, pos, head_index, d_model]."
550 )
552 def _apply_logn_scaling(self, q: torch.Tensor, kv_cache_pos_offset: int) -> torch.Tensor:
553 """Qwen-1's log-n query scaling past the training length (eval only).
555 Thresholds on the training length, never n_ctx — reaching these contexts
556 requires overriding n_ctx upward, which must not move the threshold.
557 """
558 train_length = self.cfg.train_seq_length or self.cfg.n_ctx
559 query_length = q.size(1)
560 key_length = kv_cache_pos_offset + query_length
561 if key_length <= train_length: 561 ↛ 562line 561 didn't jump to line 562 because the condition on line 561 was never true
562 return q
563 positions = torch.arange(
564 kv_cache_pos_offset + 1, key_length + 1, device=q.device, dtype=torch.float32
565 )
566 scale = torch.where(
567 positions > train_length,
568 positions.log() / math.log(train_length),
569 torch.ones((), device=q.device),
570 ).to(q.dtype)
571 return q * scale[None, :, None, None]
573 def calculate_attention_scores(
574 self,
575 q: Float[torch.Tensor, "batch query_pos head_index d_head"],
576 k: Float[torch.Tensor, "batch key_pos head_index d_head"],
577 ) -> Float[torch.Tensor, "batch head_index query_pos key_pos"]:
578 q_ = einops.rearrange(
579 q, "batch query_pos head_index d_head -> batch head_index query_pos d_head"
580 )
581 k_ = einops.rearrange(
582 k, "batch key_pos head_index d_head -> batch head_index d_head key_pos"
583 )
584 attn_scores = q_ @ k_ / self.attn_scale
585 attn_scores = apply_softcap(attn_scores, self.cfg.attn_scores_soft_cap)
586 return attn_scores
588 def calculate_z_scores(
589 self,
590 v: Float[torch.Tensor, "batch key_pos head_index d_head"],
591 pattern: Float[torch.Tensor, "batch head_index query_pos key_pos"],
592 ) -> Float[torch.Tensor, "batch query_pos head_index d_head"]:
593 v_ = einops.rearrange(
594 v, "batch key_pos head_index d_head -> batch head_index key_pos d_head"
595 )
596 pattern_ = einops.rearrange(
597 pattern,
598 "batch head_index query_pos key_pos -> batch head_index query_pos key_pos",
599 )
600 z = self.hook_z(
601 einops.rearrange(
602 pattern_ @ v_,
603 "batch head_index query_pos d_head -> batch query_pos head_index d_head",
604 )
605 )
606 return z
608 def apply_causal_mask(
609 self,
610 attn_scores: Float[torch.Tensor, "batch head_index pos pos_plus_past_kv_pos_offset"],
611 past_kv_pos_offset: int = 0,
612 attention_mask: Optional[Int[torch.Tensor, "batch offset_pos"]] = None,
613 ):
614 # The query context length is the number of positions we take queries from - if not using a past_kv_cache this is just the context length (for the current prompt), but if we're caching it can be different.
615 query_ctx_length = attn_scores.size(-2)
616 # The key context length is the number of positions in the past - this includes all positions in the cache
617 # If not caching, query_ctx_length == key_ctx_length
618 key_ctx_length = attn_scores.size(-1)
620 if query_ctx_length + past_kv_pos_offset != key_ctx_length: 620 ↛ 621line 620 didn't jump to line 621 because the condition on line 620 was never true
621 raise ValueError(
622 f"query_ctx_length {query_ctx_length} + past_kv_pos_offset {past_kv_pos_offset} != key_ctx_length {key_ctx_length} - you likely have a bug."
623 )
625 mask_device = attention_mask.device if attention_mask is not None else attn_scores.device
626 final_mask = self._make_causal_mask(
627 query_ctx_length=query_ctx_length,
628 key_ctx_length=key_ctx_length,
629 past_kv_pos_offset=past_kv_pos_offset,
630 device=mask_device,
631 )
632 if attention_mask is not None:
633 # Apply a causal mask to the attention scores considering the padding
635 # Add singleton dimensions to the attention mask to match the shape of the final mask
636 attention_mask = einops.rearrange(
637 attention_mask, "batch offset_pos -> batch 1 1 offset_pos"
638 )
640 final_mask = final_mask & attention_mask.bool() # [batch, head, pos, offset_pos]
642 attn_scores = attn_scores.to(final_mask.device)
643 ignore = cast(torch.Tensor, self.IGNORE).to(final_mask.device)
644 return torch.where(final_mask, attn_scores, ignore)
646 def _make_causal_mask(
647 self,
648 query_ctx_length: int,
649 key_ctx_length: int,
650 past_kv_pos_offset: int,
651 device: torch.device,
652 ) -> torch.Tensor:
653 """Create the causal mask for the current attention-score shape."""
654 query_positions = torch.arange(
655 past_kv_pos_offset,
656 past_kv_pos_offset + query_ctx_length,
657 device=device,
658 )
659 key_positions = torch.arange(key_ctx_length, device=device)
661 final_mask = key_positions[None, :] <= query_positions[:, None]
662 if self.attn_type == "local":
663 if not isinstance(self.cfg.window_size, int): 663 ↛ 664line 663 didn't jump to line 664 because the condition on line 663 was never true
664 raise ValueError("Window size must be an integer for local attention")
665 final_mask = final_mask & (
666 key_positions[None, :] > query_positions[:, None] - self.cfg.window_size
667 )
669 return final_mask[None, None, :, :]
671 def _rotary_base(self) -> Union[float, int]:
672 if self.cfg.rotary_base_local is not None and self.attn_type == "local":
673 return self.cfg.rotary_base_local
674 return self.cfg.rotary_base
676 def calculate_sin_cos_rotary(
677 self,
678 rotary_dim: int,
679 n_ctx: int,
680 base: Union[float, int] = 10000,
681 dtype: torch.dtype = torch.float32,
682 ) -> Tuple[Float[torch.Tensor, "n_ctx rotary_dim"], Float[torch.Tensor, "n_ctx rotary_dim"]]:
683 """
684 Calculate the sine and cosine waves to use in a rotary embedding. See https://blog.eleuther.ai/rotary-embeddings/ for details
686 Note: For some inexplicable reason, in GPT-J each ADJACENT pair of elements in k and q are rotated, in GPT-NeoX the pair of elements at k and k+n//2 are rotated (ie folding the full length in half, and then looking at pairs accordingly). I have absolutely no clue why, it should be completely equivalent.
687 To resolve this, I've coded it to default to the GPT-J mode, but to explicitly check whether it's GPT-NeoX and then do the GPT-NeoX thing if it is.
688 """
689 high_precision = torch.float32 if dtype != torch.float64 else torch.float64
690 pos = torch.arange(n_ctx, dtype=high_precision)
691 dim = torch.arange(rotary_dim // 2, dtype=high_precision)
693 use_yarn = self.cfg.use_yarn_rope and not (
694 self.cfg.yarn_global_attn_only and self.attn_type == "local"
695 )
697 # Llama-3.1 uses NTK-by-Parts Rotary Embedding introduced in Section 3.2 in https://arxiv.org/pdf/2309.00071
698 # Implementation copied from https://github.com/huggingface/transformers/blob/v4.46.0/src/transformers/modeling_rope_utils.py#L310
699 if self.cfg.use_NTK_by_parts_rope: 699 ↛ 700line 699 didn't jump to line 700 because the condition on line 699 was never true
700 inv_freq = 1.0 / (
701 base ** (torch.arange(0, rotary_dim, 2, dtype=torch.int64).float() / rotary_dim)
702 )
703 factor = self.cfg.NTK_by_parts_factor
704 low_freq_factor = self.cfg.NTK_by_parts_low_freq_factor
705 high_freq_factor = self.cfg.NTK_by_parts_high_freq_factor
706 old_context_len = self.cfg.NTK_original_ctx_len
708 low_freq_wavelen = old_context_len / low_freq_factor
709 high_freq_wavelen = old_context_len / high_freq_factor
711 wavelen = 2 * math.pi / inv_freq
712 inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)
713 smooth_factor = (old_context_len / wavelen - low_freq_factor) / (
714 high_freq_factor - low_freq_factor
715 )
716 smoothed_inv_freq = (
717 1 - smooth_factor
718 ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
719 is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
720 inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
721 freq = 1 / inv_freq_llama
722 elif use_yarn:
723 # YARN (Yet Another RoPE extensioN) from https://arxiv.org/abs/2309.00071
724 # Implementation follows HuggingFace: transformers/modeling_rope_utils.py
725 inv_freq = 1.0 / (
726 base ** (torch.arange(0, rotary_dim, 2, dtype=high_precision) / rotary_dim)
727 )
728 yarn_factor = self.cfg.yarn_factor
729 # HF uses original_max_position_embeddings (the pre-extension context length)
730 # for computing the correction range.
731 orig_max_pos = self.cfg.yarn_original_max_position_embeddings
732 beta_fast = self.cfg.yarn_beta_fast
733 beta_slow = self.cfg.yarn_beta_slow
735 def _find_correction_dim(num_rotations: float) -> float:
736 return (rotary_dim * math.log(orig_max_pos / (num_rotations * 2 * math.pi))) / (
737 2 * math.log(base)
738 )
740 low = _find_correction_dim(beta_fast)
741 high = _find_correction_dim(beta_slow)
742 if self.cfg.yarn_truncate:
743 low = math.floor(low)
744 high = math.ceil(high)
745 low = max(low, 0)
746 high = min(high, rotary_dim - 1)
748 # Linear ramp from 0 to 1 between low and high dims
749 ramp = torch.arange(rotary_dim // 2, dtype=high_precision)
750 high_f = float(high) + 0.001 if low == high else float(high)
751 ramp = torch.clamp((ramp - low) / (high_f - low), 0, 1)
753 inv_freq_interp = inv_freq / yarn_factor
754 # ramp=0 (below low) → extrapolation (original freq), ramp=1 (above high) → interpolation (scaled)
755 inv_freq = inv_freq_interp * ramp + inv_freq * (1 - ramp)
756 freq = 1.0 / inv_freq
757 else:
758 freq = base ** (dim / (rotary_dim / 2))
759 # Apply linear RoPE scaling for global attention layers if configured
760 # (e.g., Gemma 3 4B uses factor=8.0 for global layers, but not local ones)
761 scaling_factor = getattr(self.cfg, "rotary_scaling_factor", 1.0)
762 if scaling_factor != 1.0 and self.attn_type != "local": 762 ↛ 763line 762 didn't jump to line 763 because the condition on line 762 was never true
763 freq = freq * scaling_factor
764 if self.cfg.rotary_adjacent_pairs: 764 ↛ 765line 764 didn't jump to line 765 because the condition on line 764 was never true
765 freq = einops.repeat(freq, "d -> (d 2)")
766 else:
767 freq = einops.repeat(freq, "d -> (2 d)")
768 # Create a n_ctx x rotary_dim tensor, where each column is an arithmetic sequence of angles in that frequency
769 angles = pos[:, None] / freq[None, :]
770 sin, cos = torch.sin(angles).to(dtype), torch.cos(angles).to(dtype)
771 # YARN attention_factor scales the embeddings (default 1.0 is a no-op)
772 if use_yarn and self.cfg.yarn_attention_factor != 1.0:
773 sin = sin * self.cfg.yarn_attention_factor
774 cos = cos * self.cfg.yarn_attention_factor
775 return sin, cos
777 def rotate_every_two(
778 self, x: Float[torch.Tensor, "... rotary_dim"]
779 ) -> Float[torch.Tensor, "... rotary_dim"]:
780 """
781 Rotary helper function, splits x into blocks of size 2 along the final axis and maps [x0, x1] to [-x1, x0]
783 The final axis of x must have even length.
785 GPT-NeoX and GPT-J do rotary subtly differently, see calculate_sin_cos_rotary for details.
786 """
787 rot_x = x.clone()
788 if self.cfg.rotary_adjacent_pairs: 788 ↛ 789line 788 didn't jump to line 789 because the condition on line 788 was never true
789 rot_x[..., ::2] = -x[..., 1::2]
790 rot_x[..., 1::2] = x[..., ::2]
791 else:
792 n = x.size(-1) // 2
793 rot_x[..., :n] = -x[..., n:]
794 rot_x[..., n:] = x[..., :n]
796 return rot_x
798 def apply_rotary(
799 self,
800 x: Float[torch.Tensor, "batch pos head_index d_head"],
801 past_kv_pos_offset: int = 0,
802 attention_mask: Optional[Int[torch.Tensor, "batch offset_pos"]] = None,
803 ) -> Float[torch.Tensor, "batch pos head_index d_head"]:
804 # Only apply rotary to first rotary_dim dimensions (eg, if rotary_dim=64 and d_head=256, only apply to first 1/4 of dimensions)
806 if x.device != self.rotary_sin.device: 806 ↛ 807line 806 didn't jump to line 807 because the condition on line 806 was never true
807 x = x.to(cast(torch.device, self.rotary_sin.device))
809 x_pos = x.size(1)
810 if self.cfg.use_dynamic_ntk_rope and not self.training:
811 self._rescale_rotary_for_ntk(past_kv_pos_offset + x_pos)
813 x_rot = x[..., : self.cfg.rotary_dim]
814 x_pass = x[..., self.cfg.rotary_dim :]
815 x_flip = self.rotate_every_two(x_rot)
817 # Dynamically extend rotary embeddings if needed for long context
818 max_pos_needed = past_kv_pos_offset + x_pos
819 if max_pos_needed > self.rotary_cos.shape[0]:
820 new_size = min(
821 self.cfg.n_ctx,
822 max(max_pos_needed, 2 * self.rotary_cos.shape[0]),
823 )
824 self._extend_rotary_embeddings(new_size)
826 if attention_mask is None:
827 rotary_cos = cast(torch.Tensor, self.rotary_cos)[
828 None, past_kv_pos_offset : past_kv_pos_offset + x_pos, None, :
829 ]
830 rotary_sin = cast(torch.Tensor, self.rotary_sin)[
831 None, past_kv_pos_offset : past_kv_pos_offset + x_pos, None, :
832 ]
833 x_rotated = x_rot * rotary_cos + x_flip * rotary_sin
834 else:
835 offset_position_ids = get_offset_position_ids(past_kv_pos_offset, attention_mask)
836 offset_position_ids = offset_position_ids.to(cast(torch.device, self.rotary_cos.device))
837 mask_rotary_cos = cast(torch.Tensor, self.rotary_cos)[offset_position_ids, None, :]
838 mask_rotary_sin = cast(torch.Tensor, self.rotary_sin)[offset_position_ids, None, :]
839 x_rotated = x_rot * mask_rotary_cos + x_flip * mask_rotary_sin
841 return torch.cat([x_rotated, x_pass], dim=-1)
843 def _ntk_alpha(self, key_length: int, train_length: int) -> float:
844 """Qwen-1's ``get_ntk_alpha``: 1 up to the training length, then 3, 7, 15, …"""
845 if key_length <= train_length:
846 return 1.0
847 return float(max(2 ** math.ceil(math.log(key_length / train_length, 2) + 1) - 1, 1))
849 def _rescale_rotary_for_ntk(self, key_length: int) -> None:
850 """Rebuild the rotary cache on a widened base for long contexts.
852 Qwen-1 stretches the base, not the positions, so the table is rebuilt
853 rather than extended; alpha steps rarely, so this fires rarely.
854 """
855 assert self.cfg.rotary_dim is not None, "rotary_dim must be set for rotary embeddings"
856 train_length = self.cfg.train_seq_length or self.cfg.n_ctx
857 alpha = self._ntk_alpha(key_length, train_length)
858 cached_rows = self.rotary_cos.shape[0]
859 if alpha == self._ntk_alpha_cached and key_length <= cached_rows:
860 return
861 self._ntk_alpha_cached = alpha
862 base = self._rotary_base() * alpha ** (self.cfg.rotary_dim / (self.cfg.rotary_dim - 2))
863 sin, cos = self.calculate_sin_cos_rotary(
864 self.cfg.rotary_dim,
865 max(key_length, cached_rows),
866 base=base,
867 dtype=self.cfg.dtype,
868 )
869 self.rotary_sin = sin.to(self.rotary_sin.device)
870 self.rotary_cos = cos.to(self.rotary_cos.device)
872 def _extend_rotary_embeddings(self, new_size: int):
873 """Extend rotary embeddings to support longer contexts dynamically."""
874 # Ensure rotary_dim is set
875 assert self.cfg.rotary_dim is not None, "rotary_dim must be set for rotary embeddings"
877 # Calculate new embeddings
878 sin, cos = self.calculate_sin_cos_rotary(
879 self.cfg.rotary_dim,
880 new_size,
881 base=self._rotary_base(),
882 dtype=self.cfg.dtype,
883 )
885 # Update the registered buffers
886 self.rotary_sin = sin.to(self.rotary_sin.device)
887 self.rotary_cos = cos.to(self.rotary_cos.device)
889 def _extend_mask(self, new_size: int):
890 """Deprecated no-op kept for external callers."""
891 del new_size
893 def _load_from_state_dict(
894 self,
895 state_dict,
896 prefix,
897 local_metadata,
898 strict,
899 missing_keys,
900 unexpected_keys,
901 error_msgs,
902 ):
903 for buffer_name in ("mask", "rotary_sin", "rotary_cos"):
904 buffer_key = prefix + buffer_name
905 saved_buffer = state_dict.get(buffer_key)
906 current_buffer = getattr(self, buffer_name, None)
907 if (
908 isinstance(saved_buffer, torch.Tensor)
909 and isinstance(current_buffer, torch.Tensor)
910 and saved_buffer.shape != current_buffer.shape
911 ):
912 state_dict = state_dict.copy()
913 state_dict[buffer_key] = current_buffer
914 super()._load_from_state_dict(
915 state_dict,
916 prefix,
917 local_metadata,
918 strict,
919 missing_keys,
920 unexpected_keys,
921 error_msgs,
922 )
924 @staticmethod
925 def create_alibi_slope(
926 n_ctx: int, device: Optional[Union[str, torch.device]] = None
927 ) -> Float[torch.Tensor, "query key"]:
928 """Create an ALiBi Slope Matrix.
930 Create the slope matrix used in ALiBi, before it is multiplied by the head-specific scalar.
932 See :meth:`create_alibi_bias` for the full ALiBi bias calculation.
934 Examples:
936 >>> AbstractAttention.create_alibi_slope(3)
937 tensor([[ 0., 0., 0.],
938 [-1., 0., 0.],
939 [-2., -1., 0.]])
941 >>> AbstractAttention.create_alibi_slope(4)
942 tensor([[ 0., 0., 0., 0.],
943 [-1., 0., 0., 0.],
944 [-2., -1., 0., 0.],
945 [-3., -2., -1., 0.]])
947 Args:
948 n_ctx: The maximum number of tokens in a prompt.
950 Returns:
951 A tensor of shape (n_ctx, n_ctx), where the upper triangle is zero and the lower
952 triangle is decreasing by a constant slope of 1 (towards the bottom left corner).
953 """
954 # set rows as [[0,1,2...]]
955 rows = torch.arange(n_ctx, device=device).unsqueeze(0)
957 # Set cols as [[0],[1],[2]...]
958 cols = torch.arange(n_ctx, device=device).unsqueeze(1)
960 # Use broadcasting to create the desired lower triangular part of the matrix
961 slope_matrix = rows - cols
963 # Use the clamp method to set all positive values (upper right triangle) to
964 return slope_matrix.clamp(max=0).to(torch.float32)
966 @staticmethod
967 def create_alibi_multipliers(
968 n_heads: int, device: Optional[Union[str, torch.device]] = None
969 ) -> Float[torch.Tensor, "n_heads"]:
970 """Create the ALiBi Scalar Multipliers for each Head.
972 For n heads, the set of multipliers (m) is the geometric sequence that starts at 2^(-8/n), and
973 uses that same value as its ratio. For example, with 8 heads the values would be [1/(2^1),
974 1/(2^2), ... , 1/(2^8)]. With 16 heads the values would be [1/(2^0.5), 1/(2^1), ... , 1/(2^8)].
976 See :meth:`create_alibi_bias` for the full ALiBi bias calculation.
978 Examples:
980 >>> AbstractAttention.create_alibi_multipliers(8)
981 tensor([0.5000, 0.2500, 0.1250, 0.0625, 0.0312, 0.0156, 0.0078, 0.0039])
983 >>> AbstractAttention.create_alibi_multipliers(16)
984 tensor([0.7071, 0.5000, 0.3536, 0.2500, 0.1768, 0.1250, 0.0884, 0.0625, 0.0442, 0.0312,
985 0.0221, 0.0156, 0.0110, 0.0078, 0.0055, 0.0039])
987 Args:
988 n_heads: The number of heads in a layer.
989 device: The device to create the tensor on.
991 Returns:
992 A tensor of shape (n_heads,) containing the scalar multiplier for each head.
993 """
994 # Calculate the starting value
995 start = 2 ** (-8 / n_heads)
997 # Generate the indices [0, 1, ..., n_heads-1]
998 indices = torch.arange(n_heads, device=device)
1000 # Compute the multipliers, with the starting value being the same as the ratio
1001 multipliers = start * (start**indices)
1003 return multipliers
1005 @staticmethod
1006 def create_alibi_bias(
1007 n_heads: int, n_ctx: int, device: Optional[Union[torch.device, str]] = None
1008 ) -> Float[torch.Tensor, "head_idx query key"]:
1009 """Create the ALiBi Bias for all Heads.
1011 Calculate the ALiBi bias (https://arxiv.org/pdf/2108.12409.pdf) for all heads in a layer.
1013 The broad idea behind ALiBi is to remove the positional encoding from the original transformer
1014 model, and instead apply a bias to each attention score. This bias is proportional to the
1015 distance between the query and key (i.e. it encourage paying less attention to more distant
1016 tokens), and is added to the attention scores before the softmax. It is used in models such as
1017 Bloom.
1019 Examples:
1021 >>> AbstractAttention.create_alibi_bias(2, 4, torch.device('cpu'))
1022 tensor([[[ 0.0000, 0.0000, 0.0000, 0.0000],
1023 [-0.0625, 0.0000, 0.0000, 0.0000],
1024 [-0.1250, -0.0625, 0.0000, 0.0000],
1025 [-0.1875, -0.1250, -0.0625, 0.0000]],
1026 [[ 0.0000, 0.0000, 0.0000, 0.0000],
1027 [-0.0039, 0.0000, 0.0000, 0.0000],
1028 [-0.0078, -0.0039, 0.0000, 0.0000],
1029 [-0.0117, -0.0078, -0.0039, 0.0000]]])
1031 Args:
1032 n_heads: The number of heads in a layer.
1033 n_ctx: The maximum number of tokens in a prompt.
1034 device: The device to create the tensor on.
1036 Returns:
1037 The ALiBi bias that should be added to the attention scores before the softmax.
1038 """
1039 # Create the slope matrix
1040 slope: Float[torch.Tensor, "query key"] = AbstractAttention.create_alibi_slope(
1041 n_ctx, device
1042 )
1044 # Create the scalar multiplier for each head.
1045 multipliers: Float[torch.Tensor, "head_idx"] = AbstractAttention.create_alibi_multipliers(
1046 n_heads, device
1047 )
1049 # Add singleton dimensions to make shapes compatible for broadcasting:
1050 slope = einops.rearrange(slope, "query key -> 1 query key")
1051 multipliers = einops.rearrange(multipliers, "head_idx -> head_idx 1 1")
1053 # Element-wise multiplication of the slope and multipliers
1054 alibi_bias = multipliers * slope
1056 return alibi_bias