Skip to content

Fix Gemma sliding_window being halved on every config save/reload - #47940

Merged
Rocketknight1 merged 1 commit into
huggingface:mainfrom
Bluear7878:fix-gemma-sliding-window-roundtrip
Aug 13, 2026
Merged

Rocketknight1 merged 1 commit into
huggingface:mainfrom
Bluear7878:fix-gemma-sliding-window-roundtrip

Conversation

@Bluear7878

@Bluear7878 Bluear7878 commented Aug 13, 2026 •

Copy link
Copy Markdown
Contributor

CI

What does this PR do?

Gemma3TextConfig.__post_init__ converts the declared sliding window into the exclusive
half-window the attention kernels use, and writes the result back into the field it read from:

if self.use_bidirectional_attention:
    self.sliding_window = (self.sliding_window // 2) + 1  # due to fa we set exclusive bounds

save_pretrained then persists the converted value, and the next load converts it again, so the
window halves on every save/load round trip. No model download is needed to reproduce it:

from transformers import Gemma3TextConfig

config = Gemma3TextConfig(sliding_window=512, use_bidirectional_attention=True)
print(config.sliding_window)                                          # 257
config.save_pretrained("/tmp/cfg")                                    # writes "sliding_window": 257
print(Gemma3TextConfig.from_pretrained("/tmp/cfg").sliding_window)    # 129

Repeating the cycle gives 512 -> 257 -> 129 -> 65 -> 33 -> 17 -> 9 -> 5 -> 3 -> 2, and 2 is a
fixed point, so the window keeps shrinking until attention is effectively local to one token.

google/embeddinggemma-300m ships "sliding_window": 512 in its config.json, so the value
save_pretrained writes 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 the
division to config init so that flash-attention — which reads config.sliding_window for its
window_size and never builds a mask — agrees with eager/sdpa/flex. That consolidation is
correct 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_pretrained cycle 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-300m checkpoint round-tripped through save_pretrained,
with its weights verified bit-identical to the source (314/314 tensors), evaluated with
sentence-transformers on NanoBEIR:

mean nDCG@10 nfcorpus scifact
source checkpoint 0.6271558738396297 0.3929820955134049 0.8613296521658544
after one round trip 0.6184786921356261 — —
round trip, sliding_window patched back to 512 0.6271558738396297 0.3929820955134049 0.8613296521658544

The 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 to
the in-memory window instead of being halved a second time.

This mirrors ModernBertConfig, which keeps the declared window in the serialized field and
exposes the converted one as a
sliding_window property.
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) * 2 is an exact pre-image of v // 2 + 1, since ((2c - 2) // 2) + 1 == c
for every integer c. For an even declared window — every released Gemma checkpoint — it
recovers the original value exactly (512 -> 257 -> 512). An odd declared window comes back one
lower, 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_attention false / None / "vision") serialize
byte-for-byte as before.

The guard uses getattr because Gemma3nTextConfig, T5Gemma2TextConfig and
T5Gemma2DecoderConfig inherit Gemma3TextConfig through modular while overriding
__post_init__ to drop the conversion, and do not define the flag at all — the same reason
Gemma4TextConfig.__post_init__ already reads getattr(self, "attention_k_eq_v", True). Their
regenerated files carry the method as a no-op.

Compatibility

  • A newly saved config stores the declared window (512) where it previously stored the
    converted one (257). Loading it through transformers yields 257 either way.
  • Configs already saved with the converted value cannot be repaired: 257 as a declared window
    and 257 as a converted one are indistinguishable. They still load as 129, but they no
    longer decay further on subsequent saves.

Happy to add the 🚨 prefix to the title if you consider the changed serialized value breaking.

Alternatives considered

  • Moving the division back to the mask overlay would reintroduce the flash-attention / mask
    divergence that [Gemma Embedding] Fix SWA #40700 fixed.
  • Making __post_init__ idempotent needs extra state, and private attributes are serialized too
    (_sliding_window_pattern ends up in config.json), so it would add a key to every file.

Happy to move this into PreTrainedConfig instead if you would prefer a general mechanism for
__post_init__ normalisations that are not idempotent.

Tests

test_bidirectional_sliding_window_survives_save_and_reload in the Gemma3 and Gemma4 text model
tests, one per hand-written config source. Both fail on main with AssertionError: 129 != 257.

Gemma4UnifiedTextConfig and DiffusionGemmaTextConfig carry the same __post_init__ line
through 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

  • This PR fixes a bug.
  • Did you read the contributor guideline?
  • Did you write any new necessary tests? Yes, one per affected config source.
  • The generated configuration files were regenerated from their modular files.

Who can review?

@vasqu (author of #40700) and @Cyrilvallez.

@github-actions

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: diffusion_gemma, gemma3, gemma3n, gemma4, gemma4_unified, t5gemma2

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 31671521932
Result: success | Grafana metrics are not available yet.

@Rocketknight1 Rocketknight1 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Rocketknight1 added this pull request to the merge queue Aug 13, 2026
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

Merged via the queue into huggingface:main with commit f4062e2 Aug 13, 2026
63 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.

3 participants