transformer_lens.model_bridge.generalized_components.moe module

Mixture of Experts bridge component.

This module contains the bridge component for Mixture of Experts layers.

class transformer_lens.model_bridge.generalized_components.moe.MoEBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = {}, optional: bool = False, sparse_required: Tuple[str, ...] = ())

Bases: GeneralizedComponent

Bridge component for Mixture of Experts layers.

This component wraps a Mixture of Experts layer from a remote model and provides a consistent interface for accessing its weights and performing MoE operations.

hook_router_scores fires only when the wrapped block returns a tuple (gpt_oss, LLaDA2 remote); 5.13-native SparseMoeBlocks return a plain tensor, so router observability comes from the gate submodule’s hook_out instead.

DENSE_GATE_KEY = 'dense_gate'
DENSE_SUBMODULE_KEYS = ('dense_in', 'dense_out')
property W_gate: Tensor

Gated dense layer’s gate weight in TL orientation [d_model, d_mlp].

property W_in: Tensor

Dense-layer input weight in TL orientation [d_model, d_mlp].

property W_out: Tensor

Dense-layer output weight in TL orientation [d_mlp, d_model].

__init__(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = {}, optional: bool = False, sparse_required: Tuple[str, ...] = ())

Initialize the MoE bridge.

Parameters:

  • name – The name of the component in the model

  • config – Optional configuration (unused for MoEBridge)

  • submodules – Dictionary of GeneralizedComponent submodules to register

  • optional – If True, setup skips this subtree when absent (dense layers)

  • sparse_required – Submodule keys that must be declared optional (dense layers of an interleaved stack do not have them) but whose absence on a SPARSE layer is an error rather than a silent skip. Routers belong here: HF creates them unconditionally on sparse blocks, so a skip means the attribute was renamed or moved, and plain optional would drop their hooks without a word.

property bound_dense: bool

Whether this layer bound the dense variant of an interleaved MoE stack.

Public so weight-collection helpers can find the projections under DENSE_SUBMODULE_KEYS instead of the sparse gate (the router).

forward(*args: Any, **kwargs: Any) → Any

Forward pass through the MoE bridge.

Parameters:

  • *args – Input arguments

  • **kwargs – Input keyword arguments

Returns:

Same return type as original component (tuple or tensor). For MoE models that return (hidden_states, router_scores), preserves the tuple. Router scores are also captured via hook for inspection.

get_random_inputs(batch_size: int = 2, seq_len: int = 8, device: device | None = None, dtype: dtype | None = None) → Dict[str, Any]

Generate random inputs for component testing.

Parameters:

  • batch_size – Batch size for generated inputs

  • seq_len – Sequence length for generated inputs

  • device – Device to place tensors on

  • dtype – Dtype for generated tensors (defaults to float32)

Returns:

Dictionary of input tensors matching the component’s expected input signature

hook_aliases: Dict[str, str | List[str]] = {'hook_post': 'hook_out', 'hook_pre': 'hook_in'}
real_components: Dict[str, tuple]
set_original_component(component: Module) → None

Bind the layer, adopting gated-MLP semantics on dense layers.

Dense layers of interleaved MoE stacks get neuron-basis hooks and weight accessors instead of MoE boundary tensors under those names.

training: bool
validate_after_setup(skipped_optional: list[str]) → None

Fail loudly when a sparse layer is missing a submodule only dense layers may lack (see sparse_required).

Called by setup_submodules once every submodule has resolved — the skipped set is not knowable at bind time.

class transformer_lens.model_bridge.generalized_components.moe.MoERouterBridge(*args: Any, logits_index: int = 0, weights_index: int | None = 1, indices_index: int | None = 2, **kwargs: Any)

Bases: LinearBridge

Bridge MoE router logits while preserving HF’s tuple return.

5.13 TopKRouters return (router_logits, topk_weights, topk_indices); hook_out fires on the logits (element logits_index — JetMoe puts them last) and the tuple is re-packed so HF’s unpacking is undisturbed.

hook_expert_weights / hook_expert_indices mirror the HookedTransformer MoE routing hooks. HF routers hand back top-k-shaped weights [tokens, top_k], so the weights are scattered to HT’s [tokens, num_experts] before firing and gathered back afterwards — an unedited round trip returns the values bit-for-bit. Any weight edit re-derives the top-k selection from the edited tensor, so boosting a suppressed expert re-routes the token (HT’s pre-top-k contract); unlike HT’s mixtral component, edited weights are used as-is with no renormalization after the hook.

forward(input: Tensor, *args: Any, **kwargs: Any) → Any

Forward pass through the linear layer with hooks.

Parameters:

  • input – Input tensor

  • *args – Additional positional arguments

  • **kwargs – Additional keyword arguments

Returns:

Output tensor after linear transformation

set_processed_weights(weights: Mapping[str, Tensor | None], verbose: bool = False) → None

Copy router weights onto nested params by dotted path (JetMoe nests its Linear at router.layer.weight); router weights are never processed.

exception transformer_lens.model_bridge.generalized_components.moe.UnfoldableMoEParameter

Bases: Exception

A parameter whose relationship to the MoE block’s input cannot be established.

transformer_lens.model_bridge.generalized_components.moe.fold_scale_into_moe_block(module: Module, scale: Tensor) → bool

Scale every parameter of a MoE block that reads the block’s input, in place.

Routed experts, shared experts and the router all read the preceding norm’s output; only the down-projections read the expert intermediate. Returns False without touching anything when any parameter’s role is unclear, because a fold that reaches some readers and not others changes what the model computes.

transformer_lens.model_bridge.generalized_components.moe.has_batched_experts(module: Module) → bool

Whether this MoE block stores its experts as batched 3-D Parameters.

Those parameters are not weight/bias leaves of a declared bridge submodule, so TransformerBridge.state_dict() drops them and no state-dict pass can reach them.

transformer_lens.model_bridge.generalized_components.moe.unwrap_bridge(module: Module) → Module

Descend through bridge wrappers to the module that owns the weights.