Hook extensions for specialized bridges

[jlarson4]

Description

Each adapters touched in this PR mapped a component to a bridge class that cannot express it, so hooks users expect either did not exist or resolved to something that never fired. None of it changes numerics, the swaps keep the math with HF and only restores observability.

• gemma3n and gemma4 attention were bare GeneralizedComponents whose entire surface is hook_in/hook_out. No hook_pattern, no hook_attn_scores, no hook_result; hook_q/k/v/z and W_Q..b_O did not resolve; and q/k/v fired flat [batch, seq, d_model] rather than the per-head shape every gemma1/2/3 sibling gives. Now usesAttentionBridge(maintain_native_attention=True), the same delegation 20+ adapters use, so HF still runs the attention. The optional k/v/k_norm/v_norm flags are preserved: KV-shared layers (E2B/E4B) and K==V layers legitimately lack those modules, and the alias set prunes itself accordingly. gemma3n's MLP takes the one-line self._gated_mlp() that gemma4 received in Fix Gemma 4 MLP hook aliases #1650.
• Nine encoder-decoder and decoder adapters mapped the MLP as SymbolicBridge – opt, xglm, and the seven that inherit BART's _mlp() (bart, marian, pegasus, blenderbot, mbart, m2m100, led). That class exposes no hook_pre/hook_post, and on OPT/XGLM the block's hook_mlp_out pointed at a HookPoint that never fired. MLPBridge(name=None) fixes both, copying BERT, which maps the containerless fc1/fc2 layout correctly. Safe because component setup keys promotion on name is None rather than on the bridge type, so the blocks.{i}.mlp.in/out weight paths are unchanged and MLPBridge.forward isn't invoked. hook_mlp_in is deliberately not overridden on OPT/XGLM: they are pre-norm, so its HookedTransformer meaning is the pre-ln2 residual the block already supplies. BART overrides it legitimately because it is post-LN.
• Two fused gate+up projections were served whole. – OpenELM's proj_1 and GraniteMoeHybrid's shared_mlp.input_linear each concatenate gate and up, so under a plain MLPBridge hook_pre was the 2*d_mlp pre-GLU tensor and hook_pre_linear did not exist, the neuron basis was unreachable. Both now use JointGateUpMLPBridge with an architecture-specific splitter, guarded by require_readable_weight like Phi-3's. The split direction is pinned by value in both tests, with a negative control, because shapes cannot catch a gate/up swap. Both halves are [d_mlp, d_model].
• OpenELM carried three dead aliases. Its qkv_proj is fused, so hook_q/hook_k/hook_v could never resolve. They are stripped the same way raven strips its Wqkv equivalents. hook_z stays, since it resolves through out_proj. OpenELM is correspondingly removed from _KNOWN_DEAD_ALIASES — that marker is xfail(strict=True), so leaving it would XPASS and fail the suite.
• granite_moe and granite_moe_hybrid had no router hook at all – their sparse blocks declared no submodules, so the routing logits every MoE analysis needs were unreachable. Both now map MoERouterBridge(name="router", logits_index=-1). The index is load-bearing and measured: HF returns (top_k_index, top_k_weights, router_logits), so the default index 0 hooks an int64 index tensor rather than the logits.
• RWKV-7 is deliberately left as a delegated component. Swapping its channel-mix to MLPBridge for parity with RWKV-4 looked like it would be correct, but MLPBridge pre-fires the wrapped module's hook_in, which would suppress the in projection's own hook_in and on RWKV-7 that hook carries the post-token-shift key input. The swap would have traded a correct hook for two alias names. Instead the two aliases are declared directly on the delegated component, so hook_pre/hook_post exist with nothing moved.
• Six pre-existing tests asserted the defects (isinstance(mlp, SymbolicBridge), mlp.submodules["in"].name == "proj_1", shared_mlp.submodules == {"in", "out"}). They are rewritten to assert the hooks a user actually gets rather than the class.
• OpenELM computed ReLU instead of SiLU on every load – the reconstructed FFN's activation fell back to cfg.act_fn's "relu" default because the module capture probed activation_fn/act_fn (OpenELM stores self.act) and the config conversion probed four names that don't include OpenELM's activation_fn_name. The probe list now also covers GraniteMoeHybrid's self.activation.
• gemma3n compat mode skipped activation sparsity – process_weights flipped the MLP onto the functional branch, which computes plain act(gate)*up without HF's _gaussian_topk. The gemma3n MLP now keeps delegating under processed weights, which is correct because distribution writes the processed tensors into the wrapped Linears, and the hooks the swap wanted still fire from inside HF's forward.
• OPT's newly-live MLP hooks fired flattened [batch*seq, d] – OPTDecoderLayer reshapes to 2D around ln2/fc1/fc2. A stamped unflatten conversion now presents [batch, seq, d] to hooks and reverts edits back to 2D, verified read side (shapes) and write side (a position-indexed edit lands positionally; the test fake mirrors modeling_opt's 2D residual so a broken revert crashes it as it would real OPT). Inert when no hooks are attached.
• Weight accessors refuse heterogeneous-geometry layers –gemma4's per-layer KV geometry divides evenly under the majority head count, so W_K/W_V would silently mis-factorize on minority layers; a width check now raises naming the per-layer geometry. Zero false positives across the adapter suite.
• MLPBridge(name=None) direct calls raise – component setup binds the parent decoder layer as the containerless MLP's original_component, so block.mlp(x) silently executed the whole layer (attention included) with a plausible output shape; the old SymbolicBridge raised here, and now MLPBridge does too.
• Delegated-attention comments corrected (under maintain_native_attention, hook_pattern and hook_attn_scores both fire HF's post-softmax weights and writes to them don't reach the output. The shared test fake now uses real o_proj widths so hook_z's per-head shape is actually asserted; class-constant alias restatements narrowed to isinstance + the repo-wide resolution suite.
• OpenELM ffn_with_glu=False (unreachable on published checkpoints) is refused at boot instead of crashing in a mat-mul at first forward.
• Consolidation: one parametrized, guarded, bias-aware default gate/up splitter (fused_attr=, skip_init so boot no longer burns kaiming inits or RNG draws) replaces the two new per-adapter copies; fused_qkv=True on AttentionBridge replaces the raven/OpenELM hand-rolled alias strips; _build_moe_bridge() on the granite base replaces the duplicated router mapping.

OpenELM re-verification surfaced three more pre-existing breaks, each fixed rather than skipped:

• OpenELM was entirely unloadable on transformers 5.13 – nothing to do with this branch. 5.x replaces config __init__ with dataclass machinery that delivers every non-base field via __post_init__(**extras); apple's frozen 4.x remote code defines it argless, so AutoConfig.from_pretrained itself raised. New autoconfig_with_remote_post_init_compat (used by boot and verify_models' estimate) retries once after wrapping the remote class's __post_init__ to run the base handler first — base-first is load-bearing, since the class's own fields only get set by the base handler under 5.x. Generic: fires only on that exact TypeError, for any 4.x-era remote-code repo.
• OpenELM generation died on the first step – apple's forward calls DynamicCache.from_legacy_cache/.to_legacy_cache, both removed in 5.x. prepare_loading now swaps the module-level name for a subclass restoring them (5.x's ddp_cache_data init is the legacy tuple format).
• The centering benchmark tripped over per-layer head counts – OpenELM varies num_query_heads per layer, so the new width guard (correctly) refuses the cfg-scalar factorization. The centering statistic is head-agnostic, so the benchmark falls back to the raw 2D projection; the probe also had to stop using hasattr(attn, "W_O"), which invokes the property and only swallows AttributeError.

All seven apple/OpenELM-* checkpoints (270M/270M-I/450M-I/1.1B/1.1B-I/3B/3B-I) verified with the SiLU fix in place, replacing the invalidated February entries.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility