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, **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.

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.