mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
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:
parent
7467064c66
commit
d24ee77f17
3 changed files with 436 additions and 20 deletions
21
.github/workflows/consolidated-tests-ci.yml
vendored
21
.github/workflows/consolidated-tests-ci.yml
vendored
|
|
@ -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/
|
||||
|
|
|
|||
312
tests/utils/test_rope_scaling_drift.py
Normal file
312
tests/utils/test_rope_scaling_drift.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue