mirror of https://github.com/razor-ai/soup.git
112 lines
4.5 KiB
Python
112 lines
4.5 KiB
Python
"""LongLoRA S² shifted-sparse attention support (v0.49.0 Part C).
|
|
|
|
Schema-level helpers only — the live forward-override patch that mirrors
|
|
LlamaFactory ``model/model_utils/longlora.py`` is deferred to v0.49.1
|
|
(stub-then-live pattern, mirrors v0.27.0 MII / v0.37.0 multipack / v0.41.0
|
|
LLaMA Pro / v0.46.0 Quant-Lobotomy auto-measure).
|
|
|
|
LongLoRA's S² (shifted-sparse) attention groups tokens into local windows and
|
|
shifts half the heads by ``window_size // 2`` so adjacent groups can exchange
|
|
information. The implementation in LlamaFactory replaces
|
|
``LlamaAttention.forward`` with a custom kernel that materialises this pattern
|
|
via two cheap masked attentions; FlashAttention v3's custom-mask API and
|
|
Ring-FlashAttention's sequence-parallel attention are not yet compatible.
|
|
|
|
Architecture allowlist: Llama 3.x only initially (per the upstream
|
|
implementation). Extend per request.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import Any
|
|
|
|
# Word-boundary regex matching the v0.39.0 ``is_gemma4_model`` / v0.44.0
|
|
# ``is_llama4_model`` policy — substring ``"llama"`` inside an unrelated identifier
|
|
# (e.g. ``my-llama-style-finetune``) must NOT match silently. We accept ``llama``
|
|
# with optional version digit suffix and an optional ``code-`` prefix for the
|
|
# ``codellama`` family.
|
|
_LLAMA_REGEX = re.compile(r"(?:^|[^a-z0-9])(?:code)?-?llama(?:-?\d+(?:\.\d+)?)?(?:[^a-z0-9]|$)")
|
|
|
|
_MAX_MODEL_NAME_LEN = 512
|
|
|
|
|
|
def is_llama_model(model_name: str) -> bool:
|
|
"""Return True if ``model_name`` belongs to the Llama / CodeLlama family.
|
|
|
|
Args:
|
|
model_name: HuggingFace model id or local path.
|
|
|
|
Raises:
|
|
TypeError: If ``model_name`` is not a string.
|
|
ValueError: If ``model_name`` contains a null byte.
|
|
"""
|
|
if not isinstance(model_name, str):
|
|
raise TypeError(f"model_name must be a string, got {type(model_name).__name__}")
|
|
if "\x00" in model_name:
|
|
raise ValueError("model_name must not contain null bytes")
|
|
if len(model_name) > _MAX_MODEL_NAME_LEN:
|
|
return False
|
|
return _LLAMA_REGEX.search(model_name.lower()) is not None
|
|
|
|
|
|
def validate_longlora_compat(
|
|
*,
|
|
model_name: str,
|
|
task: str,
|
|
backend: str,
|
|
use_ring_attention: bool,
|
|
) -> None:
|
|
"""Raise ``ValueError`` if the LongLoRA prerequisites are not satisfied.
|
|
|
|
LongLoRA requires:
|
|
* ``task='sft'`` (preference / RL trainers have a different forward path)
|
|
* ``backend='transformers'`` (Unsloth has its own attention; MLX has none)
|
|
* Llama-family base model (initial architecture allowlist)
|
|
* ``use_ring_attention=False`` (mutually exclusive custom attention kernel)
|
|
"""
|
|
if backend == "mlx":
|
|
# Distinct error message per project policy (matches v0.34.0 review-fix
|
|
# convention of naming the specific incompatible backend).
|
|
raise ValueError(
|
|
"LongLoRA is not supported on the mlx backend (the S^2 forward "
|
|
"override is HF-Transformers specific). Use backend='transformers' "
|
|
"or set use_longlora=false."
|
|
)
|
|
if backend != "transformers":
|
|
raise ValueError(
|
|
f"LongLoRA requires backend='transformers' (got backend={backend!r})."
|
|
)
|
|
if task != "sft":
|
|
raise ValueError(
|
|
f"LongLoRA is currently restricted to task='sft' (got task={task!r}). "
|
|
"Multi-trainer expansion is planned for a future release."
|
|
)
|
|
if not is_llama_model(model_name):
|
|
raise ValueError(
|
|
"LongLoRA architecture allowlist currently covers only the Llama "
|
|
f"family (got base={model_name!r}). Open a feature request to add "
|
|
"another architecture."
|
|
)
|
|
if use_ring_attention:
|
|
raise ValueError(
|
|
"LongLoRA is incompatible with use_ring_attention (both rewrite the "
|
|
"attention kernel — pick one)."
|
|
)
|
|
|
|
|
|
def apply_longlora_forward_override(model: Any) -> None:
|
|
"""Install the LongLoRA S^2 ``LlamaAttention.forward`` override.
|
|
|
|
Deferred to v0.49.1 (matches the stub-then-live pattern used by v0.27.0
|
|
MII / v0.37.0 multipack / v0.41.0 LLaMA Pro / v0.46.0 Quant-Lobotomy).
|
|
The v0.49.0 release ships the schema gate (so ``soup train`` fails fast
|
|
when LongLoRA is misconfigured) and the kernel math; the live forward
|
|
monkeypatch + tests on the real model graph land in v0.49.1.
|
|
"""
|
|
raise NotImplementedError(
|
|
"LongLoRA S^2 forward override lands in v0.49.1 — the v0.49.0 release "
|
|
"ships the schema gate (so misconfigured runs fail fast) and the math "
|
|
"kernel only."
|
|
)
|