Fix Llama 3.1+ rope scaling dropped on the FastLanguageModel path (long inputs become gibberish past ~29K tokens) (#6197)

* Fix config.rope_scaling being dropped by the replaced rotary embedding (#2405)

On modern transformers, LlamaModel builds its rotary embedding from config
using unsloth's replacement LlamaRotaryEmbedding class, whose config path
computed vanilla inv_freq and ignored config.rope_scaling entirely. The
llama3/linear/longrope dispatch in patch_llama_rope_scaling rewrites
LlamaAttention.__init__, which no longer constructs rotary embeddings, so it
never fires; the model-level rotary is then copied onto every attention
layer. Result: Llama-3.1/3.2/3.3 ran with unscaled RoPE on the
FastLanguageModel path and collapsed into repetition loops past roughly 29K
tokens (PASS at 28867, FAIL at 31767 in needle retrieval). FastModel was
unaffected because vision.py keeps transformers' own rotary. qwen2, qwen3,
qwen3_moe, mistral and cohere assign the same base class, so any rope-scaled
config of those families was equally exposed.

The fix makes the base class config path compute inv_freq and
attention_scaling via transformers' ROPE_INIT_FUNCTIONS (covers llama3,
linear, dynamic, yarn, longrope), with an inline llama3 fallback reading
factors from config for older transformers, degrading to prior behavior on
any failure. attention_scaling is applied in _set_cos_sin_cache (1.0 default,
exact no-op for unscaled paths) and persists across extend_rope_embedding.
A type(self) guard prevents double-scaling via the legacy scaled subclasses.

Adds tests/utils/test_rope_scaling_drift.py (AST tripwire + behavioral
inv_freq/cos-cache/extension checks, validated to fail 4 of 5 on the unfixed
code) and wires it into the existing consolidated CI HARD GATE step.

Verified on GPU: 48K-token needle retrieval flips FAIL to PASS for
FastLanguageModel in bf16 and 4bit, 20K stays PASS, scaled inv_freq matches
transformers exactly, and the left-padded batch generation guard still gets
exact solo-vs-batched token matches.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: normalize object-style rope_scaling, vectorize llama3 fallback

config.rope_scaling can be a config object rather than a dict on newer
transformers; _rope_scaling_as_dict normalizes it (to_dict/dict/vars
fallbacks) before any .get() access, with a regression test using a
dataclass stand-in. The inline llama3 fallback now uses torch.where instead
of a per-frequency Python loop; verified bit-for-bit equal to transformers
ROPE_INIT_FUNCTIONS for factor 8 (Llama-3.1) and factor 32 (Llama-3.2).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: CPU-safe rope guard tests, normalized config for delegation

The rotary constructor builds per-device CUDA caches, so the behavioral tests
that instantiate it cannot run on GPU-less CI. Restructured into three layers:
the AST tripwire now also asserts the constructor stays wired to
_compute_config_rope_inv_freq; the CPU layer tests that pure helper directly
(llama3 dict, llama3 object, linear object, default type) with no
instantiation; the instantiation and cache tests are gated behind a real CUDA
probe (actual tensor allocation, so import-time CUDA spoofs cannot fool the
gate). Verified: 9 passed with GPU; 5 passed 4 skipped with CUDA hidden; 5
failed 4 skipped on the unfixed code in CPU mode.

Delegation to ROPE_INIT_FUNCTIONS now retries with a shallow config copy
carrying the normalized rope_scaling dict when the original was an object the
installed transformers cannot read; covered by a linear-object test, which has
no inline fallback and passes only through that retry path.

* Tighten comments in rope scaling fix and guard test

Comment and docstring reduction only; verified code-identical with
scripts/comment_tools.py check --strip-docstrings (AST signature match on
both Python files). All guard tests unchanged: 20 passed with GPU, 5 passed
4 skipped with CUDA hidden.

* Apply repo kwarg-spacing format

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-11 07:48:21 -07:00 committed by GitHub
commit d24ee77f17
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 436 additions and 20 deletions

View file

@ -333,19 +333,16 @@ jobs:
run: |
python -m pytest -v --tb=short tests/test_callback_signature_drift.py
- name: batched left-padding generation guard (HARD GATE)
# Guards _fast_prepare_inputs_for_generation against the bug class of
# issues #1066 / #3699: position_ids taken from cache_position (which
# counts left-pad tokens) or the 2D attention mask truncated to its
# last column. Both shipped in cc4c5d77 and were fixed by #2216 and
# #4100; nothing tested this path, so each regression reached users.
# Layer 1 in the file is stdlib-ast-only (survives unsloth import
# breakage), layer 2 calls the real function on CPU via the
# tests/conftest.py CUDA spoof. Validated to fail on the pre-#2216
# and pre-#4100 code states; staging proof on GPU-less runners:
# danielhanchen/unsloth-staging-2 PR 170 (green, gate passed in all combos) / PR 172 (red, gate failed in all combos).
- name: generation correctness guards (HARD GATE)
# Deterministic CPU guards, each validated to fail on its pre-fix code:
# leftpad = batched left-padded generation (#1066/#3699, fixed by
# #2216 + #4100; staging proof: unsloth-staging-2 PRs 170/172);
# rope_scaling_drift = config.rope_scaling dropped by replaced rotary
# classes (#2405). AST checks run first so import breakage cannot mask them.
run: |
python -m pytest -v --tb=short tests/utils/test_prepare_inputs_leftpad.py
python -m pytest -v --tb=short \
tests/utils/test_prepare_inputs_leftpad.py \
tests/utils/test_rope_scaling_drift.py
- name: unsloth Bucket-A — CPU tests not in Repo tests (CPU)
# CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/