transformer_lens.model_bridge.supported_architectures.laguna module

Laguna architecture adapter.

poolside’s Laguna coding models (LagunaForCausalLM, remote code): Llama-shaped causal decoders with two first-of-kind mechanisms — heterogeneous per-layer attention head counts (num_attention_heads_per_layer) and per-head softplus output gating (g_proj) — over a batched-expert MoE (FlexOlmo-style 3D expert parameters, top-k router, always-on shared experts) with per-layer dense/sparse selection via mlp_layer_types.

Attention delegates to HF: the bridge reimplementation assumes one uniform head count, and the softplus gate has no reconstruction. The per-layer head heterogeneity also rules out uniform Q/K/V reshape conversions, so no HookedTransformer-format conversions ship and LN folding is disabled.

class transformer_lens.model_bridge.supported_architectures.laguna.LagunaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for LagunaForCausalLM models.

__init__(cfg: Any) → None

Initialize the Laguna architecture adapter.

component_mapping: ComponentMapping | None

prepare_loading(model_name: str, model_kwargs: dict) → None

User-register Laguna’s native conversion mapping so the per-expert->batched expert merge runs under remote code (transformers skips it for custom-code modules, leaving the batched experts at random init).

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

supports_fold_ln = False

uses_split_attention: bool

weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None