diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index b976654f87..98f7e2db62 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -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() diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 564be09578..c25a031b82 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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 diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index ba23197861..13342157b0 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -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: