Fix Gemma 4 gated MLP hook aliases - #1672
Closed
LarryHu0217 wants to merge 1 commit into
Closed
Conversation
Collaborator
|
Hi @LarryHu0217! Sorry I didn't see your message on the issue before you posted this. The issue is for |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Gemma 4 dense MLPs use the same gate/up/down projection structure as other gated MLP architectures, but the adapter mapped them as a bare GeneralizedComponent. This prevented the standard hook_pre, hook_pre_linear, and hook_post aliases from being registered.\n\nThis change uses the shared gated-MLP bridge for Gemma 4 and adds unit and integration coverage that verifies the aliases map to the expected projections and are available through run_with_cache across the PLE/KV-sharing, dense K=V, and MoE fixtures.\n\nFixes #1646\n\n## Type of change\n\n- [x] Bug fix (non-breaking change which fixes an issue)\n\n## Testing\n\n- Unit adapter test: 10 passed\n- Gemma 4 integration test: 15 passed\n- Full non-slow unit suite: 5266 passed, 56 skipped, 10 xfailed\n- Format check passed\n- mypy passed for 431 source files\n\n# Checklist\n\n- [x] The implementation is self-explanatory and does not require additional comments\n- [x] No documentation change is required for restoring the existing hook interface\n- [x] My changes generate no new warnings\n- [x] I have added tests that prove the fix is effective\n- [x] New and existing unit tests pass locally with my changes\n- [x] I have not rewritten tests relating to key interfaces