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,