Skip to content

Fix Gemma 4 MLP hook aliases - #1650

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
koriyoshi2041:fix/gemma4-mlp-hooks
Aug 12, 2026
Merged

Fix Gemma 4 MLP hook aliases#1650
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
koriyoshi2041:fix/gemma4-mlp-hooks

Conversation

@koriyoshi2041

Copy link
Copy Markdown
Contributor

Description

Fixes #1646.

Problem

The Gemma 4 adapter decomposes its dense MLP into gate_proj, up_proj, and down_proj, but wraps them in a plain GeneralizedComponent. The standard gated-MLP aliases (hook_pre, hook_pre_linear, and hook_post) are therefore missing even though the underlying module has the same projection structure as other supported gated MLPs.

Fix

Use the adapter's existing _gated_mlp() helper for Gemma 4. This preserves the same projection paths while installing the standard hook aliases. The unit test now checks both the bridge type and the exact alias mapping.

Test

  • uv run pytest tests/unit/model_bridge/supported_architectures/test_gemma4_adapter.py -q (10 passed)
  • uv run pytest tests/integration/model_bridge/test_gemma4_bridge.py -q (12 passed across PLE/KV-shared, dense K=V, and MoE fixtures)
  • uv run mypy transformer_lens/model_bridge/supported_architectures/gemma4.py
  • make check-format
  • git diff --check

Risk

Low. The component retains the same mlp path and gate/up/down projection names. Gemma 4's delegated blocks still execute the Hugging Face forward path, and the three existing architecture fixtures retain logit parity. The change only adds the gated-MLP compatibility interface around the already-decomposed projections.

Type of change

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

Checklist

  • My changes generate no new warnings
  • I have added tests that prove my fix is effective
  • New and existing focused tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4

Copy link
Copy Markdown
Collaborator

Looks great, thanks for putting this together @koriyoshi2041!

@jlarson4
jlarson4 merged commit 532ee1b into TransformerLensOrg:dev Aug 12, 2026
25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug Report] Gemma 4 MLPs omit hook_pre, hook_pre_linear, and hook_post

2 participants