Repository navigation
Fix Gemma sliding_window being halved on every config save/reload - #47940
Merged
Rocketknight1 merged 1 commit intoAug 13, 2026
Merged
Rocketknight1 merged 1 commit into
Rocketknight1 merged 1 commit into
Conversation
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: diffusion_gemma, gemma3, gemma3n, gemma4, gemma4_unified, t5gemma2 |
Contributor
CI recapDashboard: View test results in Grafana |
Rocketknight1
approved these changes
Aug 13, 2026
Rocketknight1
left a comment
Member
There was a problem hiding this comment.
Verified the bug and the fix looks clean, so I'm happy to approve! This only affects embeddinggemma checkpoints, but it is a real bug and so worth fixing.
Rocketknight1
enabled auto-merge
August 13, 2026 13:27
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
3 of 6 tasks
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.
What does this PR do?
Gemma3TextConfig.__post_init__converts the declared sliding window into the exclusivehalf-window the attention kernels use, and writes the result back into the field it read from:
save_pretrainedthen persists the converted value, and the next load converts it again, so thewindow halves on every save/load round trip. No model download is needed to reproduce it:
Repeating the cycle gives
512 -> 257 -> 129 -> 65 -> 33 -> 17 -> 9 -> 5 -> 3 -> 2, and 2 is afixed point, so the window keeps shrinking until attention is effectively local to one token.
google/embeddinggemma-300mships"sliding_window": 512in itsconfig.json, so the valuesave_pretrainedwrites back is not the convention the checkpoint itself uses.Where it comes from
Introduced in #40700 (
948bc0fa). Before it, #40694 divided in the mask overlay(
abs(q_idx - kv_idx) <= sliding_window // 2) and left the config untouched. #40700 moved thedivision to config init so that flash-attention — which reads
config.sliding_windowfor itswindow_sizeand never builds a mask — agrees with eager/sdpa/flex. That consolidation iscorrect and this PR keeps it; the only problem is that the converted value is also the one that
gets serialized.
Impact
Every
save_pretrained/from_pretrainedcycle on a bidirectional Gemma config is affected:fine-tuning checkpoints, PEFT merges, quantization exports, and re-uploads to the Hub, which then
carry the halved window permanently.
Measured on a
google/embeddinggemma-300mcheckpoint round-tripped throughsave_pretrained,with its weights verified bit-identical to the source (314/314 tensors), evaluated with
sentence-transformerson NanoBEIR:sliding_windowpatched back to 512The weights are identical in all three rows, so the 0.87 pp is the config scalar alone; patching
only that field reproduces the baseline exactly to 16 significant digits on both datasets.
Fix
to_dict()serializes the value__post_init__converted from, so a saved config reloads tothe in-memory window instead of being halved a second time.
This mirrors
ModernBertConfig, which keeps the declared window in the serialized field andexposes the converted one as a
sliding_windowproperty.Gemma cannot use that split because its checkpoints already serialize the key as
sliding_window, so the conversion is undone at serialization time instead.(sliding_window - 1) * 2is an exact pre-image ofv // 2 + 1, since((2c - 2) // 2) + 1 == cfor every integer
c. For an even declared window — every released Gemma checkpoint — itrecovers the original value exactly (
512 -> 257 -> 512). An odd declared window comes back onelower, which reloads to the same in-memory value.
In-memory behaviour is unchanged: masking, attention and cache still see the converted window.
Configs that do not convert (
use_bidirectional_attentionfalse /None/"vision") serializebyte-for-byte as before.
The guard uses
getattrbecauseGemma3nTextConfig,T5Gemma2TextConfigandT5Gemma2DecoderConfiginheritGemma3TextConfigthrough modular while overriding__post_init__to drop the conversion, and do not define the flag at all — the same reasonGemma4TextConfig.__post_init__already readsgetattr(self, "attention_k_eq_v", True). Theirregenerated files carry the method as a no-op.
Compatibility
512) where it previously stored theconverted one (
257). Loading it throughtransformersyields257either way.257as a declared windowand
257as a converted one are indistinguishable. They still load as129, but they nolonger decay further on subsequent saves.
Happy to add the 🚨 prefix to the title if you consider the changed serialized value breaking.
Alternatives considered
divergence that [
Gemma Embedding] Fix SWA #40700 fixed.__post_init__idempotent needs extra state, and private attributes are serialized too(
_sliding_window_patternends up inconfig.json), so it would add a key to every file.Happy to move this into
PreTrainedConfiginstead if you would prefer a general mechanism for__post_init__normalisations that are not idempotent.Tests
test_bidirectional_sliding_window_survives_save_and_reloadin the Gemma3 and Gemma4 text modeltests, one per hand-written config source. Both fail on
mainwithAssertionError: 129 != 257.Gemma4UnifiedTextConfigandDiffusionGemmaTextConfigcarry the same__post_init__linethrough modular inheritance from
Gemma4TextConfig; their configuration files were regenerated.All seven config classes that end up with the new method — the four that convert and the three
that inherit it inert — were checked for save/reload stability over three consecutive round trips.
Before submitting
Who can review?
@vasqu (author of #40700) and @Cyrilvallez.