unsloth/tests/utils/test_rope_scaling_drift.py

"""Guard against config.rope_scaling being silently dropped (issue #2405):
the replacement rotary classes ignored it on the config path, so Llama-3.1
ran with unscaled RoPE and produced gibberish past ~32K tokens.

Three layers: (1) AST tripwire; (2) CPU checks of the pure helper
_compute_config_rope_inv_freq vs ROPE_INIT_FUNCTIONS; (3) CUDA checks on the
real class (skipped without a real device). Layers 2-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"
LOADER_PY = REPO_ROOT / "unsloth" / "models" / "loader.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 _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()
    # 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, called, _ = _iter_names_and_calls(scope)
    assert "rope_scaling" in names, (
        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)."
    )


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)."
    )


# --- 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: scaled inv_freq must differ from vanilla.
    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 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()
    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})
    try:
        object_config.rope_scaling = FakeLinearRopeScalingConfig()
    except Exception:
        pytest.skip(
            "transformers strict-validates rope_scaling to dict/RopeParameters/None, "
            "so object-style config.rope_scaling (and the delegation retry it "
            "exercises) is unreachable on this version."
        )
    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)