Coverage for transformer_lens/model_bridge/generalized_components/ssm_mixer.py: 93%
154 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Wrap-don't-reimplement bridge for HF's MambaMixer (Mamba-1), plus S6 effective attention."""
2from typing import Any, NamedTuple, Optional
4import torch
6from transformer_lens.ActivationCache import ActivationCache
7from transformer_lens.hook_points import HookPoint
8from transformer_lens.model_bridge.generalized_components.base import (
9 GeneralizedComponent,
10)
11from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
12 SSMStateHookMixin,
13)
16def _weightless_rms(x: torch.Tensor, eps: float) -> torch.Tensor:
17 """FalconMamba's parameter-free RMS over the last dim (HF rms_forward)."""
18 dtype = x.dtype
19 x = x.float()
20 return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)).to(dtype)
23class _S6Terms(NamedTuple):
24 """S6 intermediates reconstructed from cached Mamba-1 hooks (per-channel)."""
26 dt: torch.Tensor # [batch, channels, seq]
27 decay: torch.Tensor # [batch, channels, state, i, j] causal (upper triangle zero)
28 B: torch.Tensor # [batch, seq, state] (shared across channels)
29 C: torch.Tensor # [batch, seq, state] (shared across channels)
32class SSMMixerBridge(SSMStateHookMixin, GeneralizedComponent):
33 """Opaque wrapper around Mamba-1's MambaMixer.
35 Submodules (in_proj, conv1d, x_proj, dt_proj, out_proj) are swapped into
36 the HF mixer by ``replace_remote_component``, so their hooks fire when
37 slow_forward accesses them. ``A_log`` and ``D`` reach the user via
38 ``GeneralizedComponent.__getattr__`` delegation.
40 Decode-step caveat: ``conv1d.hook_out`` fires only on prefill during
41 stateful generation; see ``DepthwiseConv1DBridge`` for the reason.
42 """
44 hook_aliases = {
45 "hook_in_proj": "in_proj.hook_out",
46 "hook_conv": "conv1d.hook_out",
47 "hook_x_proj": "x_proj.hook_out",
48 "hook_dt_proj": "dt_proj.hook_out",
49 "hook_ssm_out": "hook_out",
50 # Canonical SSM vocabulary (additive): Mamba-1 exposes the discrete time
51 # step via dt_proj. B/C are bundled in x_proj (not separately hooked).
52 "hook_ssm_dt": "dt_proj.hook_out",
53 }
55 def __init__(self, *args: Any, **kwargs: Any) -> None:
56 super().__init__(*args, **kwargs) # mixin adds hook_ssm_state + eager_scan
57 # Real per-step write term dt·x·B, per-channel [batch, channels, seq, state].
58 self.hook_ssm_write = HookPoint()
60 def forward(self, *args: Any, **kwargs: Any) -> Any:
61 """Hook the input, run HF slow_forward (or the eager scan), hook the output."""
62 if self.original_component is None: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true
63 raise RuntimeError(
64 f"Original component not set for {self.name}. "
65 "Call set_original_component() first."
66 )
68 hidden_states: Optional[torch.Tensor] = None
69 if len(args) > 0 and isinstance(args[0], torch.Tensor):
70 hidden_states = self.hook_in(args[0])
71 args = (hidden_states,) + args[1:]
72 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): 72 ↛ 78line 72 didn't jump to line 78 because the condition on line 72 was always true
73 hidden_states = self.hook_in(kwargs["hidden_states"])
74 kwargs["hidden_states"] = hidden_states
76 # Eager-scan slow path: opt-in, prefill only (no cache_params). Keyed on
77 # explicit state, never on hook-registry introspection.
78 if self.eager_scan and hidden_states is not None and kwargs.get("cache_params") is None:
79 output: Any = self._eager_scan_forward(hidden_states, kwargs.get("attention_mask"))
80 else:
81 output = self.original_component(*args, **kwargs)
83 # Hook the primary output tensor, preserving tuple structure
84 if isinstance(output, tuple) and len(output) > 0: 84 ↛ 85line 84 didn't jump to line 85 because the condition on line 84 was never true
85 first = output[0]
86 if isinstance(first, torch.Tensor):
87 return (self.hook_out(first),) + output[1:]
88 return output
89 if isinstance(output, torch.Tensor): 89 ↛ 91line 89 didn't jump to line 91 because the condition on line 89 was always true
90 return self.hook_out(output)
91 return output
93 def _eager_scan_forward(
94 self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]
95 ) -> torch.Tensor:
96 """Reimplement the Mamba-1 mixer prefill forward with an eager Python S6 scan.
98 Reuses HF's in_proj / conv1d / x_proj / dt_proj / out_proj submodules (so
99 their hooks still fire) and reimplements ONLY the S6 recurrence::
101 write_t[c, s] = dt_t[c] · x_t[c] · B_t[s] -> hook_ssm_write [b, c, seq, s]
102 S_t[c, s] = exp(A[c,s]·dt_t[c]) · S_{t-1} + write_t
103 S = stack_t S_t -> hook_ssm_state [b, c, seq, s]
104 y_t[c] = sum_s C_t[s] · S_t[c, s] + D[c] · x_t[c]
106 Intervening on hook_ssm_write re-runs the recurrence (propagates to later
107 states); hook_ssm_state is the post-scan trajectory, so a patch there
108 changes only the same-position output y_t = C_t·S_t — patch hook_ssm_write
109 for a propagating state edit.
111 When the wrapped mixer exposes ``dt_layernorm`` / ``b_layernorm`` /
112 ``c_layernorm`` (Jamba), those norms run after ``x_proj`` so the scan
113 matches HF's selective-param path.
115 Kernel-divergence caveat: reproduces HF's scan only to fp tolerance
116 (≈1e-6 fp32), never bit-for-bit. O(seq) Python; materializes an
117 O(b·channels·seq·state) write/state tensor. Prefill only.
118 """
119 oc: Any = self.original_component
120 batch, seq_len, _ = hidden_states.shape
121 d_inner = oc.intermediate_size
122 state = oc.ssm_state_size
123 dt_rank = oc.time_step_rank
125 def _mask_cf(states: torch.Tensor) -> torch.Tensor:
126 # HF MambaMixer masks whenever attention_mask is present (no batch guard,
127 # unlike Mamba-2), so match that — batch-1 padded inputs included.
128 if attention_mask is not None:
129 return states * attention_mask.unsqueeze(1).to(states.dtype)
130 return states
132 # 1-2. in_proj + gate split, conv — reuse the HF submodule bridges (hooks fire).
133 # HF masks padding on the channel-first x both before and after the conv.
134 projected = oc.in_proj(hidden_states).transpose(1, 2) # [b, 2*d_inner, seq]
135 x_raw, gate = projected.chunk(2, dim=1) # [b, d_inner, seq]
136 x = oc.act(oc.conv1d(_mask_cf(x_raw))[..., :seq_len]) # [b, d_inner, seq]
137 x = _mask_cf(x)
139 # 3. x_proj -> dt low-rank / B / C ; dt_proj -> dt (softplus, channel-first).
140 # Jamba (and forks) apply RMSNorm to dt/B/C after x_proj; stock Mamba-1 does not.
141 ssm_params = oc.x_proj(x.transpose(1, 2)) # [b, seq, dt_rank + 2*state]
142 time_step, B, C = ssm_params.split([dt_rank, state, state], dim=-1) # B, C: [b, seq, state]
143 # FalconMamba applies a parameter-free RMS to B/C/dt before dt_proj.
144 rms_eps = getattr(oc, "rms_eps", None)
145 if rms_eps is not None:
146 B = _weightless_rms(B, rms_eps)
147 C = _weightless_rms(C, rms_eps)
148 time_step = _weightless_rms(time_step, rms_eps)
150 # Jamba adds module-based dt/B/C layernorms (no-op when absent).
151 time_step, B, C = self._apply_selective_param_norms(oc, time_step, B, C)
152 dt = torch.nn.functional.softplus(oc.dt_proj(time_step)).transpose(1, 2).float()
153 A = -torch.exp(self.A_log.float()) # [d_inner, state]
154 x_f, B_f, C_f = x.float(), B.float(), C.float()
156 # 4. Eager recurrence with intervention hooks.
157 writes = (dt * x_f)[:, :, :, None] * B_f[:, None, :, :] # [b, d_inner, seq, state]
158 writes = self.hook_ssm_write(writes)
160 ssm_state = torch.zeros(batch, d_inner, state, dtype=writes.dtype, device=writes.device)
161 states = []
162 for t in range(seq_len):
163 decay = torch.exp(A[None] * dt[:, :, t, None]) # [batch, d_inner, state]
164 ssm_state = decay * ssm_state + writes[:, :, t]
165 states.append(ssm_state)
166 state_traj = torch.stack(states, dim=2) # [batch, d_inner, seq, state]
167 state_traj = self.hook_ssm_state(state_traj)
169 y = torch.einsum("bcts,bts->bct", state_traj, C_f) # [batch, d_inner, seq]
170 y = y + self.D.float()[None, :, None] * x_f
171 scan_output = y * oc.act(gate.float()) # HF gates with self.act, not hardcoded silu
173 # 5. Output projection — reuse the HF submodule bridge.
174 contextualized: torch.Tensor = oc.out_proj(
175 scan_output.transpose(1, 2).to(hidden_states.dtype)
176 )
177 return contextualized
179 def compute_effective_attention(
180 self,
181 cache: ActivationCache,
182 layer_idx: int,
183 include_dt_scaling: bool = False,
184 per_state_coord: bool = False,
185 ) -> torch.Tensor:
186 """Materialize Mamba-1's per-channel effective attention from cached hooks.
188 Mamba-1's S6 selective scan is equivalent to causal attention with a
189 per-channel, per-state-coordinate learned decay — see "The Hidden
190 Attention of Mamba" (Ali et al., ACL 2025). Unlike Mamba-2 there is no
191 head grouping: each of the ``intermediate_size`` channels has its own
192 ``A`` row, so the "head" axis here is the channel axis::
194 M[c, i, j] = sum_n C[i, n] · prod_{k=j+1..i} exp(A[c,n]·dt[c,k]) · B[j, n]
196 Reads B/C from ``x_proj.hook_out`` (or post-norm
197 ``b_layernorm`` / ``c_layernorm`` hooks when present, as on Jamba) and
198 dt from ``dt_proj.hook_out`` (softplus of that output); A via
199 ``__getattr__``. Read-only: no ``forward()`` re-run.
201 Args:
202 cache: ActivationCache from ``run_with_cache`` with this layer's
203 x_proj and dt_proj hooks.
204 layer_idx: Block index for this mixer.
205 include_dt_scaling: False (default) returns the attention-like form;
206 True multiplies column j by dt[c, j], giving the reconstruction
207 form satisfying ``y[c,i] = sum_j M[c,i,j]·x[c,j] + D[c]·x[c,i]``
208 (x is the post-conv SiLU input; y the pre-gate scan output).
209 per_state_coord: False (default) sums over the state coordinate and
210 returns ``[batch, channels, seq, seq]``. True returns the
211 unsummed ``[batch, channels, state, seq, seq]`` tensor (the
212 paper's D·N matrices) — OFF by default.
214 Returns:
215 ``[batch, intermediate_size, seq, seq]``, or
216 ``[batch, intermediate_size, state_size, seq, seq]`` when
217 ``per_state_coord`` is True.
219 Peak memory is O(batch · intermediate_size · state_size · seq²) — the
220 per-(channel, state) decay tensor — even for the summed default; use on
221 short sequences.
222 """
223 t = self._s6_terms(cache, layer_idx)
225 # M_coord[c, s, i, j] = C[i, s] · decay[c, s, i, j] · B[j, s]
226 C_col = t.C.permute(0, 2, 1)[:, None, :, :, None] # [batch, 1, state, i, 1]
227 B_col = t.B.permute(0, 2, 1)[:, None, :, None, :] # [batch, 1, state, 1, j]
228 M_coord = C_col * t.decay * B_col # [batch, channels, state, i, j]
230 if include_dt_scaling:
231 M_coord = M_coord * t.dt[:, :, None, None, :] # × dt[c, j]
233 if per_state_coord:
234 return M_coord
235 return M_coord.sum(dim=2) # [batch, channels, seq, seq]
237 def compute_ssm_state(
238 self,
239 cache: ActivationCache,
240 layer_idx: int,
241 time_step: Optional[int] = None,
242 ) -> torch.Tensor:
243 """Reconstruct Mamba-1's recurrent state ``S`` from cached hook values.
245 S6 recurrence ``h_t[c,s] = exp(A[c,s]·dt_t[c])·h_{t-1}[c,s] + dt_t[c]·x_t[c]·B_t[s]``
246 unrolls to::
248 S_t[c, s] = sum_{j<=t} decay[c, s, t, j] · dt_j[c] · x_j[c] · B_j[s]
250 with the same per-(channel, state) ``decay`` used by
251 ``compute_effective_attention``. ``x`` is the post-conv SiLU input
252 (``SiLU(conv1d.hook_out)``). Read-only: no ``forward()`` re-run. Verify
253 with ``y_t[c] = sum_s C_t[s]·S_t[c,s] + D[c]·x_t[c]``.
255 On padded batches the cached hooks are unmasked, so ``S`` is exact only at
256 non-pad positions; pad-position state is out of contract.
258 Args:
259 cache: ActivationCache from ``run_with_cache`` with this layer's
260 x_proj, dt_proj, and conv1d hooks.
261 layer_idx: Block index for this mixer.
262 time_step: If given, return only ``S`` at that position
263 (``[batch, channels, state]``); None returns every step.
265 Returns:
266 ``[batch, channels, seq, state]`` for all steps, or
267 ``[batch, channels, state]`` for a single ``time_step``.
269 Peak memory is O(batch · channels · state · seq²) — the per-(channel,
270 state) decay tensor, built in full by ``_s6_terms`` regardless of
271 ``time_step`` (which bounds only the returned tensor, not this peak); use
272 on short sequences.
273 """
274 conv_key = f"blocks.{layer_idx}.mixer.conv1d.hook_out"
275 if conv_key not in cache:
276 raise RuntimeError(
277 f"compute_ssm_state needs {conv_key!r} in cache. Run "
278 "`run_with_cache()` on the bridge before calling this method."
279 )
280 t = self._s6_terms(cache, layer_idx)
281 seq_len = t.dt.shape[-1]
282 oc: Any = self.original_component
283 # x = act(conv output), the SSM scan input; channel-first [batch, channels, seq].
284 x = oc.act(cache[conv_key].float()[..., :seq_len])
285 dtx = t.dt * x # dt_j[c]·x_j[c], [batch, channels, seq]
286 # S_t[c, s] = sum_{j<=t} decay[c, s, t, j] · dtx[c, j] · B_j[s]
287 if time_step is not None:
288 return torch.einsum("bcsj,bcj,bjs->bcs", t.decay[:, :, :, time_step, :], dtx, t.B)
289 return torch.einsum("bcsij,bcj,bjs->bcis", t.decay, dtx, t.B)
291 @staticmethod
292 def _apply_selective_param_norms(
293 oc: Any,
294 time_step: torch.Tensor,
295 B: torch.Tensor,
296 C: torch.Tensor,
297 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
298 """Apply Jamba-style RMSNorms on dt/B/C when the HF mixer exposes them.
300 Stock Mamba-1 has no ``dt_layernorm`` / ``b_layernorm`` / ``c_layernorm``;
301 Jamba's ``JambaMambaMixer`` does. Calling through the (possibly bridged)
302 attributes keeps their hooks firing on the eager-scan path.
303 """
304 dt_ln = getattr(oc, "dt_layernorm", None)
305 if dt_ln is not None:
306 time_step = dt_ln(time_step)
307 b_ln = getattr(oc, "b_layernorm", None)
308 if b_ln is not None:
309 B = b_ln(B)
310 c_ln = getattr(oc, "c_layernorm", None)
311 if c_ln is not None:
312 C = c_ln(C)
313 return time_step, B, C
315 def _s6_terms(self, cache: ActivationCache, layer_idx: int) -> _S6Terms:
316 """Reconstruct the shared S6 intermediates (dt, per-(channel,state) decay, B, C).
318 Reused by ``compute_effective_attention`` and ``compute_ssm_state``. Reads
319 B/C from ``x_proj.hook_out`` (shared across channels), or — when present —
320 from Jamba's post-norm ``b_layernorm`` / ``c_layernorm`` hooks. ``dt`` comes
321 from ``dt_proj.hook_out`` (softplus); that already sits after Jamba's
322 ``dt_layernorm``. A via ``__getattr__``. Dims come from the wrapped HF
323 mixer; cfg is the fallback (mirrors the Mamba-2 bridge).
324 """
325 if self.config is None: 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true
326 raise RuntimeError("SSMMixerBridge.config must be set")
328 x_proj_key = f"blocks.{layer_idx}.mixer.x_proj.hook_out"
329 dt_proj_key = f"blocks.{layer_idx}.mixer.dt_proj.hook_out"
330 for key in (x_proj_key, dt_proj_key):
331 if key not in cache:
332 raise RuntimeError(
333 f"S6 reconstruction needs {key!r} in cache. Run "
334 "`run_with_cache()` on the bridge before calling this."
335 )
337 cfg = self.config
338 oc = self.original_component
340 def _mamba_dim(module_attr: str, cfg_attr: str, default: Any) -> Any:
341 if oc is not None: 341 ↛ 345line 341 didn't jump to line 345 because the condition on line 341 was always true
342 val = getattr(oc, module_attr, None)
343 if val is not None: 343 ↛ 345line 343 didn't jump to line 345 because the condition on line 343 was always true
344 return val
345 return getattr(cfg, cfg_attr, default)
347 state_size = int(_mamba_dim("ssm_state_size", "state_size", 16))
348 dt_rank = int(_mamba_dim("time_step_rank", "time_step_rank", 0))
350 x_proj_out = cache[x_proj_key].float() # [batch, seq, dt_rank + 2*state]
351 dt_proj_out = cache[dt_proj_key].float() # [batch, seq, channels]
352 seq_len = dt_proj_out.shape[1]
354 # Prefer post-norm B/C when the adapter mapped Jamba's selective-param LNs.
355 b_ln_key = f"blocks.{layer_idx}.mixer.b_layernorm.hook_out"
356 c_ln_key = f"blocks.{layer_idx}.mixer.c_layernorm.hook_out"
357 if b_ln_key in cache and c_ln_key in cache:
358 B = cache[b_ln_key].float()
359 C = cache[c_ln_key].float()
360 else:
361 # Stock Mamba-1: B, C are the raw x_proj tails (shared across channels).
362 _time_step, B, C = x_proj_out.split([dt_rank, state_size, state_size], dim=-1)
364 # FalconMamba normalizes B/C after x_proj; dt_proj.hook_out is already
365 # post-RMS (HF applied it before dt_proj), so dt needs no correction.
366 rms_eps = getattr(oc, "rms_eps", None) if oc is not None else None
367 if rms_eps is not None:
368 B = _weightless_rms(B, rms_eps)
369 C = _weightless_rms(C, rms_eps)
371 dt = torch.nn.functional.softplus(dt_proj_out).transpose(1, 2) # [batch, channels, seq]
372 A = -torch.exp(self.A_log.float()) # [channels, state]
374 # decay[c, s, i, j] = exp(A[c,s]·(cumdt[c,i]-cumdt[c,j])) for i >= j, else 0.
375 cumdt = torch.cumsum(dt, dim=-1) # [batch, channels, seq]
376 dcs = cumdt[:, :, :, None] - cumdt[:, :, None, :] # [batch, channels, i, j]
377 decay_exp = (
378 A[None, :, :, None, None] * dcs[:, :, None, :, :]
379 ) # [batch, channels, state, i, j]
380 causal_mask = torch.tril(
381 torch.ones(seq_len, seq_len, dtype=torch.bool, device=decay_exp.device)
382 )
383 decay = torch.where(
384 causal_mask[None, None, None, :, :],
385 torch.exp(decay_exp),
386 torch.zeros_like(decay_exp),
387 )
388 return _S6Terms(dt=dt, decay=decay, B=B, C=C)