Guard RoPE scaling against the transformers v5 buffer blank; honor extended RoPE factor - #6925
Conversation
…tended factor Add a family-agnostic guard that builds each rotary from a scaled config, blanks its non-persistent buffers (what transformers v5 does on load), runs loader._fix_rope_inv_freq, and asserts every buffer is restored to its scaled value (llama3 and longrope). This catches the whole bug class, not just the one call site, and is validated to fail on the pre-fix repair. Also make LlamaExtendedRotaryEmbedding read the llama3 factor from the config instead of hardcoding 8 (wrong for Llama-3.2, factor 32), falling back to the Llama-3.1 defaults when built without a config.
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
for more information, see https://pre-commit.ci
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4fb24d3c9c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| rope_scaling = _rope_scaling_as_dict( | ||
| getattr(getattr(self, "_unsloth_rope_config", None), "rope_scaling", None) or {} | ||
| ) | ||
| scale_factor = rope_scaling.get("factor", 8) |
There was a problem hiding this comment.
When this subclass is instantiated through the patched attention path, _utils.patch_llama_rope_scaling still constructs LlamaExtendedRotaryEmbedding(dim=..., max_position_embeddings=..., base=...) without passing config, so _unsloth_rope_config is None here and these new reads always fall back to the Llama-3.1 defaults. In that pre_patch path, Llama-3.2 configs with factor=32 will still use factor 8 and produce incorrect long-context RoPE; the config or rope_scaling values need to be wired into the subclass construction.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. The extended rotary read the config but the per-layer codegen never passed one. Fixed in patch_llama_rope_scaling (unsloth/models/_utils.py): the llama3 branch now builds the extended rotary with config=self.config, so it resolves the real factor (32 for Llama-3.2) instead of falling back to 8. self.config is already referenced in that template so nothing new is required at that call site.
| return rot, config | ||
|
|
||
|
|
||
| @requires_cuda |
There was a problem hiding this comment.
This test runs on any CUDA install, but loader._fix_rope_inv_freq immediately returns when _NEEDS_ROPE_FIX is false, i.e. with supported transformers 4.x installs. In that environment the test blanks the buffers and then calls a deliberate no-op, so the restoration assertions fail even though the v5 repair path should not be active; gate the test on loader._NEEDS_ROPE_FIX or force the flag inside the test.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Agreed. On transformers 4.x _fix_rope_inv_freq returns immediately (_NEEDS_ROPE_FIX is False), so the blanked buffers are never restored and the round-trip cannot pass. Added an early pytest.skip on not loader._NEEDS_ROPE_FIX so the test only runs where the v5 repair actually fires. Verified it skips with the gate off and still passes on 5.5.0.
| rope_scaling = _rope_scaling_as_dict( | ||
| getattr(getattr(self, "_unsloth_rope_config", None), "rope_scaling", None) or {} | ||
| ) |
There was a problem hiding this comment.
For transformers 5.x configs, this repo documents rope_parameters as the replacement for rope_scaling in unsloth/models/_utils.py, but this lookup only checks _unsloth_rope_config.rope_scaling. If LlamaExtendedRotaryEmbedding is constructed with a v5 config carrying Llama-3.2 parameters there, the lookup becomes {} and silently falls back to factor 8 instead of the configured factor 32; include config.rope_parameters in the lookup so the config-provided factor is honored in v5 as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Checked this empirically. On transformers 5.x, config.rope_scaling is a backward compatible property that returns rope_parameters (configuration_utils.py), so getattr(config, rope_scaling) already yields the v5 factors. Built LlamaExtendedRotaryEmbedding with a real transformers 5.5.0 LlamaConfig carrying factor 32 and _apply_inv_freq_scaling divides by 32, not 8. The base LlamaRotaryEmbedding config path and _compute_config_rope_inv_freq read rope_scaling the same way, so reading it here stays consistent and does not fall back to empty on v5.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4fb24d3c9c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| rope_scaling = _rope_scaling_as_dict( | ||
| getattr(getattr(self, "_unsloth_rope_config", None), "rope_scaling", None) or {} | ||
| ) | ||
| scale_factor = rope_scaling.get("factor", 8) |
There was a problem hiding this comment.
In the patched Llama attention constructor, the llama3 branch still creates LlamaExtendedRotaryEmbedding with only dim, max_position_embeddings, and base (unsloth/models/_utils.py:2832-2837), so _unsloth_rope_config is None on the normal codegen path. As a result, this lookup always falls back to factor=8 for that path, and Llama-3.2-style configs with rope_scaling.factor=32 still get the wrong extended RoPE scaling despite this change.
Useful? React with 👍 / 👎.
| wrapper = torch.nn.Module() | ||
| wrapper.add_module("rotary_emb", rot) | ||
| wrapper.config = config | ||
| loader._fix_rope_inv_freq(wrapper) |
There was a problem hiding this comment.
When this CUDA test runs under supported transformers 4.x, loader._fix_rope_inv_freq returns immediately because _NEEDS_ROPE_FIX is false (unsloth/models/loader.py:258), so the buffers blanked just above are never restored and the subsequent equality check fails. Since this is explicitly simulating the transformers v5 buffer-blank path, the test needs to skip or force the guard in non-v5 environments.
Useful? React with 👍 / 👎.
…formers 4.x - patch_llama_rope_scaling now builds the llama3 extended rotary with config=self.config so it reads the real factor (32 for Llama-3.2) instead of falling back to 8; the template already references self.config. - test_v5_blank_repair_roundtrip now skips when loader._NEEDS_ROPE_FIX is False, since _fix_rope_inv_freq is a no-op on transformers 4.x and cannot restore the blanked buffers there.
|
@codex review |
|
Codex Review: Didn't find any major issues. 🚀 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
… tests These asyncio.wait_for guards bound test setup and cross-task event signaling that complete near-instantly on success; the 0.2s budget is a latency assertion in disguise and times out under CI scheduling load (seen on the 3.11 matrix leg while 3.10/3.12/3.13 pass the same commit). 5.0s matches the timeout used elsewhere in the suite and still fails fast on a real hang. No test relies on the guard expiring.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bdeb2bf05c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| rope_scaling = _rope_scaling_as_dict( | ||
| getattr(getattr(self, "_unsloth_rope_config", None), "rope_scaling", None) or {} | ||
| ) | ||
| scale_factor = rope_scaling.get("factor", 8) |
There was a problem hiding this comment.
When this is built from a Transformers 5 config, the scaling fields live under config.rope_parameters; this repo also notes in unsloth/models/_utils.py that rope_parameters replaces rope_scaling for v5 configs. This new lookup only inspects _unsloth_rope_config.rope_scaling, so a Llama-3.2-style v5 config with rope_parameters.factor == 32 falls through to the default factor 8 and still computes the wrong extended RoPE frequencies for long-context runs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Verified on real v5 configs (transformers 5.0-5.13): config.rope_scaling is a back-compat property that returns rope_parameters, so a Llama-3.2 v5 config resolves factor 32 here today, not 8. But relying on that shim is fragile since a future 5.x may drop it, so I now read rope_parameters directly with a rope_scaling fallback. Added test_extended_rotary_reads_rope_parameters_v5, which sets rope_scaling=None and rope_parameters.factor=32 and fails on the old single-field read (resolves to 8).
transformers v5 stores llama3 scaling under config.rope_parameters and exposes rope_scaling only as a back-compat property. Reading that property works on 5.0-5.13 (verified: factor resolves to 32 for Llama-3.2), but a future release may drop the shim, after which the subclass path would fall back to factor 8. Read either field so the factor survives the rename. Adds test_extended_rotary_reads_rope_parameters_v5 (fails on the old single-field read: rope_parameters-only config resolves to 8, not 32).
|
@codex review |
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Breezy! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Summary
Follow-up to #6907. Adds a generic guard for the whole "RoPE scaling dropped when transformers v5 blanks buffers on load" bug class, and fixes a related hardcoded factor in the extended rotary.
Detection: a v5-blank round-trip guard
#6907 fixed the llama3 case and added a guard that the repair routes through the shared helper, but that only checks wiring at one call site. This adds a behavioral guard that generalizes:
build each rotary from a scaled config, snapshot its non-persistent buffers, blank them (exactly what transformers v5 meta-load does), run
loader._fix_rope_inv_freq, then assert every buffer is restored to its scaled value.It encodes no scaling math, so it guards any rotary that keeps scaling in a buffer. Parametrized over the two families that do: llama3 (
LlamaRotaryEmbedding,inv_freq) and longrope (LongRopeRotaryEmbedding,short_inv_freq/long_inv_freq). Validated to fail on the pre-fix repair and pass on the fix.I audited the other rotaries while here: Gemma / Gemma2 write cos/sin caches directly and register no
inv_freqbuffer, so they are immune; native transformers rotaries (Gemma3, etc.) carryoriginal_inv_freqand are restored by transformers' own_init_weights.Fix: extended rotary honors the config factor
LlamaExtendedRotaryEmbedding._apply_inv_freq_scalinghardcoded factor 8 / low 1 / high 4 / original_max 8192, which is wrong for a model with a different factor (Llama-3.2 uses 32). It now reads these from the config (stashed on the module in #6907), falling back to the Llama-3.1 defaults when built without a config (the legacy per-layer codegen path), so there is no behavior change on that path.Tests
tests/utils/test_rope_scaling_drift.py:test_v5_blank_repair_roundtrip[llama3, longrope](CUDA): the round-trip above.test_extended_rotary_reads_config_factor(CPU): the extended rotary divides by the config factor (32), not 8.Full suite: 13 passed, 1 skipped. Both new guards were confirmed to fail on their pre-fix code.
Also: de-flake the passthrough stream tests
Unrelated to RoPE but folded in to unblock this PR's own CI.
studio/backend/tests/test_openai_tool_passthrough.pyguarded async setup and cross-task event signaling withasyncio.wait_for(..., timeout = 0.2). Those guards catch a genuine hang, but 0.2s is a latency assertion in disguise: it times out under CI scheduling load (this PR's Backend CI failed only on the 3.11 matrix leg while 3.10 / 3.12 / 3.13 passed the same commit). Raised the 9 guards to 5.0s, matching the timeout used elsewhere in the suite. No test relies on the guard expiring, and forcing a tiny timeout reproduces the exact CI failure locally, confirming the value (not the logic) was the fragile part.