Coverage for transformer_lens/model_bridge/generalized_components/moe.py: 90%
153 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Mixture of Experts bridge component.
3This module contains the bridge component for Mixture of Experts layers.
4"""
6from __future__ import annotations
8from typing import Any, Dict, Mapping, Optional, Tuple
10import torch
12from transformer_lens.hook_points import HookPoint
13from transformer_lens.model_bridge.generalized_components.base import (
14 GeneralizedComponent,
15)
16from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
17from transformer_lens.model_bridge.generalized_components.mlp import (
18 normalize_mlp_weight,
19 weight_layout_in_out,
20)
23class MoEBridge(GeneralizedComponent):
24 """Bridge component for Mixture of Experts layers.
26 This component wraps a Mixture of Experts layer from a remote model and provides a consistent interface
27 for accessing its weights and performing MoE operations.
29 hook_router_scores fires only when the wrapped block returns a tuple
30 (gpt_oss, LLaDA2 remote); 5.13-native SparseMoeBlocks return a plain
31 tensor, so router observability comes from the ``gate`` submodule's
32 hook_out instead.
33 """
35 hook_aliases = {"hook_pre": "hook_in", "hook_post": "hook_out"}
37 # Deliberately NOT gate/in/out: ``gate`` is the ROUTER on sparse layers of
38 # the same model, so reusing it would flip the meaning of
39 # blocks.N.mlp.gate.hook_out per layer (#1645's own confusion). dense_gate
40 # is absent on ungated feed-forwards (Switch's wi/wo); the alias set
41 # follows the bound shape.
42 DENSE_SUBMODULE_KEYS = ("dense_in", "dense_out")
43 DENSE_GATE_KEY = "dense_gate"
44 _DENSE_HOOK_ALIASES = {
45 "hook_pre": "dense_in.hook_out",
46 "hook_post": "dense_out.hook_in",
47 }
48 _DENSE_GATED_HOOK_ALIASES = {
49 "hook_pre": "dense_gate.hook_out",
50 "hook_pre_linear": "dense_in.hook_out",
51 "hook_post": "dense_out.hook_in",
52 }
53 _DENSE_PROPERTY_ALIASES = {
54 "b_gate": "dense_gate.bias",
55 "b_in": "dense_in.bias",
56 "b_out": "dense_out.bias",
57 }
59 def __init__(
60 self,
61 name: str,
62 config: Optional[Any] = None,
63 submodules: Optional[Dict[str, GeneralizedComponent]] = {},
64 optional: bool = False,
65 sparse_required: Tuple[str, ...] = (),
66 ):
67 """Initialize the MoE bridge.
69 Args:
70 name: The name of the component in the model
71 config: Optional configuration (unused for MoEBridge)
72 submodules: Dictionary of GeneralizedComponent submodules to register
73 optional: If True, setup skips this subtree when absent (dense layers)
74 sparse_required: Submodule keys that must be declared ``optional``
75 (dense layers of an interleaved stack do not have them) but whose
76 absence on a SPARSE layer is an error rather than a silent skip.
77 Routers belong here: HF creates them unconditionally on sparse
78 blocks, so a skip means the attribute was renamed or moved, and
79 plain ``optional`` would drop their hooks without a word.
80 """
81 super().__init__(name, config, submodules=submodules, optional=optional)
82 self.hook_router_scores = HookPoint()
83 self._bound_dense = False
84 self._bound_dense_gate = False
85 # A misspelled key would silently disable the very guard that exists to
86 # stop silent degradation, so the opt-in is validated on construction.
87 unknown = set(sparse_required) - set(submodules or {})
88 if unknown:
89 raise ValueError(
90 f"{name}: sparse_required {sorted(unknown)} are not declared "
91 f"submodules (declared: {sorted(submodules or {})})"
92 )
93 self._sparse_required = sparse_required
95 def _binds_dense_projections(self, component: torch.nn.Module) -> bool:
96 """Whether this layer is the dense variant of an interleaved MoE stack.
98 Positive detection only: the adapter must have declared the dense
99 projections AND the wrapped module must actually expose them. Guessing
100 from the absence of ``experts`` would risk silently stripping hooks off
101 a sparse block whose experts are named differently.
102 """
103 declared = [self.submodules.get(key) for key in self.DENSE_SUBMODULE_KEYS]
104 if not all(declared):
105 return False
106 return all(
107 sub is not None and sub.name is not None and hasattr(component, sub.name)
108 for sub in declared
109 )
111 def _binds_dense_gate(self, component: torch.nn.Module) -> bool:
112 """Whether the bound dense MLP is gated (SiLU-gated) rather than plain.
114 A declared-but-unresolvable gate means the HF attribute was renamed,
115 not that the MLP is ungated — binding it as ungated would silently
116 alias hook_pre to the UP projection, so raise instead.
117 """
118 gate = self.submodules.get(self.DENSE_GATE_KEY)
119 if gate is None or gate.name is None:
120 return False # never declared: a genuinely ungated dense MLP
121 if hasattr(component, gate.name):
122 return True
123 raise ValueError(
124 f"{self.name}: dense layer wrapped {type(component).__name__} which has "
125 f"no {gate.name!r}, but this adapter declares {self.DENSE_GATE_KEY!r} — "
126 "so its dense MLPs are gated and the attribute was renamed or moved. "
127 "Binding it as ungated would alias hook_pre to the up projection. "
128 "Update the adapter's submodule name, or drop the declaration if this "
129 "architecture's dense layers really are ungated."
130 )
132 def set_original_component(self, component: torch.nn.Module) -> None:
133 """Bind the layer, adopting gated-MLP semantics on dense layers.
135 Dense layers of interleaved MoE stacks get neuron-basis hooks and
136 weight accessors instead of MoE boundary tensors under those names.
137 """
138 super().set_original_component(component)
139 is_dense = self._binds_dense_projections(component)
140 if is_dense == self._bound_dense:
141 # First sparse bind, or an idempotent rebind: nothing to morph.
142 if not is_dense: 142 ↛ 144line 142 didn't jump to line 144 because the condition on line 142 was always true
143 return
144 self._bound_dense = is_dense
145 if is_dense:
146 self._bound_dense_gate = self._binds_dense_gate(component)
147 self.hook_aliases = dict(
148 self._DENSE_GATED_HOOK_ALIASES
149 if self._bound_dense_gate
150 else self._DENSE_HOOK_ALIASES
151 )
152 dense_props = dict(self._DENSE_PROPERTY_ALIASES)
153 if not self._bound_dense_gate:
154 dense_props.pop("b_gate", None)
155 self.property_aliases = {**self.property_aliases, **dense_props}
156 # A dense layer has no router; leaving the hook in hook_dict would
157 # advertise an intervention point that can never fire.
158 if hasattr(self, "hook_router_scores"): 158 ↛ exitline 158 didn't return from function 'set_original_component' because the condition on line 158 was always true
159 self._hook_registry.pop("hook_router_scores", None)
160 del self.hook_router_scores
161 else:
162 # Symmetric restore so a rebinding harness cannot leave a chimera.
163 self.hook_aliases = dict(type(self).hook_aliases)
164 self.property_aliases = {
165 key: value
166 for key, value in self.property_aliases.items()
167 if key not in self._DENSE_PROPERTY_ALIASES
168 }
169 self._bound_dense_gate = False
170 if not hasattr(self, "hook_router_scores"): 170 ↛ exitline 170 didn't return from function 'set_original_component' because the condition on line 170 was always true
171 self.hook_router_scores = HookPoint()
173 def validate_after_setup(self, skipped_optional: list[str]) -> None:
174 """Fail loudly when a sparse layer is missing a submodule only dense
175 layers may lack (see ``sparse_required``).
177 Called by setup_submodules once every submodule has resolved — the
178 skipped set is not knowable at bind time.
179 """
180 if self._bound_dense or not self._sparse_required:
181 return
182 missing = [key for key in self._sparse_required if key in skipped_optional]
183 if missing:
184 component = type(self.original_component).__name__
185 raise ValueError(
186 f"{self.name}: sparse MoE layer wrapped {component} but required "
187 f"submodule(s) {missing} were absent. These are optional only so "
188 "dense layers of an interleaved stack can skip them; on a sparse "
189 "layer their absence means the HF attribute was renamed or moved. "
190 "Update the adapter's submodule name(s) rather than losing the hooks."
191 )
193 @property
194 def bound_dense(self) -> bool:
195 """Whether this layer bound the dense variant of an interleaved MoE stack.
197 Public so weight-collection helpers can find the projections under
198 ``DENSE_SUBMODULE_KEYS`` instead of the sparse ``gate`` (the router).
199 """
200 return self._bound_dense
202 def _dense_projection(self, key: str) -> Any:
203 """Return a bound dense projection.
205 Raises AttributeError on sparse layers (per-expert weights, no single
206 W_*) so ``hasattr`` stays False and weight collection skips them.
207 """
208 if not self._bound_dense:
209 raise AttributeError(f"{self.name}: {key} exists only on dense layers")
210 module = getattr(self, key, None)
211 if module is None: 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true
212 raise AttributeError(f"{self.name}: dense projection {key!r} is not bound")
213 return module
215 @property
216 def W_gate(self) -> torch.Tensor:
217 """Gated dense layer's gate weight in TL orientation [d_model, d_mlp]."""
218 if not self._bound_dense_gate:
219 raise AttributeError(f"{self.name}: this dense layer is ungated (no W_gate)")
220 module = self._dense_projection("dense_gate")
221 return normalize_mlp_weight(
222 module.weight, weight_layout_in_out(module), module, pattern="in"
223 )
225 @property
226 def W_in(self) -> torch.Tensor:
227 """Dense-layer input weight in TL orientation [d_model, d_mlp]."""
228 module = self._dense_projection("dense_in")
229 return normalize_mlp_weight(
230 module.weight, weight_layout_in_out(module), module, pattern="in"
231 )
233 @property
234 def W_out(self) -> torch.Tensor:
235 """Dense-layer output weight in TL orientation [d_mlp, d_model]."""
236 module = self._dense_projection("dense_out")
237 return normalize_mlp_weight(
238 module.weight, weight_layout_in_out(module), module, pattern="out"
239 )
241 def get_random_inputs(
242 self,
243 batch_size: int = 2,
244 seq_len: int = 8,
245 device: Optional[torch.device] = None,
246 dtype: Optional[torch.dtype] = None,
247 ) -> Dict[str, Any]:
248 """Generate random inputs for component testing.
250 Args:
251 batch_size: Batch size for generated inputs
252 seq_len: Sequence length for generated inputs
253 device: Device to place tensors on
254 dtype: Dtype for generated tensors (defaults to float32)
256 Returns:
257 Dictionary of input tensors matching the component's expected input signature
258 """
259 if device is None: 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true
260 device = torch.device("cpu")
261 if dtype is None: 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true
262 dtype = torch.float32
263 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 768
264 # Use positional args to avoid parameter name mismatches across MoE implementations
265 # (e.g., Mixtral uses "hidden_states", GraniteMoe uses "layer_input")
266 return {"args": (torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype),)}
268 def forward(self, *args: Any, **kwargs: Any) -> Any:
269 """Forward pass through the MoE bridge.
271 Args:
272 *args: Input arguments
273 **kwargs: Input keyword arguments
275 Returns:
276 Same return type as original component (tuple or tensor).
277 For MoE models that return (hidden_states, router_scores), preserves the tuple.
278 Router scores are also captured via hook for inspection.
279 """
280 if self.original_component is None: 280 ↛ 281line 280 didn't jump to line 281 because the condition on line 280 was never true
281 raise RuntimeError(
282 f"Original component not set for {self.name}. Call set_original_component() first."
283 )
284 if len(args) > 0:
285 hooked = self.hook_in(args[0])
286 args = (hooked,) + args[1:]
287 elif "hidden_states" in kwargs: 287 ↛ 290line 287 didn't jump to line 290 because the condition on line 287 was always true
288 hooked = self.hook_in(kwargs["hidden_states"])
289 kwargs = {**kwargs, "hidden_states": hooked}
290 output = self.original_component(*args, **kwargs)
291 if isinstance(output, tuple):
292 if not output:
293 raise TypeError(
294 f"{self.name}: expected a non-empty tuple whose first "
295 "element is a torch.Tensor from the wrapped MoE component, "
296 "got an empty tuple."
297 )
299 hidden_states = output[0]
300 if not isinstance(hidden_states, torch.Tensor):
301 raise TypeError(
302 f"{self.name}: expected the first tuple element from the "
303 f"wrapped MoE component to be a torch.Tensor, got "
304 f"{type(hidden_states).__name__}."
305 )
306 if len(output) > 1: 306 ↛ 319line 306 didn't jump to line 319 because the condition on line 306 was always true
307 router_scores = output[1]
308 # Some MoEs pack extras with the logits (LLaDA2 returns
309 # (router_logits, topk_idx)); hook the first tensor.
310 if isinstance(router_scores, tuple):
311 router_scores = next(
312 (t for t in router_scores if isinstance(t, torch.Tensor)), None
313 )
314 # The hook is removed on dense binds (no router exists there), so
315 # a dense layer whose wrapped module still returns a tuple must
316 # pass the extras through untouched rather than raise.
317 if isinstance(router_scores, torch.Tensor) and hasattr(self, "hook_router_scores"):
318 self.hook_router_scores(router_scores)
319 hidden_states = self.hook_out(hidden_states)
320 return (hidden_states,) + output[1:]
321 else:
322 hidden_states = self.hook_out(output)
323 return hidden_states
326class MoERouterBridge(LinearBridge):
327 """Bridge MoE router logits while preserving HF's tuple return.
329 5.13 TopKRouters return ``(router_logits, topk_weights, topk_indices)``;
330 hook_out fires on the logits (element ``logits_index`` — JetMoe puts them
331 last) and the tuple is re-packed so HF's unpacking is undisturbed.
332 """
334 def __init__(self, *args: Any, logits_index: int = 0, **kwargs: Any):
335 super().__init__(*args, **kwargs)
336 self.logits_index = logits_index
338 def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> Any:
339 if self.original_component is None: 339 ↛ 340line 339 didn't jump to line 340 because the condition on line 339 was never true
340 raise RuntimeError(
341 f"Original component not set for {self.name}. Call set_original_component() first."
342 )
343 input = self.hook_in(input)
344 output = self.original_component(input, *args, **kwargs)
345 if not isinstance(output, tuple) or len(output) == 0: 345 ↛ 346line 345 didn't jump to line 346 because the condition on line 345 was never true
346 return self.hook_out(output)
347 idx = self.logits_index % len(output)
348 router_logits = self.hook_out(output[idx])
349 return output[:idx] + (router_logits,) + output[idx + 1 :]
351 def set_processed_weights(
352 self, weights: Mapping[str, Optional[torch.Tensor]], verbose: bool = False
353 ) -> None:
354 """Copy router weights onto nested params by dotted path (JetMoe nests its
355 Linear at ``router.layer.weight``); router weights are never processed."""
356 if "weight" in weights:
357 super().set_processed_weights(weights, verbose=verbose)
358 return
359 if self.original_component is None: 359 ↛ 360line 359 didn't jump to line 360 because the condition on line 359 was never true
360 raise RuntimeError(f"Original component not set for {self.name}")
361 for key, tensor in weights.items():
362 if tensor is None: 362 ↛ 363line 362 didn't jump to line 363 because the condition on line 362 was never true
363 continue
364 target: Any = self.original_component
365 *path, leaf = key.split(".")
366 for part in path:
367 target = getattr(target, part)
368 param = getattr(target, leaf)
369 if param.shape != tensor.shape:
370 raise ValueError(
371 f"Router weight {key} shape {tuple(tensor.shape)} does not match "
372 f"parameter shape {tuple(param.shape)} on {self.name}"
373 )
374 with torch.no_grad():
375 param.copy_(tensor.to(dtype=param.dtype, device=param.device))