Merge remote-tracking branch 'origin/main' into ig_merge
# Conflicts: # scripts/scan_packages_baseline.json
This commit is contained in:
Daniel Han
commit
b9ebfe089b
46 changed files with 8229 additions and 2289 deletions
122
tests/python/test_fast_sentence_transformer_embedding_parity.py
Normal file
122
tests/python/test_fast_sentence_transformer_embedding_parity.py
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team.
|
||||
"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text
|
||||
like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a
|
||||
"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building
|
||||
via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings;
|
||||
`_create_transformer_module` uses `Transformer.load(...)` instead.
|
||||
|
||||
Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST
|
||||
is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity,
|
||||
opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_transformer_load_signature_supports_unsloth_kwargs():
|
||||
"""Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs
|
||||
the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back
|
||||
to Transformer(...) there, so mirror that gate and skip."""
|
||||
models = pytest.importorskip("sentence_transformers.models")
|
||||
load = getattr(models.Transformer, "load", None)
|
||||
assert callable(load), (
|
||||
"sentence_transformers Transformer.load is missing; the #6881 fix in "
|
||||
"unsloth.models.sentence_transformer._create_transformer_module depends on it."
|
||||
)
|
||||
params = inspect.signature(load).parameters
|
||||
accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
|
||||
# Mirror _create_transformer_module's hub_capable gate.
|
||||
hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision"))
|
||||
if not hub_capable:
|
||||
pytest.skip(
|
||||
"legacy Transformer.load(input_path); production path falls back to Transformer(...)"
|
||||
)
|
||||
unsupported = [
|
||||
k
|
||||
for k in ("token", "cache_folder", "revision", "trust_remote_code")
|
||||
if not (accepts_var_kw or k in params)
|
||||
]
|
||||
assert not unsupported, (
|
||||
f"installed sentence_transformers Transformer.load no longer accepts {unsupported} "
|
||||
f"and has no **kwargs; update _create_transformer_module (#6881) before it silently "
|
||||
f"falls back to Transformer(...)."
|
||||
)
|
||||
|
||||
|
||||
def _probe_texts():
|
||||
return [
|
||||
"roasted chickpeas in 20 kg bags",
|
||||
"The capital of France is Paris.",
|
||||
"A fast brown fox jumps over the lazy dog.",
|
||||
"recette de tarte aux pommes traditionnelle",
|
||||
]
|
||||
|
||||
|
||||
def test_fast_sentence_transformer_matches_stock_st():
|
||||
"""End-to-end: FastSentenceTransformer embeddings and tokenization must match a
|
||||
stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and
|
||||
GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners."""
|
||||
model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL")
|
||||
if not model_id:
|
||||
pytest.skip(
|
||||
"set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model "
|
||||
"(HF id or local path) to run the #6881 parity test"
|
||||
)
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner")
|
||||
np = pytest.importorskip("numpy")
|
||||
pytest.importorskip("sentence_transformers")
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
device = "cuda"
|
||||
# Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native
|
||||
# embedders such as EmbeddingGemma (Gemma3), which would mask real parity.
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
texts = _probe_texts()
|
||||
max_seq_length = 256
|
||||
|
||||
# Control FIRST, before importing unsloth, so its global import patches never
|
||||
# touch the stock reference (mirrors the issue's "restart runtime" repro).
|
||||
ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype})
|
||||
ctrl.max_seq_length = max_seq_length
|
||||
ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist()
|
||||
ctrl_emb = np.asarray(
|
||||
ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
|
||||
)
|
||||
|
||||
import unsloth # noqa: F401
|
||||
from unsloth import FastSentenceTransformer
|
||||
|
||||
fast = FastSentenceTransformer.from_pretrained(
|
||||
model_id,
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
load_in_4bit = False,
|
||||
load_in_16bit = True,
|
||||
)
|
||||
fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist()
|
||||
fast_emb = np.asarray(
|
||||
fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
|
||||
)
|
||||
|
||||
# Identical tokenization = no chat-template wrapping slipped in (the #6881 defect).
|
||||
assert fast_ids == ctrl_ids, (
|
||||
f"tokenization diverged (chat-template wrapping regressed?):\n"
|
||||
f" stock: {ctrl_ids}\n fast: {fast_ids}"
|
||||
)
|
||||
|
||||
cos = (ctrl_emb * fast_emb).sum(1) / (
|
||||
np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1)
|
||||
)
|
||||
assert float(cos.min()) > 0.99, (
|
||||
f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 "
|
||||
f"(per-text {[round(float(c), 5) for c in cos]})"
|
||||
)
|
||||
|
|
@ -14,6 +14,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/
|
|||
_REPO_ROOT = _TESTS_DIR.parent # unsloth/
|
||||
_INSTALL_SH = _REPO_ROOT / "install.sh"
|
||||
_INSTALL_PS1 = _REPO_ROOT / "install.ps1"
|
||||
_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1"
|
||||
_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
|
||||
|
||||
|
||||
|
|
@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged:
|
|||
assert '"torch>=2.4,<2.11.0"' in self._ps1
|
||||
|
||||
|
||||
class TestInstallPs1UvDefaultIndex:
|
||||
"""Installer-managed torch indexes must override inherited uv defaults."""
|
||||
|
||||
_ps1 = _read(_INSTALL_PS1)
|
||||
|
||||
def test_torch_installs_use_default_index(self):
|
||||
assert "--default-index $TorchIndexUrl" in self._ps1
|
||||
assert "--default-index $ROCmIndexUrl" in self._ps1
|
||||
|
||||
def test_torch_installs_do_not_use_deprecated_index_url(self):
|
||||
assert "--index-url $TorchIndexUrl" not in self._ps1
|
||||
assert "--index-url $ROCmIndexUrl" not in self._ps1
|
||||
|
||||
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
|
||||
# Extra-index vars outrank --default-index, so pinned installs must clear them.
|
||||
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
|
||||
assert var in self._ps1
|
||||
assert 'Remove-Item "Env:$n"' in self._ps1
|
||||
|
||||
|
||||
class TestSetupPs1FastInstallIndex:
|
||||
"""setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning."""
|
||||
|
||||
_ps1 = _read(_SETUP_PS1)
|
||||
|
||||
def test_fast_install_clears_all_uv_index_env_vars(self):
|
||||
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
|
||||
assert var in self._ps1
|
||||
# Must truly remove the vars (child sees no value), not set them empty.
|
||||
assert 'Remove-Item "Env:$n"' in self._ps1
|
||||
|
||||
|
||||
class TestInstallShUvDefaultIndex:
|
||||
"""Linux/Mac installer torch indexes must override inherited uv defaults."""
|
||||
|
||||
_sh = _read(_INSTALL_SH)
|
||||
|
||||
def test_torch_installs_use_default_index(self):
|
||||
assert '--default-index "$TORCH_INDEX_URL"' in self._sh
|
||||
|
||||
def test_torch_installs_do_not_use_deprecated_index_url(self):
|
||||
assert '--index-url "$TORCH_INDEX_URL"' not in self._sh
|
||||
|
||||
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
|
||||
# --default-index installs run with all uv index env vars unset via `env -u`.
|
||||
assert (
|
||||
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
|
||||
)
|
||||
|
||||
|
||||
# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
|
||||
class TestTorchConstraintShell:
|
||||
"""Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""
|
||||
|
|
|
|||
|
|
@ -312,7 +312,7 @@ if [ "$SKIP_TORCH" = true ]; then
|
|||
else
|
||||
echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
|
||||
uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
--default-index "$TORCH_INDEX_URL"
|
||||
fi
|
||||
TORCH_EOF
|
||||
|
||||
|
|
|
|||
|
|
@ -2684,68 +2684,6 @@ class TestBlackwellCuda124Exclusion:
|
|||
assert kept == [cpu]
|
||||
|
||||
|
||||
# N.1c3. direct_linux_release_plan -- no silent CPU on NVIDIA hosts
|
||||
|
||||
|
||||
class TestDirectLinuxNvidiaCpuGate:
|
||||
"""A linux-cpu-only release on an NVIDIA host must raise (caller walks back to a usable CUDA line), not silently CPU-install. CPU-only hosts keep the CPU bundle."""
|
||||
|
||||
def _bundle_cpu_only(self):
|
||||
return make_release(
|
||||
[
|
||||
make_artifact(
|
||||
"llama-b8508-bin-ubuntu-x64.tar.gz",
|
||||
install_kind = "linux-cpu",
|
||||
runtime_line = None,
|
||||
coverage_class = None,
|
||||
supported_sms = [],
|
||||
min_sm = None,
|
||||
max_sm = None,
|
||||
bundle_profile = None,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def _patch(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
INSTALL_LLAMA_PREBUILT,
|
||||
"parse_direct_linux_release_bundle",
|
||||
lambda repo, release: self._bundle_cpu_only(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
INSTALL_LLAMA_PREBUILT,
|
||||
"detect_torch_cuda_runtime_preference",
|
||||
lambda host: CudaRuntimePreference(runtime_line = None, selection_log = []),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
INSTALL_LLAMA_PREBUILT,
|
||||
"detected_linux_runtime_lines",
|
||||
lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
|
||||
)
|
||||
|
||||
def test_nvidia_host_without_cuda_line_raises_for_walkback(self, monkeypatch):
|
||||
self._patch(monkeypatch)
|
||||
host = make_host(driver_cuda_version = (13, 1), compute_caps = ["100"])
|
||||
with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt"):
|
||||
INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
|
||||
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
|
||||
)
|
||||
|
||||
def test_cpu_host_still_gets_cpu_bundle(self, monkeypatch):
|
||||
self._patch(monkeypatch)
|
||||
host = make_host(
|
||||
nvidia_smi = None,
|
||||
driver_cuda_version = None,
|
||||
compute_caps = [],
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
)
|
||||
plan = INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
|
||||
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
|
||||
)
|
||||
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
|
||||
|
||||
|
||||
class TestLinuxPublishedAttemptsNvidiaCpuGate:
|
||||
"""Live fork-manifest path: an NVIDIA host whose CUDA selection finds nothing gets an empty attempt list (source-builds with CUDA), not the manifest CPU bundle. CPU-only hosts still get the CPU bundle."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1296,16 +1296,92 @@ with sync_playwright() as p:
|
|||
# still abort or interrupt this navigation, so the field wait below is the
|
||||
# final confirmation that we reached /login.
|
||||
_tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
|
||||
try:
|
||||
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
|
||||
except Exception as exc:
|
||||
if not any(t in str(exc) for t in _tolerated_nav):
|
||||
raise
|
||||
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
|
||||
pw_field = page.locator("#password")
|
||||
pw_field.wait_for(state = "visible", timeout = 60_000)
|
||||
pw_field.fill(NEW2)
|
||||
page.locator('button[type="submit"]').click()
|
||||
# A slow CI runner can make this re-login navigation time out even with the
|
||||
# server healthy, so retry the whole goto/wait/fill/submit sequence (mirrors
|
||||
# the change-password retry above). wait_for_health is a diagnostic pre-gate.
|
||||
wait_for_health(BASE, timeout = 30.0, info = info)
|
||||
relogin_err: Exception | None = None
|
||||
for _relogin_attempt in range(3):
|
||||
try:
|
||||
try:
|
||||
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
|
||||
except Exception as exc:
|
||||
if not any(t in str(exc) for t in _tolerated_nav):
|
||||
raise
|
||||
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
|
||||
pw_field = page.locator("#password")
|
||||
pw_field.wait_for(state = "visible", timeout = 60_000)
|
||||
pw_field.fill(NEW2)
|
||||
# Wait on the login POST so a transient 4xx/5xx is caught and retried
|
||||
# here, not swallowed until the out-of-loop composer wait.
|
||||
status, _ = click_and_wait_for_response(
|
||||
page,
|
||||
url_substr = "/api/auth/login",
|
||||
method = "POST",
|
||||
do_click = lambda: page.locator('button[type="submit"]').click(),
|
||||
timeout_ms = 30_000,
|
||||
info = lambda m: print(f"[ui] {m}", flush = True),
|
||||
)
|
||||
if status is not None and status >= 400:
|
||||
raise AssertionError(
|
||||
f"login POST returned {status}; see console_errors={console_errors[:1]!r}"
|
||||
)
|
||||
relogin_err = None
|
||||
break
|
||||
except Exception as e:
|
||||
relogin_err = e
|
||||
try:
|
||||
cur_url = page.url
|
||||
except Exception:
|
||||
cur_url = "<page closed>"
|
||||
print(
|
||||
f"[ui] re-login attempt {_relogin_attempt + 1} failed: "
|
||||
f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; "
|
||||
f"page_errors={len(page_errors)} console_errors={len(console_errors)}",
|
||||
flush = True,
|
||||
)
|
||||
if console_errors:
|
||||
print(
|
||||
f"[ui] first console.error: {console_errors[0][:200]!r}",
|
||||
flush = True,
|
||||
)
|
||||
if page_errors:
|
||||
print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True)
|
||||
try:
|
||||
shoot(f"18-relogin-attempt-{_relogin_attempt + 1}-fail")
|
||||
except Exception:
|
||||
pass
|
||||
if _relogin_attempt < 2:
|
||||
# ERR_NO_BUFFER_SPACE needs the OS to recover socket
|
||||
# buffers; back off 5s then 15s before retrying.
|
||||
if "ERR_NO_BUFFER_SPACE" in str(e):
|
||||
backoff_s = 5 if _relogin_attempt == 0 else 15
|
||||
print(
|
||||
f"[ui] ENOBUFS detected; sleeping {backoff_s}s "
|
||||
f"before retry to let OS recover socket buffers...",
|
||||
flush = True,
|
||||
)
|
||||
time.sleep(backoff_s)
|
||||
# Replace the page if it died; otherwise next iteration's
|
||||
# page.goto() handles the reload.
|
||||
old_page = page
|
||||
page = recover_or_replace_page(
|
||||
page,
|
||||
ctx,
|
||||
default_timeout_ms = 60_000,
|
||||
info = lambda m: print(f"[ui] recovery: {m}", flush = True),
|
||||
)
|
||||
# A freshly created replacement page loses the pageerror/console
|
||||
# listeners; re-attach so error tracking survives recovery.
|
||||
if page is not old_page:
|
||||
page.on("pageerror", lambda e: page_errors.append(str(e)))
|
||||
page.on("console", _on_console)
|
||||
if relogin_err is not None:
|
||||
raise relogin_err
|
||||
# Composer mount confirms the rotated session is authenticated. Kept OUTSIDE the
|
||||
# retry: the loop breaks right after submit, so we never re-goto /login once login
|
||||
# has set tokens -- that would hit the guest guard, redirect to /chat, and make a
|
||||
# merely-slow composer look like a broken login.
|
||||
composer = page.locator('textarea[aria-label="Message input"]')
|
||||
composer.wait_for(state = "visible", timeout = 60_000)
|
||||
shoot("18-relogin-with-NEW2")
|
||||
|
|
|
|||
63
tests/test_fast_gemv_dispatch.py
Normal file
63
tests/test_fast_gemv_dispatch.py
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is
|
||||
already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the
|
||||
bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up
|
||||
# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection
|
||||
# errors instead of producing a skip. Any other import error still surfaces as a failure.
|
||||
pytest.importorskip("bitsandbytes")
|
||||
|
||||
import unsloth # noqa: F401 (sets UNSLOTH_IS_PRESENT before transformers)
|
||||
from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES
|
||||
|
||||
_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None
|
||||
|
||||
|
||||
def _proj(weight, weight_scale = None):
|
||||
proj = SimpleNamespace(weight = weight, bias = None, merged = False)
|
||||
if weight_scale is not None:
|
||||
proj.weight_scale = weight_scale
|
||||
return proj
|
||||
|
||||
|
||||
def test_bf16_weight_scale_not_used_as_quant_state():
|
||||
"""A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None."""
|
||||
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2))
|
||||
W, W_quant = get_lora_parameters_bias(proj)[:2]
|
||||
assert W_quant is None
|
||||
|
||||
|
||||
def test_fp8_weight_keeps_scale():
|
||||
"""An actual fp8 weight still resolves its weight_scale as the quant state."""
|
||||
if _FP8 is None:
|
||||
pytest.skip("no float8 dtype in this torch build")
|
||||
scale = torch.rand(2, 2)
|
||||
proj = _proj(torch.randn(4, 4).to(_FP8), scale)
|
||||
W, W_quant = get_lora_parameters_bias(proj)[:2]
|
||||
assert W_quant is scale
|
||||
|
||||
|
||||
def test_plain_bf16_has_no_quant_state():
|
||||
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16))
|
||||
W, W_quant = get_lora_parameters_bias(proj)[:2]
|
||||
assert W_quant is None
|
||||
358
tests/test_fp8_restore_dropped_scale.py
Normal file
358
tests/test_fp8_restore_dropped_scale.py
Normal file
|
|
@ -0,0 +1,358 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Restoring dropped block-fp8 `weight_scale_inv` tensors on load (#6200).
|
||||
|
||||
Some block-scale fp8 checkpoints leave a Linear (e.g. `mlp.gate_proj`) unconverted, so its raw
|
||||
quantized values land in a plain bf16 weight and its `weight_scale_inv` is dropped, producing a
|
||||
garbage un-scaled weight. `_restore_dropped_fp8_scales` dequantizes such orphaned weights in place
|
||||
using the scale from the checkpoint. Runs offline on CPU with synthetic checkpoints.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from safetensors.torch import save_file
|
||||
|
||||
# Import unsloth first to set UNSLOTH_IS_PRESENT env var.
|
||||
import unsloth
|
||||
from unsloth.models.loader_utils import _restore_dropped_fp8_scales, _FP8_DTYPES
|
||||
|
||||
|
||||
_SHARD = "model-00001-of-00001.safetensors"
|
||||
_FP8 = _FP8_DTYPES[0] if _FP8_DTYPES else None
|
||||
|
||||
|
||||
def _write_checkpoint(
|
||||
path,
|
||||
tensors,
|
||||
filename = _SHARD,
|
||||
include_index = True,
|
||||
):
|
||||
save_file(tensors, os.path.join(path, filename))
|
||||
if include_index:
|
||||
weight_map = {name: filename for name in tensors}
|
||||
with open(os.path.join(path, "model.safetensors.index.json"), "w") as f:
|
||||
json.dump({"weight_map": weight_map}, f)
|
||||
|
||||
|
||||
def _fp8_config(block = (2, 2)):
|
||||
return SimpleNamespace(
|
||||
quantization_config = {
|
||||
"quant_method": "fp8",
|
||||
"weight_block_size": list(block),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _fp8_anchor():
|
||||
"""A module carrying a real fp8 weight, so the model looks like a genuine fp8 load."""
|
||||
m = nn.Linear(2, 2, bias = False)
|
||||
m.weight = nn.Parameter(torch.randn(2, 2).to(_FP8), requires_grad = False)
|
||||
return m
|
||||
|
||||
|
||||
def _bf16_linear(out_f, in_f, raw):
|
||||
m = nn.Linear(in_f, out_f, bias = False).to(torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
m.weight.copy_(raw)
|
||||
return m
|
||||
|
||||
|
||||
def _expand(scale, block, shape):
|
||||
bs0, bs1 = block
|
||||
expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
|
||||
return expanded[: shape[0], : shape[1]]
|
||||
|
||||
|
||||
def test_restore_dequantizes_orphaned_scale():
|
||||
"""A plain bf16 weight whose scale was dropped is dequantized in place."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
torch.manual_seed(0)
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(
|
||||
d,
|
||||
{
|
||||
"layer.weight": raw.to(torch.float32),
|
||||
"layer.weight_scale_inv": scale,
|
||||
},
|
||||
)
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_skips_already_fp8_weight():
|
||||
"""A correctly converted fp8 weight is skipped, never double-scaled."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
weight = torch.randn(4, 4).to(_FP8)
|
||||
before = weight.clone()
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.layer = nn.Linear(4, 4, bias = False)
|
||||
model.layer.weight = nn.Parameter(weight, requires_grad = False)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": torch.rand(2, 2)})
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 0 and skipped == 1
|
||||
assert torch.equal(model.layer.weight.data.float(), before.float())
|
||||
|
||||
|
||||
def test_skips_offloaded_meta_weight():
|
||||
"""A disk-offloaded layer (weight on the meta device) is skipped without error or restore."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = nn.Linear(4, 4, bias = False)
|
||||
# Simulate an offloaded weight living on the meta device.
|
||||
model.layer.weight = nn.Parameter(
|
||||
torch.empty(4, 4, dtype = torch.bfloat16, device = "meta"), requires_grad = False
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(
|
||||
d,
|
||||
{
|
||||
"layer.weight": raw.to(torch.float32),
|
||||
"layer.weight_scale_inv": scale,
|
||||
},
|
||||
)
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 0
|
||||
assert model.layer.weight.device.type == "meta"
|
||||
|
||||
|
||||
def test_noop_when_fully_dequantized():
|
||||
"""If the model has no fp8 weights at all (e.g. load_in_16bit dequantize), do not rescale."""
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.layer = _bf16_linear(4, 4, raw) # no fp8 anchor -> looks dequantized
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert (restored, skipped) == (0, 0)
|
||||
assert torch.equal(model.layer.weight.data, raw) # untouched
|
||||
|
||||
|
||||
def test_non_block_divisible_shape():
|
||||
"""Block scale is expanded then sliced to a non-divisible weight shape."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(3, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(3, 4, raw) # weight shape [3, 4]
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (3, 4))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_transposed_scale_layout():
|
||||
"""A scale stored in the transposed block grid is transposed before use."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 2, dtype = torch.bfloat16) # weight [4, 2] -> grid (2, 1)
|
||||
scale_correct = torch.rand(2, 1, dtype = torch.float32) + 0.1
|
||||
scale_stored = scale_correct.t().contiguous() # stored transposed as (1, 2)
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 2, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale_stored})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale_correct, (2, 2), (4, 2))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_single_file_checkpoint_without_index():
|
||||
"""Unsharded model.safetensors (no index) is still scanned for dropped scales."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(
|
||||
d, {"layer.weight_scale_inv": scale}, filename = "model.safetensors", include_index = False
|
||||
)
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_scalar_block_size_config():
|
||||
"""A scalar weight_block_size (not a list) is handled without error."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = SimpleNamespace(
|
||||
quantization_config = {"quant_method": "fp8", "weight_block_size": 2}
|
||||
)
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
|
||||
|
||||
def test_text_only_prefix_mapping():
|
||||
"""Checkpoint keys with a language_model prefix match the stripped text-only module names."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(2, 2, dtype = torch.bfloat16)
|
||||
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.model = nn.Module()
|
||||
model.model.gate_proj = _bf16_linear(2, 2, raw) # module lacks the language_model prefix
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
# checkpoint key carries the language_model wrapper the text-only load stripped
|
||||
_write_checkpoint(d, {"model.language_model.gate_proj.weight_scale_inv": scale})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
|
||||
assert torch.equal(model.model.gate_proj.weight.data, expected)
|
||||
|
||||
|
||||
def test_skips_variant_load():
|
||||
"""A variant load (variant="fp8") is skipped to avoid applying default-checkpoint scales."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
result = _restore_dropped_fp8_scales(model, d, local_files_only = True, variant = "fp8")
|
||||
assert result == (0, 0)
|
||||
assert torch.equal(model.layer.weight.data, raw) # untouched
|
||||
|
||||
|
||||
def test_vlm_language_model_model_alias():
|
||||
"""A checkpoint key language_model.model.* matches a model.language_model.* module."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(2, 2, dtype = torch.bfloat16)
|
||||
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.model = nn.Module()
|
||||
model.model.language_model = nn.Module()
|
||||
model.model.language_model.gate_proj = _bf16_linear(
|
||||
2, 2, raw
|
||||
) # -> model.language_model.gate_proj
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"language_model.model.gate_proj.weight_scale_inv": scale})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
|
||||
assert torch.equal(model.model.language_model.gate_proj.weight.data, expected)
|
||||
|
||||
|
||||
def test_noop_without_scale_keys():
|
||||
if _FP8 is None:
|
||||
return
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight": torch.randn(4, 4)})
|
||||
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
|
||||
|
||||
|
||||
def test_noop_without_index_or_single_file():
|
||||
if _FP8 is None:
|
||||
return
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
|
||||
|
||||
|
||||
def test_noop_when_not_block_fp8():
|
||||
"""A non-fp8 (or non-block) quantization config is ignored."""
|
||||
scale = torch.rand(2, 2)
|
||||
model = nn.Module()
|
||||
model.config = SimpleNamespace(quantization_config = {"quant_method": "compressed-tensors"})
|
||||
model.layer = nn.Linear(4, 4, bias = False)
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
|
||||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ ST_TAGS = [
|
|||
"v5.2.3",
|
||||
"v5.3.0",
|
||||
"v5.4.1",
|
||||
"v5.5.1",
|
||||
"v5.6.0",
|
||||
"master",
|
||||
]
|
||||
|
||||
|
|
@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str):
|
|||
)
|
||||
|
||||
|
||||
# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881).
|
||||
@pytest.mark.parametrize("tag", ST_TAGS)
|
||||
def test_st_transformer_load_accepts_unsloth_kwargs(tag: str):
|
||||
"""unsloth builds saved ST models via Transformer.load(...) so the saved
|
||||
modality_config is honored (#6881). If .load stops accepting the hub kwargs it
|
||||
passes (and has no **kwargs), update the fix before it silently regresses. Not
|
||||
locating .load is a SKIP (may be inherited); the live test guards the install."""
|
||||
candidates = [
|
||||
"sentence_transformers/models/Transformer.py",
|
||||
"sentence_transformers/models/transformer.py",
|
||||
"sentence_transformers/base/modules/transformer.py",
|
||||
"sentence_transformers/base/modules/module.py",
|
||||
]
|
||||
for p in candidates:
|
||||
src = fetch_text("UKPLab/sentence-transformers", tag, p)
|
||||
if src is None or not has_def(src, "load", "func"):
|
||||
continue
|
||||
m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S)
|
||||
if m is None:
|
||||
continue
|
||||
sig = m.group(1)
|
||||
accepts_var_kw = "**" in sig
|
||||
missing = [
|
||||
kw
|
||||
for kw in ("token", "cache_folder", "revision", "trust_remote_code")
|
||||
if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig))
|
||||
]
|
||||
assert not missing, (
|
||||
f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no "
|
||||
f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module "
|
||||
f"(#6881) before it silently falls back to Transformer(...)."
|
||||
)
|
||||
return
|
||||
pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)")
|
||||
|
||||
|
||||
# sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls.
|
||||
@pytest.mark.parametrize("tag", ST_TAGS)
|
||||
def test_st_util_helpers(tag: str):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue