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
parent 7467064c66
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/

View file

@ -0,0 +1,312 @@
"""Guard for config.rope_scaling being silently dropped (issue #2405).
Unsloth's replacement rotary classes ignored rope_scaling when constructed
from a config (the modern-transformers path), so Llama-3.1 ran with unscaled
RoPE and collapsed into gibberish past ~32K tokens.
Layers: (1) AST tripwire, stdlib only; (2) CPU checks of the pure helper
_compute_config_rope_inv_freq against transformers' ROPE_INIT_FUNCTIONS;
(3) CUDA checks instantiating the real class (skipped without a real device,
probed by allocating a tensor so import-time CUDA spoofs cannot fool the gate).
Layers 2 and 3 fail on the unfixed code.
"""
import ast
import math
from pathlib import Path
import pytest
import torch
def _has_real_cuda():
try:
torch.zeros(1).to("cuda")
return True
except Exception:
return False
HAS_REAL_CUDA = _has_real_cuda()
requires_cuda = pytest.mark.skipif(
not HAS_REAL_CUDA,
reason = "LlamaRotaryEmbedding builds per-device CUDA caches in __init__",
)
REPO_ROOT = Path(__file__).resolve().parents[2]
LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py"
CLASS_NAME = "LlamaRotaryEmbedding"
# Llama-3.1-style rope_scaling.
LLAMA3_ROPE_SCALING = {
"rope_type": "llama3",
"factor": 8.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
}
ROPE_THETA = 500000.0
HEAD_DIM = 128
MAX_POS = 131072
# --- Layer 1: AST structural tripwire (stdlib only, no unsloth import) ---
def _load_class_init():
tree = ast.parse(LLAMA_PY.read_text())
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == CLASS_NAME:
for sub in node.body:
if isinstance(sub, ast.FunctionDef) and sub.name == "__init__":
return sub
raise AssertionError(
f"{CLASS_NAME}.__init__ not found in {LLAMA_PY}; if it was renamed or "
"moved, update this guard so RoPE scaling stays protected (issue #2405)"
)
def _config_branch(init_fn):
"""The `if config is not None:` block at the top of __init__."""
for node in init_fn.body:
if isinstance(node, ast.If):
test = node.test
is_config_test = (
isinstance(test, ast.Compare)
and isinstance(test.left, ast.Name)
and test.left.id == "config"
)
if is_config_test:
return node
return None
def test_config_path_inspects_rope_scaling():
init_fn = _load_class_init()
branch = _config_branch(init_fn)
assert branch is not None, (
f"{CLASS_NAME}.__init__ no longer has an `if config is not None:` "
"branch; the config constructor path must read config.rope_scaling so "
"scaled models (llama3/linear/longrope) are not silently unscaled "
"(issue #2405)"
)
names = set()
for stmt in branch.body:
for sub in ast.walk(stmt):
if isinstance(sub, ast.Attribute):
names.add(sub.attr)
elif isinstance(sub, ast.Constant) and isinstance(sub.value, str):
names.add(sub.value)
assert "rope_scaling" in names, (
f"{CLASS_NAME}.__init__ config path does not reference `rope_scaling`. "
"When a rotary class is built straight from a config (the path modern "
"transformers takes, since rotary moved to LlamaModel), the llama3 / "
"linear / longrope scaling must still be applied; otherwise long inputs "
"produce repeated-pattern gibberish (issue #2405)."
)
called = {
sub.func.id
for stmt in branch.body
for sub in ast.walk(stmt)
if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)
}
assert "_compute_config_rope_inv_freq" in called, (
f"{CLASS_NAME}.__init__ config path no longer calls "
"_compute_config_rope_inv_freq; the CPU behavioral tests below cover "
"that helper directly, so the constructor must stay wired to it or "
"scaled configs silently lose RoPE scaling again (issue #2405)."
)
# --- Layer 2: CPU behavioral guard (pure helper, no instantiation) ---
def _make_config(rope_scaling):
from transformers import LlamaConfig
return LlamaConfig(
hidden_size = 256,
num_attention_heads = 2,
num_key_value_heads = 2,
head_dim = HEAD_DIM,
rope_theta = ROPE_THETA,
max_position_embeddings = MAX_POS,
rope_scaling = rope_scaling,
)
def _unsloth_rotary(config):
from unsloth.models import llama as llama_mod
return llama_mod.LlamaRotaryEmbedding(config = config)
def _reference_inv_freq(config, rope_type):
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
inv_freq, _attention_factor = ROPE_INIT_FUNCTIONS[rope_type](config, "cpu")
return inv_freq.float().cpu()
def _vanilla_inv_freq():
return 1.0 / (
ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2, dtype = torch.int64).float() / HEAD_DIM)
)
def _compute_helper(config, rope_scaling):
from unsloth.models.llama import _compute_config_rope_inv_freq
return _compute_config_rope_inv_freq(config, rope_scaling)
def test_llama3_scaling_applied_to_inv_freq():
config = _make_config(LLAMA3_ROPE_SCALING)
got, attention_scaling = _compute_helper(config, config.rope_scaling)
expected = _reference_inv_freq(config, "llama3")
vanilla = _vanilla_inv_freq()
# Guard against a vacuous test.
assert not torch.allclose(
expected, vanilla, rtol = 1e-4
), "test setup error: llama3-scaled inv_freq should differ from vanilla"
assert got is not None, (
"_compute_config_rope_inv_freq returned None for a llama3 config; the "
"config path is dropping config.rope_scaling, so long-context inference "
"degrades into repeated-pattern gibberish (issue #2405)."
)
got = got.float().cpu()
assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), (
"inv_freq for a llama3 config does not match transformers' llama3 RoPE "
"scaling (issue #2405).\n"
f"got[:6]={got[:6].tolist()}\nexpected[:6]={expected[:6].tolist()}"
)
def test_default_rope_type_matches_vanilla_inv_freq():
config = _make_config(None)
got, attention_scaling = _compute_helper(config, {"rope_type": "default"})
assert got is not None
vanilla = _vanilla_inv_freq()
assert torch.allclose(got.float().cpu(), vanilla, rtol = 1e-4, atol = 1e-6), (
"default rope_type must equal the vanilla inv_freq; "
f"got[:6]={got[:6].tolist()} vanilla[:6]={vanilla[:6].tolist()}"
)
def _cos_at_position(rot, position):
"""cos row at one position, built like _set_cos_sin_cache but CPU-only."""
inv_freq = rot.inv_freq.float().cpu()
t = torch.tensor([position], dtype = torch.float32)
t = rot._apply_time_scaling(t.clone()) if hasattr(rot, "_apply_time_scaling") else t
freqs = torch.outer(t, inv_freq)
emb = torch.cat((freqs, freqs), dim = -1)
return emb.cos().squeeze(0)
# --- Layer 3: CUDA behavioral guard (real instantiation needs a device) ---
@requires_cuda
def test_constructor_applies_llama3_scaling():
config = _make_config(LLAMA3_ROPE_SCALING)
rot = _unsloth_rotary(config)
got = rot.inv_freq.float().cpu()
expected = _reference_inv_freq(config, "llama3")
assert torch.allclose(
got, expected, rtol = 1e-4, atol = 1e-6
), "LlamaRotaryEmbedding built from a llama3 config produced unscaled inv_freq (issue #2405)."
@requires_cuda
def test_constructor_unscaled_config_uses_vanilla_inv_freq():
rot = _unsloth_rotary(_make_config(None))
got = rot.inv_freq.float().cpu()
vanilla = _vanilla_inv_freq()
assert torch.allclose(
got, vanilla, rtol = 1e-4, atol = 1e-6
), "LlamaRotaryEmbedding with no rope_scaling must use the vanilla inv_freq"
@requires_cuda
def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position():
scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
unscaled = _unsloth_rotary(_make_config(None))
pos = 10000
cos_scaled = _cos_at_position(scaled, pos)
cos_unscaled = _cos_at_position(unscaled, pos)
assert not torch.allclose(cos_scaled, cos_unscaled, rtol = 1e-4, atol = 1e-5), (
f"cos values at position {pos} are identical for a llama3-scaled and an "
"unscaled rotary embedding, which means scaling was dropped (issue "
"#2405). With correct llama3 scaling the low-frequency bands shrink by "
"up to 8x and must change the angles at long positions."
)
@requires_cuda
def test_extended_cache_keeps_scaling_after_growth():
scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
# Grow past the initial cache size (mirrors long-context decode).
dummy = torch.zeros(1, dtype = torch.float32)
scaled.extend_rope_embedding(dummy, seq_len = 40960)
config = _make_config(LLAMA3_ROPE_SCALING)
expected = _reference_inv_freq(config, "llama3")
got = scaled.inv_freq.float().cpu()
assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), (
"growing the RoPE cache (extend_rope_embedding) must preserve llama3 "
"scaling of inv_freq; long-context decode loses scaling otherwise "
"(issue #2405)."
)
def test_object_style_rope_scaling_does_not_crash():
# Object-style rope_scaling must be normalized, not .get()'d directly.
from dataclasses import dataclass
from unsloth.models.llama import _compute_config_rope_inv_freq
@dataclass
class FakeRopeScalingConfig:
rope_type: str = "llama3"
factor: float = 8.0
low_freq_factor: float = 1.0
high_freq_factor: float = 4.0
original_max_position_embeddings: int = 8192
config = _make_config(LLAMA3_ROPE_SCALING)
inv_freq, attention_scaling = _compute_config_rope_inv_freq(config, FakeRopeScalingConfig())
assert inv_freq is not None, (
"object-style (non-dict) config.rope_scaling must be normalized, not "
"dropped; otherwise scaled models silently lose RoPE scaling again "
"(issue #2405)."
)
expected = _reference_inv_freq(config, "llama3")
assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6)
def test_object_style_rope_scaling_on_config_delegates_correctly():
# 'linear' has no inline fallback; only the normalized-config retry passes this.
from dataclasses import dataclass
from unsloth.models.llama import _compute_config_rope_inv_freq
@dataclass
class FakeLinearRopeScalingConfig:
rope_type: str = "linear"
factor: float = 4.0
dict_config = _make_config({"rope_type": "linear", "factor": 4.0})
expected = _reference_inv_freq(dict_config, "linear")
object_config = _make_config({"rope_type": "linear", "factor": 4.0})
object_config.rope_scaling = FakeLinearRopeScalingConfig()
inv_freq, attention_scaling = _compute_config_rope_inv_freq(
object_config, object_config.rope_scaling
)
assert inv_freq is not None, (
"linear rope_scaling exposed as a config object was silently dropped; "
"delegation must retry with a config copy carrying the normalized dict "
"(issue #2405)."
)
assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6)

View file

@ -1622,6 +1622,93 @@ def _get_rope_theta(config, default = 10000.0):
return default
def _rope_scaling_as_dict(rope_scaling):
"""Normalize config.rope_scaling (dict or config object) to a dict; {} on failure."""
if isinstance(rope_scaling, dict):
return rope_scaling
for converter in ("to_dict", "dict"):
fn = getattr(rope_scaling, converter, None)
if callable(fn):
try:
d = fn()
if isinstance(d, dict):
return d
except Exception:
pass
try:
return {k: v for k, v in vars(rope_scaling).items() if not k.startswith("_")}
except TypeError:
return {}
def _llama3_inv_freq_from_config(
config,
rope_scaling,
device = "cpu",
):
"""llama3 inv_freq with factors from config; fallback when modeling_rope_utils is missing."""
base = _get_rope_theta(config, default = 10000.0)
dim = getattr(config, "head_dim", None)
if dim is None:
dim = int(config.hidden_size // config.num_attention_heads)
inv_freq = 1.0 / (
base ** (torch.arange(0, dim, 2, dtype = torch.int64, device = device).float() / dim)
)
scale_factor = rope_scaling.get("factor", 8.0)
low_freq_factor = rope_scaling.get("low_freq_factor", 1.0)
high_freq_factor = rope_scaling.get("high_freq_factor", 4.0)
old_context_len = rope_scaling.get("original_max_position_embeddings", 8192)
low_freq_wavelen = old_context_len / low_freq_factor
high_freq_wavelen = old_context_len / high_freq_factor
assert low_freq_wavelen != high_freq_wavelen
# Vectorized meta-llama bands: high freqs kept, low divided by factor, medium blended.
wavelen = 2 * math.pi / inv_freq
scaled = torch.where(wavelen > low_freq_wavelen, inv_freq / scale_factor, inv_freq)
smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
smoothed = (1 - smooth) * inv_freq / scale_factor + smooth * inv_freq
is_medium = (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen)
return torch.where(is_medium, smoothed, scaled)
def _compute_config_rope_inv_freq(config, rope_scaling):
"""(inv_freq, attention_scaling) per config.rope_scaling via transformers'
ROPE_INIT_FUNCTIONS, with an inline llama3 fallback; (None, 1.0) on failure."""
original_rope_scaling = rope_scaling
rope_scaling = _rope_scaling_as_dict(rope_scaling)
rope_type = rope_scaling.get("rope_type", None) or rope_scaling.get("type", None)
try:
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type]
try:
inv_freq, attention_scaling = rope_init_fn(config, torch.device("cpu"))
except Exception:
# Object-style rope_scaling: retry with a config copy carrying the plain dict.
if isinstance(original_rope_scaling, dict):
raise
import copy as _copy
config_copy = _copy.copy(config)
config_copy.rope_scaling = rope_scaling
inv_freq, attention_scaling = rope_init_fn(config_copy, torch.device("cpu"))
return inv_freq.to(dtype = torch.float32, device = "cpu"), float(attention_scaling)
except Exception as exception:
if rope_type == "llama3":
try:
return _llama3_inv_freq_from_config(config, rope_scaling), 1.0
except Exception:
pass
logger.warning_once(
f"Unsloth: Could not apply RoPE scaling '{rope_type}' from config "
f"({type(exception).__name__}: {exception}); falling back to unscaled RoPE. "
"Long-context generation may degrade."
)
return None, 1.0
# Solves https://github.com/unslothai/unsloth/issues/168
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
@ -1640,6 +1727,12 @@ class LlamaRotaryEmbedding(torch.nn.Module):
config = None, # [TODO] Hack to pass in config - need to remove later
):
super().__init__()
# cos/sin multiplier (1.0 except yarn / longrope); set before any cache build.
self.attention_scaling = 1.0
# Base-class-from-config path (modern transformers): derive inv_freq like
# transformers so config.rope_scaling is not dropped (#2405). Scaled
# subclasses are excluded to avoid double-scaling.
config_inv_freq = None
if config is not None:
# [TODO] Hack to pass in config - need to remove later
base = _get_rope_theta(config, default = base)
@ -1652,6 +1745,13 @@ class LlamaRotaryEmbedding(torch.nn.Module):
device = DEVICE_TYPE_TORCH
max_position_embeddings = config.max_position_embeddings
rope_scaling = getattr(config, "rope_scaling", None)
if rope_scaling is not None and type(self) is LlamaRotaryEmbedding:
config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq(
config,
rope_scaling,
)
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
@ -1660,12 +1760,17 @@ class LlamaRotaryEmbedding(torch.nn.Module):
self.multi_gpu_cos_cached = [None] * DEVICE_COUNT
self.multi_gpu_sin_cached = [None] * DEVICE_COUNT
# Normal Llama-3 RoPE
inv_freq = 1.0 / (
self.base
** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim)
)
inv_freq = self._apply_inv_freq_scaling(inv_freq)
if config_inv_freq is not None:
inv_freq = config_inv_freq # already scaled; skip subclass scaling
else:
# Normal Llama-3 RoPE
inv_freq = 1.0 / (
self.base
** (
torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim
)
)
inv_freq = self._apply_inv_freq_scaling(inv_freq)
self.register_buffer("inv_freq", inv_freq, persistent = False)
# Build here to make `torch.jit.trace` work.
@ -1704,8 +1809,10 @@ class LlamaRotaryEmbedding(torch.nn.Module):
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim = -1)
cos = emb.cos().to(dtype = dtype, device = device, non_blocking = True)
sin = emb.sin().to(dtype = dtype, device = device, non_blocking = True)
# Applied here so attention_scaling survives extend_rope_embedding rebuilds;
# default 1.0 keeps unscaled paths bit-identical.
cos = (emb.cos() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True)
sin = (emb.sin() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True)
self.multi_gpu_cos_cached[device.index] = cos
self.multi_gpu_sin_cached[device.index] = sin
return cos, sin