Coverage for transformer_lens/model_bridge/generalized_components/ssm_mixer.py: 92%
156 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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()
59 self.hook_ssm_write.pos_dim = 2
60 self.hook_ssm_state.pos_dim = 2
62 def forward(self, *args: Any, **kwargs: Any) -> Any:
63 """Hook the input, run HF slow_forward (or the eager scan), hook the output."""
64 if self.original_component is None: 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true
65 raise RuntimeError(
66 f"Original component not set for {self.name}. "
67 "Call set_original_component() first."
68 )
70 hidden_states: Optional[torch.Tensor] = None
71 if len(args) > 0 and isinstance(args[0], torch.Tensor):
72 hidden_states = self.hook_in(args[0])
73 args = (hidden_states,) + args[1:]
74 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): 74 ↛ 80line 74 didn't jump to line 80 because the condition on line 74 was always true
75 hidden_states = self.hook_in(kwargs["hidden_states"])
76 kwargs["hidden_states"] = hidden_states
78 # Eager-scan slow path: opt-in, prefill only (no cache_params). Keyed on
79 # explicit state, never on hook-registry introspection.
80 if self.eager_scan and hidden_states is not None and kwargs.get("cache_params") is None:
81 output: Any = self._eager_scan_forward(hidden_states, kwargs.get("attention_mask"))
82 else:
83 output = self.original_component(*args, **kwargs)
85 # Hook the primary output tensor, preserving tuple structure
86 if isinstance(output, tuple) and len(output) > 0: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true
87 first = output[0]
88 if isinstance(first, torch.Tensor):
89 return (self.hook_out(first),) + output[1:]
90 return output
91 if isinstance(output, torch.Tensor): 91 ↛ 93line 91 didn't jump to line 93 because the condition on line 91 was always true
92 return self.hook_out(output)
93 return output
95 def _eager_scan_forward(
96 self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]
97 ) -> torch.Tensor:
98 """Reimplement the Mamba-1 mixer prefill forward with an eager Python S6 scan.
100 Reuses HF's in_proj / conv1d / x_proj / dt_proj / out_proj submodules (so
101 their hooks still fire) and reimplements ONLY the S6 recurrence::
103 write_t[c, s] = dt_t[c] · x_t[c] · B_t[s] -> hook_ssm_write [b, c, seq, s]
104 S_t[c, s] = exp(A[c,s]·dt_t[c]) · S_{t-1} + write_t
105 S = stack_t S_t -> hook_ssm_state [b, c, seq, s]
106 y_t[c] = sum_s C_t[s] · S_t[c, s] + D[c] · x_t[c]
108 Intervening on hook_ssm_write re-runs the recurrence (propagates to later
109 states); hook_ssm_state is the post-scan trajectory, so a patch there
110 changes only the same-position output y_t = C_t·S_t — patch hook_ssm_write
111 for a propagating state edit.
113 When the wrapped mixer exposes ``dt_layernorm`` / ``b_layernorm`` /
114 ``c_layernorm`` (Jamba), those norms run after ``x_proj`` so the scan
115 matches HF's selective-param path.
117 Kernel-divergence caveat: reproduces HF's scan only to fp tolerance
118 (≈1e-6 fp32), never bit-for-bit. O(seq) Python; materializes an
119 O(b·channels·seq·state) write/state tensor. Prefill only.
120 """
121 oc: Any = self.original_component
122 batch, seq_len, _ = hidden_states.shape
123 d_inner = oc.intermediate_size
124 state = oc.ssm_state_size
125 dt_rank = oc.time_step_rank
127 def _mask_cf(states: torch.Tensor) -> torch.Tensor:
128 # HF MambaMixer masks whenever attention_mask is present (no batch guard,
129 # unlike Mamba-2), so match that — batch-1 padded inputs included.
130 if attention_mask is not None:
131 return states * attention_mask.unsqueeze(1).to(states.dtype)
132 return states
134 # 1-2. in_proj + gate split, conv — reuse the HF submodule bridges (hooks fire).
135 # HF masks padding on the channel-first x both before and after the conv.
136 projected = oc.in_proj(hidden_states).transpose(1, 2) # [b, 2*d_inner, seq]
137 x_raw, gate = projected.chunk(2, dim=1) # [b, d_inner, seq]
138 x = oc.act(oc.conv1d(_mask_cf(x_raw))[..., :seq_len]) # [b, d_inner, seq]
139 x = _mask_cf(x)
141 # 3. x_proj -> dt low-rank / B / C ; dt_proj -> dt (softplus, channel-first).
142 # Jamba (and forks) apply RMSNorm to dt/B/C after x_proj; stock Mamba-1 does not.
143 ssm_params = oc.x_proj(x.transpose(1, 2)) # [b, seq, dt_rank + 2*state]
144 time_step, B, C = ssm_params.split([dt_rank, state, state], dim=-1) # B, C: [b, seq, state]
145 # FalconMamba applies a parameter-free RMS to B/C/dt before dt_proj.
146 rms_eps = getattr(oc, "rms_eps", None)
147 if rms_eps is not None:
148 B = _weightless_rms(B, rms_eps)
149 C = _weightless_rms(C, rms_eps)
150 time_step = _weightless_rms(time_step, rms_eps)
152 # Jamba adds module-based dt/B/C layernorms (no-op when absent).
153 time_step, B, C = self._apply_selective_param_norms(oc, time_step, B, C)
154 dt = torch.nn.functional.softplus(oc.dt_proj(time_step)).transpose(1, 2).float()
155 A = -torch.exp(self.A_log.float()) # [d_inner, state]
156 x_f, B_f, C_f = x.float(), B.float(), C.float()
158 # 4. Eager recurrence with intervention hooks.
159 writes = (dt * x_f)[:, :, :, None] * B_f[:, None, :, :] # [b, d_inner, seq, state]
160 writes = self.hook_ssm_write(writes)
162 ssm_state = torch.zeros(batch, d_inner, state, dtype=writes.dtype, device=writes.device)
163 states = []
164 for t in range(seq_len):
165 decay = torch.exp(A[None] * dt[:, :, t, None]) # [batch, d_inner, state]
166 ssm_state = decay * ssm_state + writes[:, :, t]
167 states.append(ssm_state)
168 state_traj = torch.stack(states, dim=2) # [batch, d_inner, seq, state]
169 state_traj = self.hook_ssm_state(state_traj)
171 y = torch.einsum("bcts,bts->bct", state_traj, C_f) # [batch, d_inner, seq]
172 y = y + self.D.float()[None, :, None] * x_f
173 scan_output = y * oc.act(gate.float()) # HF gates with self.act, not hardcoded silu
175 # 5. Output projection — reuse the HF submodule bridge.
176 contextualized: torch.Tensor = oc.out_proj(
177 scan_output.transpose(1, 2).to(hidden_states.dtype)
178 )
179 return contextualized
181 def compute_effective_attention(
182 self,
183 cache: ActivationCache,
184 layer_idx: int,
185 include_dt_scaling: bool = False,
186 per_state_coord: bool = False,
187 ) -> torch.Tensor:
188 """Materialize Mamba-1's per-channel effective attention from cached hooks.
190 Mamba-1's S6 selective scan is equivalent to causal attention with a
191 per-channel, per-state-coordinate learned decay — see "The Hidden
192 Attention of Mamba" (Ali et al., ACL 2025). Unlike Mamba-2 there is no
193 head grouping: each of the ``intermediate_size`` channels has its own
194 ``A`` row, so the "head" axis here is the channel axis::
196 M[c, i, j] = sum_n C[i, n] · prod_{k=j+1..i} exp(A[c,n]·dt[c,k]) · B[j, n]
198 Reads B/C from ``x_proj.hook_out`` (or post-norm
199 ``b_layernorm`` / ``c_layernorm`` hooks when present, as on Jamba) and
200 dt from ``dt_proj.hook_out`` (softplus of that output); A via
201 ``__getattr__``. Read-only: no ``forward()`` re-run.
203 Args:
204 cache: ActivationCache from ``run_with_cache`` with this layer's
205 x_proj and dt_proj hooks.
206 layer_idx: Block index for this mixer.
207 include_dt_scaling: False (default) returns the attention-like form;
208 True multiplies column j by dt[c, j], giving the reconstruction
209 form satisfying ``y[c,i] = sum_j M[c,i,j]·x[c,j] + D[c]·x[c,i]``
210 (x is the post-conv SiLU input; y the pre-gate scan output).
211 per_state_coord: False (default) sums over the state coordinate and
212 returns ``[batch, channels, seq, seq]``. True returns the
213 unsummed ``[batch, channels, state, seq, seq]`` tensor (the
214 paper's D·N matrices) — OFF by default.
216 Returns:
217 ``[batch, intermediate_size, seq, seq]``, or
218 ``[batch, intermediate_size, state_size, seq, seq]`` when
219 ``per_state_coord`` is True.
221 Peak memory is O(batch · intermediate_size · state_size · seq²) — the
222 per-(channel, state) decay tensor — even for the summed default; use on
223 short sequences.
224 """
225 t = self._s6_terms(cache, layer_idx)
227 # M_coord[c, s, i, j] = C[i, s] · decay[c, s, i, j] · B[j, s]
228 C_col = t.C.permute(0, 2, 1)[:, None, :, :, None] # [batch, 1, state, i, 1]
229 B_col = t.B.permute(0, 2, 1)[:, None, :, None, :] # [batch, 1, state, 1, j]
230 M_coord = C_col * t.decay * B_col # [batch, channels, state, i, j]
232 if include_dt_scaling:
233 M_coord = M_coord * t.dt[:, :, None, None, :] # × dt[c, j]
235 if per_state_coord:
236 return M_coord
237 return M_coord.sum(dim=2) # [batch, channels, seq, seq]
239 def compute_ssm_state(
240 self,
241 cache: ActivationCache,
242 layer_idx: int,
243 time_step: Optional[int] = None,
244 ) -> torch.Tensor:
245 """Reconstruct Mamba-1's recurrent state ``S`` from cached hook values.
247 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]``
248 unrolls to::
250 S_t[c, s] = sum_{j<=t} decay[c, s, t, j] · dt_j[c] · x_j[c] · B_j[s]
252 with the same per-(channel, state) ``decay`` used by
253 ``compute_effective_attention``. ``x`` is the post-conv SiLU input
254 (``SiLU(conv1d.hook_out)``). Read-only: no ``forward()`` re-run. Verify
255 with ``y_t[c] = sum_s C_t[s]·S_t[c,s] + D[c]·x_t[c]``.
257 On padded batches the cached hooks are unmasked, so ``S`` is exact only at
258 non-pad positions; pad-position state is out of contract.
260 Args:
261 cache: ActivationCache from ``run_with_cache`` with this layer's
262 x_proj, dt_proj, and conv1d hooks.
263 layer_idx: Block index for this mixer.
264 time_step: If given, return only ``S`` at that position
265 (``[batch, channels, state]``); None returns every step.
267 Returns:
268 ``[batch, channels, seq, state]`` for all steps, or
269 ``[batch, channels, state]`` for a single ``time_step``.
271 Peak memory is O(batch · channels · state · seq²) — the per-(channel,
272 state) decay tensor, built in full by ``_s6_terms`` regardless of
273 ``time_step`` (which bounds only the returned tensor, not this peak); use
274 on short sequences.
275 """
276 conv_key = f"blocks.{layer_idx}.mixer.conv1d.hook_out"
277 if conv_key not in cache:
278 raise RuntimeError(
279 f"compute_ssm_state needs {conv_key!r} in cache. Run "
280 "`run_with_cache()` on the bridge before calling this method."
281 )
282 t = self._s6_terms(cache, layer_idx)
283 seq_len = t.dt.shape[-1]
284 oc: Any = self.original_component
285 # x = act(conv output), the SSM scan input; channel-first [batch, channels, seq].
286 x = oc.act(cache[conv_key].float()[..., :seq_len])
287 dtx = t.dt * x # dt_j[c]·x_j[c], [batch, channels, seq]
288 # S_t[c, s] = sum_{j<=t} decay[c, s, t, j] · dtx[c, j] · B_j[s]
289 if time_step is not None:
290 return torch.einsum("bcsj,bcj,bjs->bcs", t.decay[:, :, :, time_step, :], dtx, t.B)
291 return torch.einsum("bcsij,bcj,bjs->bcis", t.decay, dtx, t.B)
293 @staticmethod
294 def _apply_selective_param_norms(
295 oc: Any,
296 time_step: torch.Tensor,
297 B: torch.Tensor,
298 C: torch.Tensor,
299 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
300 """Apply Jamba-style RMSNorms on dt/B/C when the HF mixer exposes them.
302 Stock Mamba-1 has no ``dt_layernorm`` / ``b_layernorm`` / ``c_layernorm``;
303 Jamba's ``JambaMambaMixer`` does. Calling through the (possibly bridged)
304 attributes keeps their hooks firing on the eager-scan path.
305 """
306 dt_ln = getattr(oc, "dt_layernorm", None)
307 if dt_ln is not None:
308 time_step = dt_ln(time_step)
309 b_ln = getattr(oc, "b_layernorm", None)
310 if b_ln is not None:
311 B = b_ln(B)
312 c_ln = getattr(oc, "c_layernorm", None)
313 if c_ln is not None:
314 C = c_ln(C)
315 return time_step, B, C
317 def _s6_terms(self, cache: ActivationCache, layer_idx: int) -> _S6Terms:
318 """Reconstruct the shared S6 intermediates (dt, per-(channel,state) decay, B, C).
320 Reused by ``compute_effective_attention`` and ``compute_ssm_state``. Reads
321 B/C from ``x_proj.hook_out`` (shared across channels), or — when present —
322 from Jamba's post-norm ``b_layernorm`` / ``c_layernorm`` hooks. ``dt`` comes
323 from ``dt_proj.hook_out`` (softplus); that already sits after Jamba's
324 ``dt_layernorm``. A via ``__getattr__``. Dims come from the wrapped HF
325 mixer; cfg is the fallback (mirrors the Mamba-2 bridge).
326 """
327 if self.config is None: 327 ↛ 328line 327 didn't jump to line 328 because the condition on line 327 was never true
328 raise RuntimeError("SSMMixerBridge.config must be set")
330 x_proj_key = f"blocks.{layer_idx}.mixer.x_proj.hook_out"
331 dt_proj_key = f"blocks.{layer_idx}.mixer.dt_proj.hook_out"
332 for key in (x_proj_key, dt_proj_key):
333 if key not in cache:
334 raise RuntimeError(
335 f"S6 reconstruction needs {key!r} in cache. Run "
336 "`run_with_cache()` on the bridge before calling this."
337 )
339 cfg = self.config
340 oc = self.original_component
342 def _mamba_dim(module_attr: str, cfg_attr: str, default: Any) -> Any:
343 if oc is not None: 343 ↛ 347line 343 didn't jump to line 347 because the condition on line 343 was always true
344 val = getattr(oc, module_attr, None)
345 if val is not None: 345 ↛ 347line 345 didn't jump to line 347 because the condition on line 345 was always true
346 return val
347 return getattr(cfg, cfg_attr, default)
349 state_size = int(_mamba_dim("ssm_state_size", "state_size", 16))
350 dt_rank = int(_mamba_dim("time_step_rank", "time_step_rank", 0))
352 x_proj_out = cache[x_proj_key].float() # [batch, seq, dt_rank + 2*state]
353 dt_proj_out = cache[dt_proj_key].float() # [batch, seq, channels]
354 seq_len = dt_proj_out.shape[1]
356 # Prefer post-norm B/C when the adapter mapped Jamba's selective-param LNs.
357 b_ln_key = f"blocks.{layer_idx}.mixer.b_layernorm.hook_out"
358 c_ln_key = f"blocks.{layer_idx}.mixer.c_layernorm.hook_out"
359 if b_ln_key in cache and c_ln_key in cache:
360 B = cache[b_ln_key].float()
361 C = cache[c_ln_key].float()
362 else:
363 # Stock Mamba-1: B, C are the raw x_proj tails (shared across channels).
364 _time_step, B, C = x_proj_out.split([dt_rank, state_size, state_size], dim=-1)
366 # FalconMamba normalizes B/C after x_proj; dt_proj.hook_out is already
367 # post-RMS (HF applied it before dt_proj), so dt needs no correction.
368 rms_eps = getattr(oc, "rms_eps", None) if oc is not None else None
369 if rms_eps is not None:
370 B = _weightless_rms(B, rms_eps)
371 C = _weightless_rms(C, rms_eps)
373 dt = torch.nn.functional.softplus(dt_proj_out).transpose(1, 2) # [batch, channels, seq]
374 A = -torch.exp(self.A_log.float()) # [channels, state]
376 # decay[c, s, i, j] = exp(A[c,s]·(cumdt[c,i]-cumdt[c,j])) for i >= j, else 0.
377 cumdt = torch.cumsum(dt, dim=-1) # [batch, channels, seq]
378 dcs = cumdt[:, :, :, None] - cumdt[:, :, None, :] # [batch, channels, i, j]
379 decay_exp = (
380 A[None, :, :, None, None] * dcs[:, :, None, :, :]
381 ) # [batch, channels, state, i, j]
382 causal_mask = torch.tril(
383 torch.ones(seq_len, seq_len, dtype=torch.bool, device=decay_exp.device)
384 )
385 decay = torch.where(
386 causal_mask[None, None, None, :, :],
387 torch.exp(decay_exp),
388 torch.zeros_like(decay_exp),
389 )
390 return _S6Terms(dt=dt, decay=decay, B=B, C=C)