diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 022c796ac..f7c338d76 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -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/ diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py new file mode 100644 index 000000000..fae5a7d8a --- /dev/null +++ b/tests/utils/test_rope_scaling_drift.py @@ -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) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d5883e434..2d31f71ab 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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