ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix llama3 RoPE scaling dropped on transformers v5 (#6907)

Browse source 
* Fix llama3 RoPE scaling dropped on transformers v5

transformers v5 loads on meta then blanks non-persistent buffers, so
_fix_rope_inv_freq rebuilds inv_freq after load. It recomputed a vanilla
inv_freq and applied _apply_inv_freq_scaling, a no-op on the base
LlamaRotaryEmbedding used by the config/llama3 path, so inv_freq ended up
divided by 1 instead of the config factor (8 for Llama 3.1, 32 for Llama
3.2). This corrupts long-range positions and inflates long-context loss
about 3-5x. transformers 4.x was unaffected.

Route __init__ and the v5 repair through one _unsloth_recompute_inv_freq
so they cannot diverge, and stash the config on the rotary module so the
repair can rebuild the same scaled value.

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

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

* Add test for llama3 RoPE scaling under the transformers v5 repair

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

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

* Update RoPE drift guard for the recompute refactor and guard the v5 repair

The drift guard's AST tripwire asserted the config-scaling call lived in the
if config is not None branch of LlamaRotaryEmbedding.__init__. The fix moved
that into _unsloth_recompute_inv_freq, so follow it there (with a fallback to
the old inline branch) and add a guard that loader._fix_rope_inv_freq rebuilds
inv_freq through the same helper. Also add a CPU functional check of the helper
and drop the redundant standalone test.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-06 09:13:14 -07:00 • committed by GitHub
commit 2fada48ef5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 132 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

128
tests/utils/test_rope_scaling_drift.py
View file

@ -31,6 +31,7 @@ requires_cuda = pytest.mark.skipif(
REPO_ROOT = Path(__file__).resolve().parents[2]
LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py"
LOADER_PY = REPO_ROOT / "unsloth" / "models" / "loader.py"
CLASS_NAME = "LlamaRotaryEmbedding"
@ -78,42 +79,88 @@ def _config_branch(init_fn):
return None
def _iter_names_and_calls(node):
"""(attribute/string names, bare-name calls, method-call attrs) under node."""
names, calls, call_attrs = set(), set(), set()
for sub in ast.walk(node):
if isinstance(sub, ast.Attribute):
names.add(sub.attr)
elif isinstance(sub, ast.Constant) and isinstance(sub.value, str):
names.add(sub.value)
elif isinstance(sub, ast.Call):
if isinstance(sub.func, ast.Name):
calls.add(sub.func.id)
elif isinstance(sub.func, ast.Attribute):
call_attrs.add(sub.func.attr)
return names, calls, call_attrs
def _find_method(source_path, class_name, method_name):
for node in ast.walk(ast.parse(source_path.read_text())):
if isinstance(node, ast.ClassDef) and node.name == class_name:
for sub in node.body:
if isinstance(sub, ast.FunctionDef) and sub.name == method_name:
return sub
return None
def _find_function(source_path, function_name):
for node in ast.walk(ast.parse(source_path.read_text())):
if isinstance(node, ast.FunctionDef) and node.name == function_name:
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)"
)
# inv_freq is derived through the shared _unsloth_recompute_inv_freq helper
# (or still inlined in the config branch on older layouts); whichever scope
# holds the scaling must read config.rope_scaling and call
# _compute_config_rope_inv_freq, else scaled models run unscaled (#2405).
_, _, init_call_attrs = _iter_names_and_calls(init_fn)
scope = _find_method(LLAMA_PY, CLASS_NAME, "_unsloth_recompute_inv_freq")
if scope is not None:
assert "_unsloth_recompute_inv_freq" in init_call_attrs, (
f"{CLASS_NAME}.__init__ no longer derives inv_freq via "
"_unsloth_recompute_inv_freq; keep the constructor wired to the "
"shared scaling helper or scaled configs silently lose RoPE scaling "
"(issue #2405)."
)
else:
scope = _config_branch(init_fn)
assert scope is not None, (
f"{CLASS_NAME}.__init__ has neither a _unsloth_recompute_inv_freq "
"helper nor an `if config is not None:` branch; the config path must "
"apply llama3/linear/longrope scaling (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)
names, called, _ = _iter_names_and_calls(scope)
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)."
f"{CLASS_NAME} inv_freq computation does not reference `rope_scaling`; "
"scaled models (llama3/linear/longrope) would run unscaled and produce "
"repeated-pattern gibberish past the original context (issue #2405)."
)
assert "_compute_config_rope_inv_freq" in called, (
f"{CLASS_NAME} inv_freq computation no longer calls "
"_compute_config_rope_inv_freq; keep it wired or scaled configs silently "
"lose RoPE scaling again (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)."
def test_v5_repair_reuses_recompute():
# transformers v5 blanks non-persistent buffers on load, so
# loader._fix_rope_inv_freq rebuilds inv_freq; it must reuse the scaled
# recompute, since an unscaled rebuild re-drops llama3 scaling (#2405).
fix_fn = _find_function(LOADER_PY, "_fix_rope_inv_freq")
assert fix_fn is not None, (
"loader._fix_rope_inv_freq not found; if it was renamed, update this "
"guard so the v5 rope repair keeps applying config scaling (issue #2405)."
)
_, _, call_attrs = _iter_names_and_calls(fix_fn)
assert "_unsloth_recompute_inv_freq" in call_attrs, (
"loader._fix_rope_inv_freq no longer rebuilds inv_freq via "
"_unsloth_recompute_inv_freq; transformers v5 blanks the buffer on load "
"and an unscaled rebuild re-drops llama3 scaling (issue #2405)."
)
@ -189,6 +236,27 @@ def test_default_rope_type_matches_vanilla_inv_freq():
)
def test_recompute_helper_scales_on_cpu():
# Exercise the exact method loader._fix_rope_inv_freq calls, without CUDA.
from unsloth.models.llama import LlamaRotaryEmbedding, _get_rope_theta
def recompute(config):
rot = object.__new__(LlamaRotaryEmbedding)
rot.attention_scaling = 1.0
rot.base = _get_rope_theta(config, 10000.0)
rot.dim = config.head_dim
rot._unsloth_rope_config = config
return rot._unsloth_recompute_inv_freq().float().cpu()
config = _make_config(LLAMA3_ROPE_SCALING)
assert torch.allclose(
recompute(config), _reference_inv_freq(config, "llama3"), rtol = 1e-4, atol = 1e-6
), "_unsloth_recompute_inv_freq dropped llama3 scaling (issue #2405)."
assert torch.allclose(
recompute(_make_config(None)), _vanilla_inv_freq(), rtol = 1e-4, atol = 1e-6
), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
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()

41
unsloth/models/llama.py
View file

@ -1756,7 +1756,6 @@ class LlamaRotaryEmbedding(torch.nn.Module):
# 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)
@ -1769,32 +1768,17 @@ 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
# Kept so the v5 rope repair can rebuild the scaled inv_freq (#2405).
self._unsloth_rope_config = config
# Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this
self.current_rope_size = min(4 * 8192, self.max_position_embeddings)
self.multi_gpu_cos_cached = [None] * DEVICE_COUNT
self.multi_gpu_sin_cached = [None] * DEVICE_COUNT
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)
inv_freq = self._unsloth_recompute_inv_freq()
self.register_buffer("inv_freq", inv_freq, persistent = False)
# Build here to make `torch.jit.trace` work.
@ -1817,6 +1801,25 @@ class LlamaRotaryEmbedding(torch.nn.Module):
"""Override to apply custom inv_freq scaling (e.g., extended RoPE)."""
return inv_freq
def _unsloth_recompute_inv_freq(self):
# Config scaling (llama3/yarn) first, else vanilla + subclass scaling.
# Shared by __init__ and the v5 rope repair so they cannot diverge.
config = getattr(self, "_unsloth_rope_config", None)
config_inv_freq = None
rope_scaling = getattr(config, "rope_scaling", None) if config is not None else 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,
)
if config_inv_freq is not None:
return config_inv_freq
inv_freq = 1.0 / (
self.base
** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim)
)
return self._apply_inv_freq_scaling(inv_freq)
def _apply_time_scaling(self, t):
"""Override to apply custom time scaling (e.g., linear scaling)."""
return t

19
unsloth/models/loader.py
View file

@ -245,6 +245,7 @@ def _maybe_advise_fla_install(model_types):
"transformers will use a slower pure PyTorch path."
)
def _fix_rope_inv_freq(model):
"""Fix inv_freq corruption caused by transformers v5 meta-device loading.
@ -268,14 +269,18 @@ def _fix_rope_inv_freq(model):
and hasattr(module, "_apply_inv_freq_scaling")
and hasattr(module, "multi_gpu_cos_cached")
):
inv_freq = 1.0 / (
module.base
** (
torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float()
/ module.dim
if hasattr(module, "_unsloth_recompute_inv_freq"):
# Restore config scaling (llama3/yarn); unscaled here broke v5.
inv_freq = module._unsloth_recompute_inv_freq()
else:
inv_freq = 1.0 / (
module.base
** (
torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float()
/ module.dim
)
)
)
inv_freq = module._apply_inv_freq_scaling(inv_freq)
inv_freq = module._apply_inv_freq_scaling(inv_freq)
module.inv_freq = inv_freq
for device_idx in range(len(module.multi_gpu_cos_cached)):
if module.multi_gpu_cos_cached[device_idx] is not None: