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)

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.

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

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)

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]
training: bool
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.