mirror of https://github.com/razor-ai/soup.git
409 lines
15 KiB
Python
409 lines
15 KiB
Python
"""Long-context fine-tuning utilities — 128k+ token support.
|
|
|
|
Configures RoPE (Rotary Position Embedding) scaling to extend model context
|
|
windows beyond their pre-training length. Supports multiple scaling strategies:
|
|
|
|
- linear: Simple linear interpolation (PI) — good baseline
|
|
- dynamic: NTK-aware Dynamic scaling — better for large extensions
|
|
- yarn: YaRN (Yet another RoPE extensioN) — best quality for 4-8x extension
|
|
(v0.49.0 Part A — math kernel + config-emit; HF Transformers owns the
|
|
actual rotation under the hood)
|
|
- longrope: LongRoPE — progressive extension with search-based factors
|
|
- llama3: Llama 3.1 frequency-band NTK-aware scaling (v0.49.0 Part D)
|
|
|
|
Also handles gradient checkpointing configuration for memory efficiency
|
|
when training on very long sequences.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
# Supported RoPE scaling methods (v0.49.0 adds "llama3").
|
|
ROPE_SCALING_TYPES = ("linear", "dynamic", "yarn", "longrope", "llama3")
|
|
|
|
# Default context lengths for known model families.
|
|
MODEL_DEFAULT_CONTEXT: dict[str, int] = {
|
|
"llama-3": 8192,
|
|
"llama-2": 4096,
|
|
"mistral": 32768,
|
|
"mixtral": 32768,
|
|
"qwen2": 32768,
|
|
"qwen3": 32768,
|
|
"phi-3": 4096,
|
|
"phi-4": 16384,
|
|
"gemma": 8192,
|
|
"gemma-2": 8192,
|
|
"deepseek": 4096,
|
|
"codellama": 16384,
|
|
}
|
|
|
|
# v0.49.0 Part D — Llama 3.1 NTK-aware defaults.
|
|
LLAMA3_DEFAULT_SCALE_FACTOR: float = 8.0
|
|
LLAMA3_DEFAULT_LOW_FREQ_FACTOR: float = 1.0
|
|
LLAMA3_DEFAULT_HIGH_FREQ_FACTOR: float = 4.0
|
|
LLAMA3_DEFAULT_OLD_CONTEXT_LEN: int = 8192
|
|
|
|
|
|
def get_model_default_context(model_name: str) -> int:
|
|
"""Estimate the default context length for a model based on its name."""
|
|
model_lower = model_name.lower()
|
|
for family, ctx_len in MODEL_DEFAULT_CONTEXT.items():
|
|
if family in model_lower:
|
|
return ctx_len
|
|
# Conservative default for unknown models.
|
|
return 4096
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# YaRN math kernel (v0.49.0 Part A)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _finite_positive(value: Any, name: str) -> float:
|
|
"""Reject bool/NaN/Inf/<=0 with a typed error (mirrors v0.41.0 Part B)."""
|
|
if isinstance(value, bool):
|
|
raise ValueError(f"{name} must be a real number, not bool")
|
|
try:
|
|
v = float(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError(f"{name} must be a real number") from exc
|
|
if not math.isfinite(v):
|
|
raise ValueError(f"{name} must be finite (got {v!r})")
|
|
if v <= 0.0:
|
|
raise ValueError(f"{name} must be > 0 (got {v!r})")
|
|
return v
|
|
|
|
|
|
def yarn_find_correction_dim(
|
|
*,
|
|
num_rotations: float,
|
|
dim: int,
|
|
base: float = 10000.0,
|
|
max_position_embeddings: int,
|
|
) -> float:
|
|
"""Inverse of the rotation count for a given embedding dimension index.
|
|
|
|
Mirrors the upstream YaRN reference implementation
|
|
(https://arxiv.org/abs/2309.00071 §3.4):
|
|
|
|
dim_idx = (dim * ln(L / (n * 2π))) / (2 * ln(base))
|
|
|
|
Args:
|
|
num_rotations: Rotation count cutoff (``beta_fast`` or ``beta_slow``).
|
|
dim: Embedding dimension per head.
|
|
base: RoPE base (``theta``); typically 10_000 for Llama / Mistral.
|
|
max_position_embeddings: Original model context length.
|
|
"""
|
|
num_rotations = _finite_positive(num_rotations, "num_rotations")
|
|
if isinstance(dim, bool) or not isinstance(dim, int) or dim <= 0:
|
|
raise ValueError(f"dim must be a positive int (got {dim!r})")
|
|
base = _finite_positive(base, "base")
|
|
if (
|
|
isinstance(max_position_embeddings, bool)
|
|
or not isinstance(max_position_embeddings, int)
|
|
or max_position_embeddings <= 0
|
|
):
|
|
raise ValueError(
|
|
f"max_position_embeddings must be a positive int (got {max_position_embeddings!r})"
|
|
)
|
|
return (dim * math.log(max_position_embeddings / (num_rotations * 2.0 * math.pi))) / (
|
|
2.0 * math.log(base)
|
|
)
|
|
|
|
|
|
def yarn_find_correction_range(
|
|
*,
|
|
beta_fast: float,
|
|
beta_slow: float,
|
|
dim: int,
|
|
base: float = 10000.0,
|
|
max_position_embeddings: int,
|
|
) -> tuple[int, int]:
|
|
"""Return clamped ``(low, high)`` dim indices for YaRN piecewise interpolation."""
|
|
low = math.floor(
|
|
yarn_find_correction_dim(
|
|
num_rotations=beta_fast,
|
|
dim=dim,
|
|
base=base,
|
|
max_position_embeddings=max_position_embeddings,
|
|
)
|
|
)
|
|
high = math.ceil(
|
|
yarn_find_correction_dim(
|
|
num_rotations=beta_slow,
|
|
dim=dim,
|
|
base=base,
|
|
max_position_embeddings=max_position_embeddings,
|
|
)
|
|
)
|
|
half = dim // 2
|
|
low = max(0, min(low, half))
|
|
high = max(0, min(high, half))
|
|
if high <= low:
|
|
# Disambiguate so the ramp-mask denominator is non-zero.
|
|
high = min(low + 1, half)
|
|
if high <= low:
|
|
low = max(0, high - 1)
|
|
return low, high
|
|
|
|
|
|
def yarn_linear_ramp_mask(*, low: int, high: int, dim: int) -> list[float]:
|
|
"""Return the YaRN piecewise-linear ramp mask of length ``dim``.
|
|
|
|
Each element is clamped to ``[0, 1]``; ``low == high`` is auto-disambiguated
|
|
to avoid a div-by-zero in the upstream formula
|
|
``(idx - low) / (high - low)``.
|
|
"""
|
|
if isinstance(dim, bool) or not isinstance(dim, int) or dim <= 0:
|
|
raise ValueError(f"dim must be a positive int (got {dim!r})")
|
|
if isinstance(low, bool) or isinstance(high, bool):
|
|
raise ValueError("low/high must not be bool")
|
|
if not isinstance(low, int) or not isinstance(high, int):
|
|
raise ValueError("low/high must be ints")
|
|
if low < 0 or high < 0:
|
|
raise ValueError("low/high must be non-negative")
|
|
if high <= low:
|
|
high = low + 1
|
|
denom = float(high - low)
|
|
return [max(0.0, min(1.0, (idx - low) / denom)) for idx in range(dim)]
|
|
|
|
|
|
def yarn_get_mscale(factor: float) -> float:
|
|
"""Return the YaRN attention temperature multiplier ``m_scale``.
|
|
|
|
Per the YaRN paper §3.5, ``m_scale = 0.1 * ln(s) + 1`` for ``s > 1``;
|
|
factors at or below 1 produce no scaling, so we clamp to ``1.0``.
|
|
"""
|
|
if isinstance(factor, bool):
|
|
raise ValueError("factor must be a real number, not bool")
|
|
try:
|
|
f = float(factor)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("factor must be a real number") from exc
|
|
if not math.isfinite(f):
|
|
# Stay consistent with `_finite_positive` — explicit rejection over a
|
|
# silent identity fallback (python-review LOW finding).
|
|
raise ValueError(f"factor must be finite (got {f!r})")
|
|
if f <= 1.0:
|
|
return 1.0
|
|
return 0.1 * math.log(f) + 1.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Llama 3.1 NTK-aware kernel (v0.49.0 Part D)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def scale_inv_freq_llama3(
|
|
*,
|
|
inv_freq: float,
|
|
scale_factor: float = LLAMA3_DEFAULT_SCALE_FACTOR,
|
|
low_freq_factor: float = LLAMA3_DEFAULT_LOW_FREQ_FACTOR,
|
|
high_freq_factor: float = LLAMA3_DEFAULT_HIGH_FREQ_FACTOR,
|
|
old_context_len: int = LLAMA3_DEFAULT_OLD_CONTEXT_LEN,
|
|
) -> float:
|
|
"""Apply Llama 3.1 NTK-aware scaling to a single ``inv_freq`` value.
|
|
|
|
Replicates Unsloth ``models/llama.py:1853`` / HF transformers ``modeling_rope_utils``
|
|
``_compute_llama3_parameters``. The function classifies the frequency by its
|
|
wavelength relative to ``old_context_len``:
|
|
|
|
* wavelength < ``high_freq_wavelen`` → unchanged
|
|
* wavelength > ``low_freq_wavelen`` → divided by ``scale_factor``
|
|
* in between → smooth linear blend
|
|
"""
|
|
for _name, _val in (
|
|
("inv_freq", inv_freq),
|
|
("scale_factor", scale_factor),
|
|
("low_freq_factor", low_freq_factor),
|
|
("high_freq_factor", high_freq_factor),
|
|
):
|
|
if isinstance(_val, bool):
|
|
raise ValueError(f"{_name} must be a real number, not bool")
|
|
if not isinstance(_val, (int, float)) or not math.isfinite(float(_val)):
|
|
raise ValueError(f"{_name} must be finite (got {_val!r})")
|
|
if isinstance(old_context_len, bool) or not isinstance(old_context_len, int):
|
|
raise ValueError(f"old_context_len must be an int (got {old_context_len!r})")
|
|
if scale_factor <= 1.0:
|
|
raise ValueError(f"scale_factor must be > 1 (got {scale_factor!r})")
|
|
if high_freq_factor <= low_freq_factor:
|
|
raise ValueError(
|
|
f"high_freq_factor ({high_freq_factor}) must be greater than "
|
|
f"low_freq_factor ({low_freq_factor})"
|
|
)
|
|
if old_context_len <= 0:
|
|
raise ValueError(f"old_context_len must be > 0 (got {old_context_len})")
|
|
|
|
inv_freq_f = float(inv_freq)
|
|
if inv_freq_f <= 0.0:
|
|
return inv_freq_f
|
|
|
|
wavelen = 2.0 * math.pi / inv_freq_f
|
|
low_freq_wavelen = old_context_len / low_freq_factor
|
|
high_freq_wavelen = old_context_len / high_freq_factor
|
|
|
|
if wavelen < high_freq_wavelen:
|
|
return inv_freq_f
|
|
if wavelen > low_freq_wavelen:
|
|
return inv_freq_f / scale_factor
|
|
# Smooth blend between the two regions.
|
|
smooth = (old_context_len / wavelen - low_freq_factor) / (
|
|
high_freq_factor - low_freq_factor
|
|
)
|
|
return (1.0 - smooth) * (inv_freq_f / scale_factor) + smooth * inv_freq_f
|
|
|
|
|
|
def detect_llama3_rope_in_config(config: Mapping[str, Any]) -> bool:
|
|
"""Return True if the HF-style config dict carries a Llama 3.1 RoPE block.
|
|
|
|
Both ``rope_scaling.type`` and the newer ``rope_scaling.rope_type`` keys are
|
|
accepted (transformers >=4.43 uses ``rope_type``).
|
|
"""
|
|
if not isinstance(config, Mapping):
|
|
raise TypeError(f"config must be a Mapping (got {type(config).__name__})")
|
|
rope = config.get("rope_scaling")
|
|
if not isinstance(rope, Mapping):
|
|
return False
|
|
# Explicit ``is None`` check (mirrors v0.40.6 review-fix policy) — prevents
|
|
# a falsy-but-set ``type`` from silently falling through to ``rope_type``.
|
|
type_value = rope.get("type") if rope.get("type") is not None else rope.get("rope_type")
|
|
return isinstance(type_value, str) and type_value.lower() == "llama3"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public config emission
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def get_rope_scaling_config(
|
|
scaling_type: str,
|
|
target_length: float,
|
|
original_length: int,
|
|
*,
|
|
yarn_factor: float | None = None,
|
|
yarn_attn_factor: float | None = None,
|
|
yarn_beta_fast: int | None = None,
|
|
yarn_beta_slow: int | None = None,
|
|
llama3_scale_factor: float = LLAMA3_DEFAULT_SCALE_FACTOR,
|
|
llama3_low_freq_factor: float = LLAMA3_DEFAULT_LOW_FREQ_FACTOR,
|
|
llama3_high_freq_factor: float = LLAMA3_DEFAULT_HIGH_FREQ_FACTOR,
|
|
llama3_old_context_len: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Build the ``rope_scaling`` dict consumed by HF model configs.
|
|
|
|
The YaRN tunables are emitted only when ``scaling_type='yarn'``; the Llama
|
|
3.1 tunables are emitted only when ``scaling_type='llama3'``.
|
|
"""
|
|
if scaling_type not in ROPE_SCALING_TYPES:
|
|
raise ValueError(
|
|
f"Unknown RoPE scaling type: {scaling_type}. "
|
|
f"Options: {', '.join(ROPE_SCALING_TYPES)}"
|
|
)
|
|
|
|
# Security review fix — validate numeric inputs at the public boundary.
|
|
# The schema layer already guards Pydantic-loaded values; this protects
|
|
# direct callers from emitting ``{"factor": NaN}`` into HF model configs.
|
|
if isinstance(target_length, bool) or isinstance(original_length, bool):
|
|
raise ValueError("target_length / original_length must not be bool")
|
|
if not isinstance(target_length, (int, float)) or not math.isfinite(float(target_length)):
|
|
raise ValueError(f"target_length must be a finite number (got {target_length!r})")
|
|
if not isinstance(original_length, int) or original_length <= 0:
|
|
raise ValueError(f"original_length must be a positive int (got {original_length!r})")
|
|
if yarn_factor is not None:
|
|
yarn_factor = _finite_positive(yarn_factor, "yarn_factor")
|
|
|
|
# If target_length looks like a scaling factor (small number > 1.0 but < 64),
|
|
# treat it as a multiplier; values >= 64 are token counts.
|
|
if target_length < 64 and target_length > 1.0:
|
|
factor = float(target_length)
|
|
else:
|
|
factor = target_length / original_length
|
|
|
|
if factor <= 1.0:
|
|
return {}
|
|
|
|
factor = float(factor)
|
|
if scaling_type == "linear":
|
|
return {"type": "linear", "factor": factor}
|
|
if scaling_type == "dynamic":
|
|
return {"type": "dynamic", "factor": factor}
|
|
if scaling_type == "yarn":
|
|
cfg: dict[str, Any] = {
|
|
"type": "yarn",
|
|
"factor": yarn_factor if yarn_factor is not None else factor,
|
|
"original_max_position_embeddings": original_length,
|
|
}
|
|
if yarn_attn_factor is not None:
|
|
cfg["attention_factor"] = float(yarn_attn_factor)
|
|
if yarn_beta_fast is not None:
|
|
cfg["beta_fast"] = int(yarn_beta_fast)
|
|
if yarn_beta_slow is not None:
|
|
cfg["beta_slow"] = int(yarn_beta_slow)
|
|
return cfg
|
|
if scaling_type == "longrope":
|
|
return {
|
|
"type": "longrope",
|
|
"factor": factor,
|
|
"original_max_position_embeddings": original_length,
|
|
}
|
|
# scaling_type == "llama3" — defer to the caller-supplied original_length
|
|
# unless an explicit Llama 3.1-specific override is provided.
|
|
return {
|
|
"type": "llama3",
|
|
"factor": factor,
|
|
"original_max_position_embeddings": (
|
|
llama3_old_context_len if llama3_old_context_len is not None else original_length
|
|
),
|
|
"low_freq_factor": float(llama3_low_freq_factor),
|
|
"high_freq_factor": float(llama3_high_freq_factor),
|
|
}
|
|
|
|
|
|
def apply_long_context_config(
|
|
model_config,
|
|
target_length: int,
|
|
rope_scaling_type: str = "dynamic",
|
|
model_name: str = "",
|
|
) -> dict | None:
|
|
"""Apply long-context configuration to a model config object."""
|
|
original_length = getattr(
|
|
model_config,
|
|
"max_position_embeddings",
|
|
get_model_default_context(model_name),
|
|
)
|
|
if target_length <= original_length:
|
|
return None
|
|
rope_config = get_rope_scaling_config(
|
|
scaling_type=rope_scaling_type,
|
|
target_length=target_length,
|
|
original_length=original_length,
|
|
)
|
|
if not rope_config:
|
|
return None
|
|
model_config.rope_scaling = rope_config
|
|
model_config.max_position_embeddings = target_length
|
|
return rope_config
|
|
|
|
|
|
def validate_long_context_config(
|
|
max_length: int,
|
|
rope_scaling_type: str | None,
|
|
use_gradient_checkpointing: bool,
|
|
) -> list[str]:
|
|
"""Validate long-context configuration."""
|
|
errors: list[str] = []
|
|
if rope_scaling_type and rope_scaling_type not in ROPE_SCALING_TYPES:
|
|
errors.append(
|
|
f"Unknown RoPE scaling type: {rope_scaling_type}. "
|
|
f"Options: {', '.join(ROPE_SCALING_TYPES)}"
|
|
)
|
|
if max_length >= 65536 and not use_gradient_checkpointing:
|
|
errors.append(
|
|
f"Training with max_length={max_length} without gradient checkpointing "
|
|
"will likely cause OOM. Set gradient_checkpointing: true in config."
|
|
)
|
|
return errors
|