Coverage for transformer_lens/model_bridge/generalized_components/mpt_alibi_attention.py: 78%
72 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""MPT ALiBi attention bridge — MPT uses ``position_bias`` kwarg + bool causal mask."""
3from __future__ import annotations
5import math
6from typing import Any, Dict, Optional
8import torch
9from packaging import version
11from transformer_lens.model_bridge.generalized_components.alibi_joint_qkv_attention import (
12 ALiBiJointQKVAttentionBridge,
13)
14from transformer_lens.utilities.attention import clamp_qkv
16try:
17 import transformers as _transformers
19 _TRANSFORMERS_V5 = version.parse(_transformers.__version__) >= version.parse("5.0.0")
20except Exception:
21 _TRANSFORMERS_V5 = False
24def _build_mpt_alibi_tensor(num_heads: int, seq_len: int, alibi_bias_max: int = 8) -> torch.Tensor:
25 """MPT ALiBi bias [num_heads, 1, seq_len] — mirrors HF's ``build_mpt_alibi_tensor``."""
26 alibi = torch.arange(1 - seq_len, 1, dtype=torch.int32).view(1, 1, 1, seq_len)
27 num_heads_power_of_2 = 2 ** math.ceil(math.log2(num_heads))
29 base = torch.arange(1, num_heads_power_of_2 + 1, dtype=torch.int64).float()
30 base = base * (alibi_bias_max / num_heads_power_of_2)
31 slopes = 1.0 / torch.pow(2, base)
32 slopes = slopes.view(1, num_heads_power_of_2, 1, 1)
34 if num_heads_power_of_2 != num_heads: 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true
35 slopes = torch.concat([slopes[:, 1::2, ...], slopes[:, ::2, ...]], dim=1)[
36 :, :num_heads, ...
37 ]
39 alibi = alibi * slopes # [1, n_heads, 1, seq_len]
40 return alibi.squeeze(0) # [n_heads, 1, seq_len]
43class MPTALiBiAttentionBridge(ALiBiJointQKVAttentionBridge):
44 """ALiBi bridge for MPT: overrides ALiBi kwarg name, bias shape, mask format, and clip_qkv."""
46 _clip_qkv: Optional[float] = None
47 _softmax_scale: Optional[float] = None
49 def forward(
50 self, *args: Any, **kwargs: Any
51 ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, None]:
52 """2-tuple on transformers>=5, 3-tuple on <5 — MptBlock unpack arity changed in v5."""
53 output, attn_weights = super().forward(*args, **kwargs)
54 if _TRANSFORMERS_V5: 54 ↛ 56line 54 didn't jump to line 56 because the condition on line 54 was always true
55 return output, attn_weights
56 return output, attn_weights, None
58 def set_original_component(self, original_component: torch.nn.Module) -> None:
59 super().set_original_component(original_component)
60 if hasattr(self, "o") and hasattr(original_component, "out_proj"): 60 ↛ 62line 60 didn't jump to line 62 because the condition on line 60 was always true
61 self.o.set_original_component(original_component.out_proj)
62 clip = getattr(original_component, "clip_qkv", None)
63 self._clip_qkv = float(clip) if clip is not None else None
64 # HF resolves config.attn_config.softmax_scale (or 1/sqrt(head_dim))
65 # onto the module at init; honor it instead of recomputing the default.
66 scale = getattr(original_component, "softmax_scale", None)
67 self._softmax_scale = float(scale) if scale is not None else None
69 def _reconstruct_attention(
70 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: Any
71 ) -> tuple[torch.Tensor, torch.Tensor]:
72 # clip_qkv is post-projection, pre-head-split — must happen before reshape.
73 # Truthiness gate matches HF (`if self.clip_qkv:`): 0.0 means disabled.
74 if self._clip_qkv:
75 q, k, v = clamp_qkv(q, k, v, self._clip_qkv)
77 num_heads = self.config.n_heads if self.config else 32
78 q, k, v, batch_size, seq_len, head_dim = self._reshape_qkv_to_heads(
79 q, k, v, num_heads, num_heads
80 )
82 softmax_scale = self._softmax_scale if self._softmax_scale is not None else head_dim**-0.5
83 attn_scores = torch.matmul(q, k.transpose(-2, -1)) * softmax_scale
85 # position_bias is [n_heads, 1, max_seq_len]; slice trailing kv_len, broadcast over batch.
86 position_bias = kwargs.get("position_bias", None)
87 if position_bias is not None:
88 kv_len = attn_scores.shape[-1]
89 pb = position_bias[:, :, -kv_len:]
90 attn_scores = attn_scores + pb.unsqueeze(0)
92 # MPT passes a bool 4D mask (True = masked), not an additive float mask.
93 attention_mask = kwargs.get("attention_mask", None)
94 if attention_mask is not None:
95 mask_value = (
96 -torch.inf if self.compatibility_mode else torch.finfo(attn_scores.dtype).min
97 )
98 attn_scores = attn_scores.masked_fill(attention_mask, mask_value)
100 attn_scores = self.hook_attn_scores(attn_scores)
102 attn_weights = self._softmax_dropout_pattern(
103 attn_scores, upcast_to_fp32=True, target_dtype=q.dtype
104 )
106 attn_output = torch.matmul(attn_weights, v)
107 attn_output = self._reshape_attn_output(
108 attn_output, batch_size, seq_len, num_heads, head_dim
109 )
110 attn_output = self._apply_output_projection(attn_output)
111 return attn_output, attn_weights
113 def get_random_inputs(
114 self,
115 batch_size: int = 2,
116 seq_len: int = 8,
117 device: Optional[torch.device] = None,
118 dtype: Optional[torch.dtype] = None,
119 ) -> Dict[str, Any]:
120 """Test inputs using MPT's kwarg names: position_bias (no batch dim) + bool causal mask."""
121 if device is None:
122 device = torch.device("cpu")
123 if dtype is None:
124 dtype = torch.float32
126 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 2048
127 num_heads = self.config.n_heads if self.config and hasattr(self.config, "n_heads") else 32
129 position_bias = _build_mpt_alibi_tensor(num_heads, seq_len).to(device=device, dtype=dtype)
131 causal = torch.triu(
132 torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1
133 )
134 causal_mask = causal.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, -1, -1)
136 return {
137 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype),
138 "position_bias": position_bias,
139 "attention_mask": causal_mask,
140 }