From 534c877d2136b47b0ceec25cc45900c7f7f15e6d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 04:20:41 -0700 Subject: [PATCH] Keep native RoPE scaling when extending context; carry rope_theta for linear (#7028) * Keep native RoPE scaling when extending context; carry rope_theta for linear When max_seq_length exceeds a model's native window, the loader overwrote the model's rope_scaling with linear scaling. For models that already ship a scaled RoPE (llama3/yarn/longrope) that is far worse for long context, and on transformers v5 the linear dict omitted rope_theta (v5 keeps it under rope_parameters), so the rotary base fell back to 10000 and broke past ~8K tokens. Keep the native scaling and just widen the window; only synthesize linear for plain-RoPE models, and carry rope_theta so v5 keeps the real base. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Only preserve native llama3 when extending context; keep linear fallback otherwise The patched attention constructor (patch_llama_rope_scaling) rebuilds only linear, llama3 and longrope and its longrope branch reads a top-level original_max_position_embeddings, so preserving yarn or a nested-only longrope config would raise during construction on transformers <= 4.47.1. Keep only llama3 native; yarn/longrope/other types fall back to the linear override, still carrying rope_theta. * Correct long-context extension comment to match llama3-only preservation --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/utils/test_rope_scaling_drift.py | 33 ++++++++++++++ unsloth/models/llama.py | 61 +++++++++++++++++--------- 2 files changed, 73 insertions(+), 21 deletions(-) diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 7a738e236c..b2ec1e5a20 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -257,6 +257,39 @@ def test_recompute_helper_scales_on_cpu(): ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." +def test_extended_rope_scaling_keeps_llama3_and_carries_theta(): + # Long-context extension keeps native llama3, but falls back to linear for every other + # type (the patched attention constructor only rebuilds linear/llama3/longrope), and the + # linear dict carries rope_theta so transformers v5 does not fall back to base 10000. + from types import SimpleNamespace + + from unsloth.models.llama import _extended_rope_scaling + + # llama3 model: keep native scaling, do not synthesize linear. + scaling, native = _extended_rope_scaling(_make_config(LLAMA3_ROPE_SCALING), 2.0) + assert ( + scaling is None and native == "llama3" + ), "must keep native llama3 scaling instead of overwriting it with linear." + + # yarn is not rebuildable by the patcher -> keep the safe linear fallback, not native. + yarn = SimpleNamespace(rope_scaling = {"rope_type": "yarn", "factor": 2.0}, rope_theta = 500000.0) + scaling, _ = _extended_rope_scaling(yarn, 2.0) + assert scaling == { + "type": "linear", + "factor": 2.0, + "rope_theta": 500000.0, + }, f"yarn must fall back to linear (patcher cannot rebuild it), got {scaling}." + + # plain RoPE with theta only under v5 rope_parameters: linear must carry rope_theta. + v5 = SimpleNamespace(rope_parameters = {"rope_type": "default", "rope_theta": 1000000.0}) + scaling, _ = _extended_rope_scaling(v5, 2.0) + assert scaling == { + "type": "linear", + "factor": 2.0, + "rope_theta": 1000000.0, + }, f"linear override dropped rope_theta on v5 (got {scaling}); base would fall back to 10000." + + def test_extended_rotary_reads_config_factor(): # LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8 # (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405). diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 1f43f61443..05523bc27b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1651,6 +1651,26 @@ def _rope_scaling_as_dict(rope_scaling): return {} +def _extended_rope_scaling(config, factor): + """RoPE scaling to extend a model past its native window. Keeps native llama3 as-is + (linear extension is far worse for long context); everything else gets linear. Returns + (scaling_or_None, type): None keeps llama3. The linear dict carries rope_theta so + transformers v5 (which stores it under rope_parameters) keeps the real base, not 10000. + Only llama3 is preserved because patch_llama_rope_scaling can only rebuild linear/llama3/ + longrope and its longrope branch needs a top-level original_max_position_embeddings.""" + existing = _rope_scaling_as_dict( + getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {} + ) + existing_type = existing.get("rope_type") or existing.get("type") + if existing_type == "llama3": + return None, existing_type + return { + "type": "linear", + "factor": factor, + "rope_theta": _get_rope_theta(config), + }, existing_type + + def _llama3_inv_freq_from_config( config, rope_scaling, @@ -2518,34 +2538,33 @@ class FastLlamaModel: max_seq_length = model_max_seq_length if (rope_scaling is None) and (max_seq_length > model_max_seq_length): - rope_scaling = max_seq_length / model_max_seq_length + factor = max_seq_length / model_max_seq_length if fast_inference: raise NotImplementedError( "Unsloth: Fast inference does not yet work with RoPE Scaling." ) - logger.warning_once( - f"Unsloth: {model_name} can only handle sequence lengths of at most " - f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of " - f"{round(rope_scaling, 3)}, it can be magically be extended to " - f"{max_seq_length}!" - ) - - # Warn RoPE scaling isn't allowed - if not has_rope_scaling: - raise RuntimeError( - f"However, {model_name} doesn't support RoPE Scaling!\n" - "Please file a feature request at https://github.com/unslothai/unsloth." + linear_scaling, native_type = _extended_rope_scaling(model_config, factor) + if linear_scaling is not None: + logger.warning_once( + f"Unsloth: {model_name} can only handle sequence lengths of at most " + f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of " + f"{round(factor, 3)}, it can be magically be extended to " + f"{max_seq_length}!" + ) + if not has_rope_scaling: + raise RuntimeError( + f"However, {model_name} doesn't support RoPE Scaling!\n" + "Please file a feature request at https://github.com/unslothai/unsloth." + ) + kwargs["rope_scaling"] = linear_scaling + else: + # Native llama3 scaling already handles long context; just widen the window. + logger.warning_once( + f"Unsloth: extending {model_name} to {max_seq_length} using its native " + f"{native_type} RoPE scaling." ) - - rope_scaling = { - "type": "linear", - "factor": rope_scaling, - } - - # Add to kwargs - kwargs["rope_scaling"] = rope_scaling from .loader_utils import ( check_and_disable_bitsandbytes_loading,