Add config for Gemma3-12B on v6e-256/512/1024 #2488
Merged
Conversation
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
NuojCheng
approved these changes
Oct 13, 2025
bvandermoon
reviewed
Oct 13, 2025
396c59e to
cbe1e3c
Compare
bvandermoon
approved these changes
Oct 13, 2025
cbe1e3c to
633f360
Compare
633f360 to
291ec23
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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
Add new configs for Gemma3-12B on v6e-256, 2x-v6e-256 (512), and 4x-v6e-256 (1024) (derived from b/429070413).
Why is this change being made?
This PR adds the missing Gemma3-12B model configs on different v6e configurations.
What is the problem being solved and any relevant context?
Users required to run Gemma3-12B on different v6e configurations. We found some model parameters that results in good MFU. We will create recipes to share the findings.
Why this is a good solution?
We thoroughly tested different configurations (v6e-256, 2x-v6e-256 (512), 4x-v6e-256 (1024)) until we found solutions with good MFU that make the customer happy (b/429070413). We made further improvements by using the relatively new feature of vocab tiling (#2242).
What would be some information about the specific implementation?
Using pure FSDP on 32K sequence length with vocab tiling, custom splash attention block sizes and increased xla_tpu_scoped_vmem_limit_kib.
What are the shortcomings of the solution and possible future improvements?
We could continue looking for better splash attention block sizes.
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/429070413
FIXES: b/451397849
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Tested the following configurations using MaxText's benchmark_runner / XPK and verifying the Cloud logs:
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.