Coverage for transformer_lens/model_bridge/generalized_components/moe.py: 89%
276 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Mixture of Experts bridge component.
3This module contains the bridge component for Mixture of Experts layers.
4"""
6from __future__ import annotations
8import logging
9from typing import Any, Dict, List, Mapping, Optional, Tuple
11import torch
12from torch import nn
14from transformer_lens.hook_points import HookPoint
15from transformer_lens.model_bridge.generalized_components.base import (
16 GeneralizedComponent,
17)
18from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
19from transformer_lens.model_bridge.generalized_components.mlp import (
20 normalize_mlp_weight,
21 weight_layout_in_out,
22)
24# transformers stores a whole expert stack in one 3-D Parameter under a fixed
25# vocabulary of names. The role has to come from the name: in the untransposed
26# layout down_proj is [n_experts, d_model, d_mlp], so its d_model axis is
27# indistinguishable by shape from an input projection's.
28_BATCHED_INPUT_PROJECTIONS = frozenset({"gate_up_proj", "gate_proj", "up_proj"})
29_BATCHED_OUTPUT_PROJECTIONS = frozenset({"down_proj"})
30_BATCHED_PROJECTIONS = _BATCHED_INPUT_PROJECTIONS | _BATCHED_OUTPUT_PROJECTIONS
33class UnfoldableMoEParameter(Exception):
34 """A parameter whose relationship to the MoE block's input cannot be established."""
37def unwrap_bridge(module: nn.Module) -> nn.Module:
38 """Descend through bridge wrappers to the module that owns the weights."""
39 while True:
40 original = getattr(module, "original_component", None)
41 if not isinstance(original, nn.Module):
42 return module
43 module = original
46def has_batched_experts(module: nn.Module) -> bool:
47 """Whether this MoE block stores its experts as batched 3-D Parameters.
49 Those parameters are not ``weight``/``bias`` leaves of a declared bridge
50 submodule, so ``TransformerBridge.state_dict()`` drops them and no state-dict
51 pass can reach them.
52 """
53 return any(
54 parameter.ndim == 3 and name.rpartition(".")[2] in _BATCHED_PROJECTIONS
55 for name, parameter in unwrap_bridge(module).named_parameters()
56 )
59def _input_axis(
60 owner: nn.Module, leaf: str, parameter: torch.Tensor, d_model: int
61) -> Optional[int]:
62 """Axis the block's input flows into, or None when the parameter never reads it.
64 Raises UnfoldableMoEParameter when the role cannot be established, so the caller
65 can decline rather than guess — a fold that misses one reader is silently wrong.
66 """
67 if leaf == "bias" or leaf.endswith("_bias"):
68 return None # folding a norm's gain never touches a downstream bias
69 if parameter.ndim == 1:
70 raise UnfoldableMoEParameter(
71 f"{leaf}: 1-D weight inside the MoE block looks like a normalization gain, "
72 "which would re-normalize away the scale being folded"
73 )
74 if parameter.ndim == 2:
75 in_features = getattr(owner, "in_features", None)
76 if in_features is not None:
77 return -1 if in_features == d_model else None
78 # Routers hold a bare Parameter and consume it through F.linear, i.e. [out, in].
79 if parameter.shape[0] == d_model and parameter.shape[-1] != d_model: 79 ↛ 80line 79 didn't jump to line 80 because the condition on line 79 was never true
80 raise UnfoldableMoEParameter(
81 f"{leaf}: 2-D {tuple(parameter.shape)} on {type(owner).__name__} has "
82 "d_model on the output axis, so its orientation is not F.linear's"
83 )
84 if parameter.shape[-1] != d_model: 84 ↛ 85line 84 didn't jump to line 85 because the condition on line 84 was never true
85 return None
86 if parameter.shape[0] == d_model: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true
87 raise UnfoldableMoEParameter(f"{leaf}: square {tuple(parameter.shape)} is ambiguous")
88 return -1
89 if parameter.ndim == 3: 89 ↛ 111line 89 didn't jump to line 111 because the condition on line 89 was always true
90 if leaf in _BATCHED_OUTPUT_PROJECTIONS:
91 return None # reads the expert intermediate, not the block input
92 if leaf not in _BATCHED_INPUT_PROJECTIONS:
93 raise UnfoldableMoEParameter(f"{leaf}: unrecognized batched expert parameter")
94 transposed = getattr(owner, "is_transposed", None)
95 if transposed is not None:
96 axis = 1 if transposed else -1
97 if parameter.shape[axis] != d_model: 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true
98 raise UnfoldableMoEParameter(
99 f"{leaf}: {tuple(parameter.shape)} has no d_model on the axis "
100 f"is_transposed={transposed} implies"
101 )
102 return axis
103 # A few experts classes (Llama4) predate the layout flag; fall back to shape.
104 candidates = [axis for axis in (1, -1) if parameter.shape[axis] == d_model]
105 if len(candidates) != 1:
106 raise UnfoldableMoEParameter(
107 f"{leaf}: {type(owner).__name__} declares no is_transposed and "
108 f"{tuple(parameter.shape)} does not pin d_model to one axis"
109 )
110 return candidates[0]
111 raise UnfoldableMoEParameter(f"{leaf}: unexpected {parameter.ndim}-D parameter")
114def fold_scale_into_moe_block(module: nn.Module, scale: torch.Tensor) -> bool:
115 """Scale every parameter of a MoE block that reads the block's input, in place.
117 Routed experts, shared experts and the router all read the preceding norm's
118 output; only the down-projections read the expert intermediate. Returns False
119 without touching anything when any parameter's role is unclear, because a fold
120 that reaches some readers and not others changes what the model computes.
121 """
122 block = unwrap_bridge(module)
123 d_model = int(scale.shape[0])
124 plan: List[Tuple[torch.Tensor, int]] = []
125 try:
126 for name, parameter in block.named_parameters():
127 prefix, _, leaf = name.rpartition(".")
128 owner = block.get_submodule(prefix) if prefix else block
129 axis = _input_axis(owner, leaf, parameter, d_model)
130 if axis is not None:
131 plan.append((parameter, axis))
132 except UnfoldableMoEParameter as reason:
133 logging.warning("Not folding the layer norm into %s: %s", type(block).__name__, reason)
134 return False
135 if not plan: 135 ↛ 136line 135 didn't jump to line 136 because the condition on line 135 was never true
136 return False
137 with torch.no_grad():
138 for target, axis in plan:
139 shape = [1] * target.ndim
140 shape[axis] = d_model
141 target.mul_(scale.reshape(shape).to(dtype=target.dtype, device=target.device))
142 return True
145class MoEBridge(GeneralizedComponent):
146 """Bridge component for Mixture of Experts layers.
148 This component wraps a Mixture of Experts layer from a remote model and provides a consistent interface
149 for accessing its weights and performing MoE operations.
151 hook_router_scores fires only when the wrapped block returns a tuple
152 (gpt_oss, LLaDA2 remote); 5.13-native SparseMoeBlocks return a plain
153 tensor, so router observability comes from the ``gate`` submodule's
154 hook_out instead.
155 """
157 hook_aliases = {"hook_pre": "hook_in", "hook_post": "hook_out"}
159 # Deliberately NOT gate/in/out: ``gate`` is the ROUTER on sparse layers of
160 # the same model, so reusing it would flip the meaning of
161 # blocks.N.mlp.gate.hook_out per layer (#1645's own confusion). dense_gate
162 # is absent on ungated feed-forwards (Switch's wi/wo); the alias set
163 # follows the bound shape.
164 DENSE_SUBMODULE_KEYS = ("dense_in", "dense_out")
165 DENSE_GATE_KEY = "dense_gate"
166 _DENSE_HOOK_ALIASES = {
167 "hook_pre": "dense_in.hook_out",
168 "hook_post": "dense_out.hook_in",
169 }
170 _DENSE_GATED_HOOK_ALIASES = {
171 "hook_pre": "dense_gate.hook_out",
172 "hook_pre_linear": "dense_in.hook_out",
173 "hook_post": "dense_out.hook_in",
174 }
175 _DENSE_PROPERTY_ALIASES = {
176 "b_gate": "dense_gate.bias",
177 "b_in": "dense_in.bias",
178 "b_out": "dense_out.bias",
179 }
181 def __init__(
182 self,
183 name: str,
184 config: Optional[Any] = None,
185 submodules: Optional[Dict[str, GeneralizedComponent]] = {},
186 optional: bool = False,
187 sparse_required: Tuple[str, ...] = (),
188 ):
189 """Initialize the MoE bridge.
191 Args:
192 name: The name of the component in the model
193 config: Optional configuration (unused for MoEBridge)
194 submodules: Dictionary of GeneralizedComponent submodules to register
195 optional: If True, setup skips this subtree when absent (dense layers)
196 sparse_required: Submodule keys that must be declared ``optional``
197 (dense layers of an interleaved stack do not have them) but whose
198 absence on a SPARSE layer is an error rather than a silent skip.
199 Routers belong here: HF creates them unconditionally on sparse
200 blocks, so a skip means the attribute was renamed or moved, and
201 plain ``optional`` would drop their hooks without a word.
202 """
203 super().__init__(name, config, submodules=submodules, optional=optional)
204 self.hook_router_scores = HookPoint()
205 self._bound_dense = False
206 self._bound_dense_gate = False
207 # A misspelled key would silently disable the very guard that exists to
208 # stop silent degradation, so the opt-in is validated on construction.
209 unknown = set(sparse_required) - set(submodules or {})
210 if unknown:
211 raise ValueError(
212 f"{name}: sparse_required {sorted(unknown)} are not declared "
213 f"submodules (declared: {sorted(submodules or {})})"
214 )
215 self._sparse_required = sparse_required
216 # HookedTransformer exposes the routing observables on the MoE block
217 # itself; the bridge fires them on the router submodule, whose adapter
218 # key differs ("gate" on 5.13 SparseMoeBlocks, "router" on GPT-OSS).
219 # Alias so code migrated from HT finds them under the HT name.
220 self._router_hook_aliases = self._build_router_hook_aliases(submodules or {})
221 self.hook_aliases = {**self.hook_aliases, **self._router_hook_aliases}
223 @staticmethod
224 def _build_router_hook_aliases(
225 submodules: Mapping[str, GeneralizedComponent],
226 ) -> Dict[str, str]:
227 """Map HT's block-level routing hook names onto the router submodule."""
228 aliases: Dict[str, str] = {}
229 for key, component in submodules.items():
230 if not isinstance(component, MoERouterBridge):
231 continue
232 if component.weights_index is not None:
233 aliases["hook_expert_weights"] = f"{key}.hook_expert_weights"
234 if component.indices_index is not None:
235 aliases["hook_expert_indices"] = f"{key}.hook_expert_indices"
236 break
237 return aliases
239 def _binds_dense_projections(self, component: torch.nn.Module) -> bool:
240 """Whether this layer is the dense variant of an interleaved MoE stack.
242 Positive detection only: the adapter must have declared the dense
243 projections AND the wrapped module must actually expose them. Guessing
244 from the absence of ``experts`` would risk silently stripping hooks off
245 a sparse block whose experts are named differently.
246 """
247 declared = [self.submodules.get(key) for key in self.DENSE_SUBMODULE_KEYS]
248 if not all(declared):
249 return False
250 return all(
251 sub is not None and sub.name is not None and hasattr(component, sub.name)
252 for sub in declared
253 )
255 def _binds_dense_gate(self, component: torch.nn.Module) -> bool:
256 """Whether the bound dense MLP is gated (SiLU-gated) rather than plain.
258 A declared-but-unresolvable gate means the HF attribute was renamed,
259 not that the MLP is ungated — binding it as ungated would silently
260 alias hook_pre to the UP projection, so raise instead.
261 """
262 gate = self.submodules.get(self.DENSE_GATE_KEY)
263 if gate is None or gate.name is None:
264 return False # never declared: a genuinely ungated dense MLP
265 if hasattr(component, gate.name):
266 return True
267 raise ValueError(
268 f"{self.name}: dense layer wrapped {type(component).__name__} which has "
269 f"no {gate.name!r}, but this adapter declares {self.DENSE_GATE_KEY!r} — "
270 "so its dense MLPs are gated and the attribute was renamed or moved. "
271 "Binding it as ungated would alias hook_pre to the up projection. "
272 "Update the adapter's submodule name, or drop the declaration if this "
273 "architecture's dense layers really are ungated."
274 )
276 def set_original_component(self, component: torch.nn.Module) -> None:
277 """Bind the layer, adopting gated-MLP semantics on dense layers.
279 Dense layers of interleaved MoE stacks get neuron-basis hooks and
280 weight accessors instead of MoE boundary tensors under those names.
281 """
282 super().set_original_component(component)
283 is_dense = self._binds_dense_projections(component)
284 if is_dense == self._bound_dense:
285 # First sparse bind, or an idempotent rebind: nothing to morph.
286 if not is_dense: 286 ↛ 288line 286 didn't jump to line 288 because the condition on line 286 was always true
287 return
288 self._bound_dense = is_dense
289 if is_dense:
290 self._bound_dense_gate = self._binds_dense_gate(component)
291 self.hook_aliases = dict(
292 self._DENSE_GATED_HOOK_ALIASES
293 if self._bound_dense_gate
294 else self._DENSE_HOOK_ALIASES
295 )
296 dense_props = dict(self._DENSE_PROPERTY_ALIASES)
297 if not self._bound_dense_gate:
298 dense_props.pop("b_gate", None)
299 self.property_aliases = {**self.property_aliases, **dense_props}
300 # A dense layer has no router; leaving the hook in hook_dict would
301 # advertise an intervention point that can never fire.
302 if hasattr(self, "hook_router_scores"): 302 ↛ exitline 302 didn't return from function 'set_original_component' because the condition on line 302 was always true
303 self._hook_registry.pop("hook_router_scores", None)
304 del self.hook_router_scores
305 else:
306 # Symmetric restore so a rebinding harness cannot leave a chimera.
307 self.hook_aliases = {**type(self).hook_aliases, **self._router_hook_aliases}
308 self.property_aliases = {
309 key: value
310 for key, value in self.property_aliases.items()
311 if key not in self._DENSE_PROPERTY_ALIASES
312 }
313 self._bound_dense_gate = False
314 if not hasattr(self, "hook_router_scores"): 314 ↛ exitline 314 didn't return from function 'set_original_component' because the condition on line 314 was always true
315 self.hook_router_scores = HookPoint()
317 def validate_after_setup(self, skipped_optional: list[str]) -> None:
318 """Fail loudly when a sparse layer is missing a submodule only dense
319 layers may lack (see ``sparse_required``).
321 Called by setup_submodules once every submodule has resolved — the
322 skipped set is not knowable at bind time.
323 """
324 if self._bound_dense or not self._sparse_required:
325 return
326 missing = [key for key in self._sparse_required if key in skipped_optional]
327 if missing:
328 component = type(self.original_component).__name__
329 raise ValueError(
330 f"{self.name}: sparse MoE layer wrapped {component} but required "
331 f"submodule(s) {missing} were absent. These are optional only so "
332 "dense layers of an interleaved stack can skip them; on a sparse "
333 "layer their absence means the HF attribute was renamed or moved. "
334 "Update the adapter's submodule name(s) rather than losing the hooks."
335 )
337 @property
338 def bound_dense(self) -> bool:
339 """Whether this layer bound the dense variant of an interleaved MoE stack.
341 Public so weight-collection helpers can find the projections under
342 ``DENSE_SUBMODULE_KEYS`` instead of the sparse ``gate`` (the router).
343 """
344 return self._bound_dense
346 def _dense_projection(self, key: str) -> Any:
347 """Return a bound dense projection.
349 Raises AttributeError on sparse layers (per-expert weights, no single
350 W_*) so ``hasattr`` stays False and weight collection skips them.
351 """
352 if not self._bound_dense:
353 raise AttributeError(f"{self.name}: {key} exists only on dense layers")
354 module = getattr(self, key, None)
355 if module is None: 355 ↛ 356line 355 didn't jump to line 356 because the condition on line 355 was never true
356 raise AttributeError(f"{self.name}: dense projection {key!r} is not bound")
357 return module
359 @property
360 def W_gate(self) -> torch.Tensor:
361 """Gated dense layer's gate weight in TL orientation [d_model, d_mlp]."""
362 if not self._bound_dense_gate:
363 raise AttributeError(f"{self.name}: this dense layer is ungated (no W_gate)")
364 module = self._dense_projection("dense_gate")
365 return normalize_mlp_weight(
366 module.weight, weight_layout_in_out(module), module, pattern="in"
367 )
369 @property
370 def W_in(self) -> torch.Tensor:
371 """Dense-layer input weight in TL orientation [d_model, d_mlp]."""
372 module = self._dense_projection("dense_in")
373 return normalize_mlp_weight(
374 module.weight, weight_layout_in_out(module), module, pattern="in"
375 )
377 @property
378 def W_out(self) -> torch.Tensor:
379 """Dense-layer output weight in TL orientation [d_mlp, d_model]."""
380 module = self._dense_projection("dense_out")
381 return normalize_mlp_weight(
382 module.weight, weight_layout_in_out(module), module, pattern="out"
383 )
385 def get_random_inputs(
386 self,
387 batch_size: int = 2,
388 seq_len: int = 8,
389 device: Optional[torch.device] = None,
390 dtype: Optional[torch.dtype] = None,
391 ) -> Dict[str, Any]:
392 """Generate random inputs for component testing.
394 Args:
395 batch_size: Batch size for generated inputs
396 seq_len: Sequence length for generated inputs
397 device: Device to place tensors on
398 dtype: Dtype for generated tensors (defaults to float32)
400 Returns:
401 Dictionary of input tensors matching the component's expected input signature
402 """
403 if device is None: 403 ↛ 404line 403 didn't jump to line 404 because the condition on line 403 was never true
404 device = torch.device("cpu")
405 if dtype is None: 405 ↛ 406line 405 didn't jump to line 406 because the condition on line 405 was never true
406 dtype = torch.float32
407 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 768
408 # Use positional args to avoid parameter name mismatches across MoE implementations
409 # (e.g., Mixtral uses "hidden_states", GraniteMoe uses "layer_input")
410 return {"args": (torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype),)}
412 def forward(self, *args: Any, **kwargs: Any) -> Any:
413 """Forward pass through the MoE bridge.
415 Args:
416 *args: Input arguments
417 **kwargs: Input keyword arguments
419 Returns:
420 Same return type as original component (tuple or tensor).
421 For MoE models that return (hidden_states, router_scores), preserves the tuple.
422 Router scores are also captured via hook for inspection.
423 """
424 if self.original_component is None: 424 ↛ 425line 424 didn't jump to line 425 because the condition on line 424 was never true
425 raise RuntimeError(
426 f"Original component not set for {self.name}. Call set_original_component() first."
427 )
428 if len(args) > 0:
429 hooked = self.hook_in(args[0])
430 args = (hooked,) + args[1:]
431 elif "hidden_states" in kwargs: 431 ↛ 434line 431 didn't jump to line 434 because the condition on line 431 was always true
432 hooked = self.hook_in(kwargs["hidden_states"])
433 kwargs = {**kwargs, "hidden_states": hooked}
434 output = self.original_component(*args, **kwargs)
435 if isinstance(output, tuple):
436 if not output:
437 raise TypeError(
438 f"{self.name}: expected a non-empty tuple whose first "
439 "element is a torch.Tensor from the wrapped MoE component, "
440 "got an empty tuple."
441 )
443 hidden_states = output[0]
444 if not isinstance(hidden_states, torch.Tensor):
445 raise TypeError(
446 f"{self.name}: expected the first tuple element from the "
447 f"wrapped MoE component to be a torch.Tensor, got "
448 f"{type(hidden_states).__name__}."
449 )
450 if len(output) > 1: 450 ↛ 463line 450 didn't jump to line 463 because the condition on line 450 was always true
451 router_scores = output[1]
452 # Some MoEs pack extras with the logits (LLaDA2 returns
453 # (router_logits, topk_idx)); hook the first tensor.
454 if isinstance(router_scores, tuple):
455 router_scores = next(
456 (t for t in router_scores if isinstance(t, torch.Tensor)), None
457 )
458 # The hook is removed on dense binds (no router exists there), so
459 # a dense layer whose wrapped module still returns a tuple must
460 # pass the extras through untouched rather than raise.
461 if isinstance(router_scores, torch.Tensor) and hasattr(self, "hook_router_scores"):
462 self.hook_router_scores(router_scores)
463 hidden_states = self.hook_out(hidden_states)
464 return (hidden_states,) + output[1:]
465 else:
466 hidden_states = self.hook_out(output)
467 return hidden_states
470class MoERouterBridge(LinearBridge):
471 """Bridge MoE router logits while preserving HF's tuple return.
473 5.13 TopKRouters return ``(router_logits, topk_weights, topk_indices)``;
474 hook_out fires on the logits (element ``logits_index`` — JetMoe puts them
475 last) and the tuple is re-packed so HF's unpacking is undisturbed.
477 ``hook_expert_weights`` / ``hook_expert_indices`` mirror the HookedTransformer
478 MoE routing hooks. HF routers hand back top-k-shaped weights
479 ``[tokens, top_k]``, so the weights are scattered to HT's
480 ``[tokens, num_experts]`` before firing and gathered back afterwards — an
481 unedited round trip returns the values bit-for-bit. Any weight edit
482 re-derives the top-k selection from the edited tensor, so boosting a
483 suppressed expert re-routes the token (HT's pre-top-k contract); unlike
484 HT's mixtral component, edited weights are used as-is with no
485 renormalization after the hook.
486 """
488 def __init__(
489 self,
490 *args: Any,
491 logits_index: int = 0,
492 weights_index: Optional[int] = 1,
493 indices_index: Optional[int] = 2,
494 **kwargs: Any,
495 ):
496 super().__init__(*args, **kwargs)
497 self.logits_index = logits_index
498 self.weights_index = weights_index
499 self.indices_index = indices_index
500 # None means this router's tuple has no clean [tokens, top_k] pair
501 # (JetMoe returns a sorted-expert layout); registering the hook anyway
502 # would advertise an intervention point that can never fire.
503 if weights_index is not None:
504 self.hook_expert_weights = HookPoint()
505 if indices_index is not None:
506 self.hook_expert_indices = HookPoint()
508 def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> Any:
509 if self.original_component is None: 509 ↛ 510line 509 didn't jump to line 510 because the condition on line 509 was never true
510 raise RuntimeError(
511 f"Original component not set for {self.name}. Call set_original_component() first."
512 )
513 input = self.hook_in(input)
514 output = self.original_component(input, *args, **kwargs)
515 if not isinstance(output, tuple) or len(output) == 0: 515 ↛ 516line 515 didn't jump to line 516 because the condition on line 515 was never true
516 return self.hook_out(output)
517 parts = list(output)
518 count = len(parts)
519 logits_at = self.logits_index % count
520 parts[logits_at] = self.hook_out(parts[logits_at])
522 weights_at = None if self.weights_index is None else self.weights_index % count
523 indices_at = None if self.indices_index is None else self.indices_index % count
524 if weights_at is None and indices_at is None: 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true
525 return tuple(parts)
527 indices = None if indices_at is None else parts[indices_at]
528 expanded = None
529 if weights_at is not None: 529 ↛ 547line 529 didn't jump to line 547 because the condition on line 529 was always true
530 scattered = self._expand_expert_weights(parts[weights_at], indices, parts[logits_at])
531 expanded = self.hook_expert_weights(scattered)
532 if (
533 indices is not None
534 and expanded.shape != parts[weights_at].shape
535 and not torch.equal(expanded, scattered)
536 ):
537 # An edit outside the current top-k would otherwise be discarded
538 # by the gather below (those columns have no downstream reader in
539 # the [tokens, top_k] layout). Re-derive the selection from the
540 # edited tensor so boosting a suppressed expert re-routes, as it
541 # does on HookedTransformer's pre-top-k hook. The edited values
542 # are used as-is — no per-arch renormalization is re-applied.
543 _, new_indices = torch.topk(expanded, indices.shape[-1], dim=-1)
544 indices = new_indices.to(indices.dtype)
545 if indices_at is not None: 545 ↛ 547line 545 didn't jump to line 547 because the condition on line 545 was always true
546 parts[indices_at] = indices
547 if indices_at is not None: 547 ↛ 550line 547 didn't jump to line 550 because the condition on line 547 was always true
548 indices = self.hook_expert_indices(indices)
549 parts[indices_at] = indices
550 if weights_at is not None: 550 ↛ 554line 550 didn't jump to line 554 because the condition on line 550 was always true
551 # Gathered after the indices hook so re-routing picks up the weight
552 # sitting at the newly selected expert, as HookedTransformer does.
553 parts[weights_at] = self._collapse_expert_weights(expanded, indices, parts[weights_at])
554 return tuple(parts)
556 def _expand_expert_weights(
557 self,
558 weights: torch.Tensor,
559 indices: Optional[torch.Tensor],
560 logits: torch.Tensor,
561 ) -> torch.Tensor:
562 """Scatter top-k weights into HT's ``[tokens, num_experts]`` layout."""
563 if not self._is_top_k_shaped(weights, indices, logits): 563 ↛ 564line 563 didn't jump to line 564 because the condition on line 563 was never true
564 return weights
565 assert indices is not None
566 scattered = torch.zeros(
567 (*weights.shape[:-1], logits.shape[-1]),
568 dtype=weights.dtype,
569 device=weights.device,
570 )
571 scattered.scatter_(-1, indices.long(), weights)
572 return scattered
574 def _collapse_expert_weights(
575 self,
576 expanded: Optional[torch.Tensor],
577 indices: Optional[torch.Tensor],
578 original: torch.Tensor,
579 ) -> torch.Tensor:
580 """Gather the expanded weights back to the top-k layout HF expects."""
581 if expanded is None or indices is None or expanded.shape == original.shape: 581 ↛ 582line 581 didn't jump to line 582 because the condition on line 581 was never true
582 return expanded if expanded is not None else original
583 return expanded.gather(-1, indices.long())
585 @staticmethod
586 def _is_top_k_shaped(
587 weights: torch.Tensor,
588 indices: Optional[torch.Tensor],
589 logits: torch.Tensor,
590 ) -> bool:
591 """Whether the weights are the top-k slice rather than full expert width."""
592 return (
593 indices is not None
594 and isinstance(weights, torch.Tensor)
595 and isinstance(logits, torch.Tensor)
596 and weights.shape == indices.shape
597 and weights.shape[-1] != logits.shape[-1]
598 )
600 def set_processed_weights(
601 self, weights: Mapping[str, Optional[torch.Tensor]], verbose: bool = False
602 ) -> None:
603 """Copy router weights onto nested params by dotted path (JetMoe nests its
604 Linear at ``router.layer.weight``); router weights are never processed."""
605 if "weight" in weights:
606 super().set_processed_weights(weights, verbose=verbose)
607 return
608 if self.original_component is None: 608 ↛ 609line 608 didn't jump to line 609 because the condition on line 608 was never true
609 raise RuntimeError(f"Original component not set for {self.name}")
610 for key, tensor in weights.items():
611 if tensor is None: 611 ↛ 612line 611 didn't jump to line 612 because the condition on line 611 was never true
612 continue
613 target: Any = self.original_component
614 *path, leaf = key.split(".")
615 for part in path:
616 target = getattr(target, part)
617 param = getattr(target, leaf)
618 if param.shape != tensor.shape:
619 raise ValueError(
620 f"Router weight {key} shape {tuple(tensor.shape)} does not match "
621 f"parameter shape {tuple(param.shape)} on {self.name}"
622 )
623 with torch.no_grad():
624 param.copy_(tensor.to(dtype=param.dtype, device=param.device))