mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-11 09:54:09 +00:00
6d8c18cd1a
* Replace standalone Studio wording with Unsloth Replace the single word Studio with Unsloth wherever it is used as shorthand for Unsloth Studio in docs, CLI output, UI strings, i18n locales, workflow display names, comments and docstrings. Kept unchanged: the full name Unsloth Studio, third party product names (LM Studio, Visual Studio, Mac Studio), feature names (Recipe Studio, Fine-tuning Studio and its translations), and all identifiers such as env vars, commands, paths and filenames. * Address review feedback on the Studio wording rename Use "an" before Unsloth where the rename left the article as "a". Restore the split brand where Unsloth and Studio render as two halves of the full product name: the onboarding sidebar subtitle and the IPv6 localhost warning. Scope two messages to the full name Unsloth Studio where plain Unsloth was misleading: the AMD README bullet and the CLI studio setup error.
353 lines
16 KiB
Python
353 lines
16 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Tests for ``_estimate_compute_buffer_bytes``: it scales with ``--parallel``,
|
|
tensor exceeds pipeline, and it is a safe upper bound on the allocations measured
|
|
on real hardware (Qwen3.6-27B-MTP: parallel 1/2/4/8 -> 36/492/1388/3220 MiB single
|
|
GPU, ~600 MiB/device tensor). No GPU, subprocess, or GGUF I/O."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types as _types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|
if _BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, _BACKEND_DIR)
|
|
|
|
_loggers_stub = _types.ModuleType("loggers")
|
|
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
|
sys.modules.setdefault("loggers", _loggers_stub)
|
|
_structlog_stub = _types.ModuleType("structlog")
|
|
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
|
|
sys.modules.setdefault("structlog", _structlog_stub)
|
|
# httpx -- only stub when the real library is missing. Unconditional stubbing
|
|
# shadows HTTPError/Response that huggingface_hub.errors imports at load time,
|
|
# silently breaking the transformers introspection tier in tests collected after
|
|
# this one (the stub leaks via sys.modules for the whole session).
|
|
try:
|
|
import httpx as _httpx_real # noqa: F401
|
|
except ImportError:
|
|
_httpx_stub = _types.ModuleType("httpx")
|
|
for _exc in (
|
|
"ConnectError",
|
|
"TimeoutException",
|
|
"ReadTimeout",
|
|
"ReadError",
|
|
"RemoteProtocolError",
|
|
"CloseError",
|
|
"HTTPError",
|
|
"RequestError",
|
|
):
|
|
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
|
|
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
|
|
_httpx_stub.Response = type("Response", (), {})
|
|
_httpx_stub.Client = type(
|
|
"C",
|
|
(),
|
|
{
|
|
"__init__": lambda s, **kw: None,
|
|
"__enter__": lambda s: s,
|
|
"__exit__": lambda s, *a: None,
|
|
},
|
|
)
|
|
sys.modules["httpx"] = _httpx_stub
|
|
|
|
from core.inference.llama_cpp import LlamaCppBackend
|
|
|
|
MIB = 1024 * 1024
|
|
|
|
|
|
def _backend(
|
|
vocab = 248320,
|
|
embd = 5120,
|
|
mla = None,
|
|
arch = None,
|
|
):
|
|
"""Backend with just the dims the compute-buffer estimate reads."""
|
|
b = LlamaCppBackend.__new__(LlamaCppBackend)
|
|
b._vocab_size = vocab
|
|
b._embedding_length = embd
|
|
b._key_length_mla = mla # non-None -> MLA (compressed attention)
|
|
b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4')
|
|
return b
|
|
|
|
|
|
# Measured ground truth (MiB) the estimate must upper-bound.
|
|
_PIPELINE_MEASURED = {1: 36, 2: 492, 4: 1388, 8: 3220}
|
|
_TENSOR_MEASURED_PER_DEVICE = 600
|
|
|
|
|
|
class TestSafeUpperBound:
|
|
"""The estimate must be >= every measured allocation (never under-reserve)."""
|
|
|
|
@pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items()))
|
|
def test_pipeline_upper_bounds_measured(self, parallel, measured):
|
|
est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB
|
|
assert est >= measured, f"under-reserved at parallel={parallel}: {est:.0f} < {measured}"
|
|
|
|
@pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items()))
|
|
def test_pipeline_not_wildly_over(self, parallel, measured):
|
|
# Stay within ~2x of measured so we don't waste context (the point of
|
|
# replacing the flat reserve). parallel=1 is tiny in absolute terms.
|
|
est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB
|
|
assert est <= max(measured * 2.0, 128)
|
|
|
|
def test_tensor_upper_bounds_measured(self):
|
|
est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB
|
|
assert est >= _TENSOR_MEASURED_PER_DEVICE
|
|
|
|
def test_tensor_far_below_old_flat_reserve(self):
|
|
# The whole point: deterministic estimate << flat 5120 for this model.
|
|
est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB
|
|
assert est < LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
|
|
|
|
|
|
class TestScaling:
|
|
def test_grows_with_serving_slots(self):
|
|
b = _backend()
|
|
vals = [b._estimate_compute_buffer_bytes(n_parallel = p) for p in (1, 2, 4, 8)]
|
|
assert vals == sorted(vals) and vals[0] < vals[-1]
|
|
|
|
def test_parallel_1_is_small(self):
|
|
# Single-token decode: a few tens of MiB, not gigabytes.
|
|
est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1) / MIB
|
|
assert est < 128
|
|
|
|
def test_tensor_exceeds_pipeline_at_same_parallel(self):
|
|
b = _backend()
|
|
pipe = b._estimate_compute_buffer_bytes(n_parallel = 1)
|
|
tens = b._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True)
|
|
assert tens > pipe
|
|
|
|
def test_scales_with_vocab(self):
|
|
small = _backend(vocab = 32000)._estimate_compute_buffer_bytes(n_parallel = 4)
|
|
big = _backend(vocab = 256000)._estimate_compute_buffer_bytes(n_parallel = 4)
|
|
assert big > small
|
|
|
|
def test_scales_with_ubatch(self):
|
|
b = _backend()
|
|
lo = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 256)
|
|
hi = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 1024)
|
|
assert hi > lo
|
|
|
|
|
|
class TestFallback:
|
|
def test_zero_when_vocab_missing(self):
|
|
assert _backend(vocab = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0
|
|
|
|
def test_zero_when_embd_missing(self):
|
|
assert _backend(embd = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0
|
|
|
|
def test_zero_lets_tensor_plan_use_flat_fallback(self):
|
|
# When dims are missing, _plan_tensor_parallel must fall back to the flat
|
|
# reserve (defense-in-depth) rather than reserving 0 and OOMing.
|
|
b = _backend(vocab = None, embd = None)
|
|
b._n_layers = None # can't estimate KV -> floors ctx, still returns a plan
|
|
ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, 48000)], 8 * 1024**3, 8192)
|
|
assert gi == [0, 1] # both GPUs usable under the flat fallback
|
|
|
|
|
|
class TestParallel1Default:
|
|
"""At Unsloth's default --parallel 1 the buffer is negligible in pipeline."""
|
|
|
|
def test_default_n_parallel(self):
|
|
est = _backend()._estimate_compute_buffer_bytes() / MIB
|
|
assert est < 128
|
|
|
|
|
|
class TestContextLinearBuffer:
|
|
"""``_compute_buffer_ctx_bytes``: the flash-attn KQ-mask + attention scratch
|
|
grow ~linearly with context; the flat estimate above only covers ctx -> 0.
|
|
Measured slope (q8_0 KV, ubatch 512) was 0.74-2.02 x n_embd; 2 x n_embd is the
|
|
worst-case upper bound the term must hold to."""
|
|
|
|
# (model, n_embd, ctx, measured CUDA0 compute buffer MiB at that ctx, q8_0/ub512)
|
|
_MEASURED = [
|
|
("Qwen3.5-2B", 2048, 262144, 796),
|
|
("Qwen3.5-4B", 2560, 262144, 1330), # worst slope, 2.02 x n_embd
|
|
("Qwen3.5-9B", 4096, 262144, 1336),
|
|
("Qwen3.6-27B", 5120, 262144, 1360),
|
|
("Gemma-4-31B", 5376, 262144, 2392),
|
|
]
|
|
|
|
def test_zero_by_default(self):
|
|
# Omitted/zero ctx -> no term (keeps the flat callers unchanged).
|
|
assert _backend()._compute_buffer_ctx_bytes(0) == 0
|
|
|
|
def test_zero_when_embd_missing(self):
|
|
assert _backend(embd = None)._compute_buffer_ctx_bytes(262144) == 0
|
|
|
|
def test_grows_linearly_with_context(self):
|
|
b = _backend(embd = 4096)
|
|
a = b._compute_buffer_ctx_bytes(65536)
|
|
d = b._compute_buffer_ctx_bytes(131072)
|
|
assert d == pytest.approx(2 * a, rel = 1e-6)
|
|
|
|
def test_scales_with_embd(self):
|
|
# The quantized (dequant-scratch) rate scales with n_embd; f16 (mask) does not.
|
|
small = _backend(embd = 2048)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0")
|
|
big = _backend(embd = 5120)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0")
|
|
assert big > small
|
|
|
|
def test_scales_with_ubatch(self):
|
|
b = _backend(embd = 4096)
|
|
lo = b._compute_buffer_ctx_bytes(131072, n_ubatch = 256)
|
|
hi = b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024)
|
|
assert hi > lo
|
|
|
|
@pytest.mark.parametrize("name,embd,ctx,measured", _MEASURED)
|
|
def test_upper_bounds_measured_compute_growth(self, name, embd, ctx, measured):
|
|
# flat term + context-linear term must cover the real (q8_0) buffer at full ctx.
|
|
b = _backend(embd = embd)
|
|
flat = b._estimate_compute_buffer_bytes(n_parallel = 1)
|
|
total = (flat + b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0")) / MIB
|
|
assert total >= measured, f"{name}: under-reserved {total:.0f} < {measured}"
|
|
|
|
def test_worst_case_rate_covers_two_x_embd(self):
|
|
# >= 2 x n_embd bytes per context token at the default micro-batch (the worst
|
|
# measured quantized slope, Qwen3.5-4B), so flat + term upper-bounds the buffer.
|
|
embd = 4096
|
|
b = _backend(embd = embd)
|
|
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "q8_0") / 100000
|
|
assert per_tok >= 2 * embd
|
|
|
|
|
|
class TestContextBufferKVQuant:
|
|
"""The context-linear rate depends on the KV cache type: a quantized cache adds a
|
|
context-sized dequant scratch (heavy); f16/bf16/f32 only pays the KQ mask (light).
|
|
Measured Qwen3.5-4B at 256k: 1.30 GiB (q8_0) vs 0.31 GiB (f16)."""
|
|
|
|
def test_quantized_heavier_than_f16(self):
|
|
b = _backend(embd = 4096)
|
|
q = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0")
|
|
f = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16")
|
|
assert q > f
|
|
|
|
def test_none_cache_type_is_f16(self):
|
|
# None -> f16 (llama.cpp's default); the env-quantized case is covered by the
|
|
# KV budget's f16 over-reservation, so we take the lighter mask-only rate.
|
|
b = _backend(embd = 4096)
|
|
assert b._compute_buffer_ctx_bytes(
|
|
131072, cache_type_kv = None
|
|
) == b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16")
|
|
|
|
@pytest.mark.parametrize("ct", ["f16", "bf16", "f32"])
|
|
def test_unquantized_uses_mask_only_rate(self, ct):
|
|
# f16/bf16/f32: KQ mask only, n_ubatch*2 B/tok, independent of n_embd.
|
|
b_small = _backend(embd = 2048)
|
|
b_big = _backend(embd = 8192)
|
|
per_small = b_small._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000
|
|
per_big = b_big._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000
|
|
assert per_small == per_big # no n_embd scaling on the f16 path
|
|
expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY # ubatch 512
|
|
assert per_small == pytest.approx(expected, rel = 1e-6)
|
|
|
|
@pytest.mark.parametrize("ct", ["q8_0", "q5_1", "q4_0", "iq4_nl"])
|
|
def test_quantized_types_use_heavy_rate(self, ct):
|
|
embd = 4096
|
|
b = _backend(embd = embd)
|
|
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000
|
|
assert per_tok == pytest.approx(
|
|
LlamaCppBackend._CTX_COMPUTE_BYTES_PER_EMBD * embd, rel = 1e-6
|
|
)
|
|
|
|
def test_f16_covers_measured_mask(self):
|
|
# f16 buffer is ~mask only (~n_ubatch*2 B/tok); 0.5 x n_embd must cover the
|
|
# measured Qwen3.5-4B f16 slope (~0.4 x n_embd = 0.31 GiB at 256k).
|
|
b = _backend(embd = 2560) # Qwen3.5-4B
|
|
est = b._compute_buffer_ctx_bytes(262144, cache_type_kv = "f16") / MIB
|
|
assert est >= 320 # measured 0.31 GiB growth
|
|
|
|
|
|
class TestContextBufferMLA:
|
|
"""MLA (compressed attention) needs a smaller quantized dequant scratch than
|
|
regular attention: measured 0.94 x n_embd on GLM-5.2 and Kimi-K2.7 vs up to
|
|
2.02x on Qwen/Gemma. Charging the regular rate would badly over-reserve a tight
|
|
multi-GPU MLA pin (per-device scaling multiplies the error)."""
|
|
|
|
def test_mla_lighter_than_regular(self):
|
|
reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
|
|
mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
|
|
assert mla < reg
|
|
|
|
@pytest.mark.parametrize(
|
|
"name,embd,ctx,measured",
|
|
[
|
|
("GLM-5.2", 6144, 754688, 4141), # per-device compute MiB at q8_0
|
|
("Kimi-K2.7", 7168, 262144, 1690),
|
|
],
|
|
)
|
|
def test_mla_rate_covers_measured(self, name, embd, ctx, measured):
|
|
b = _backend(embd = embd, mla = 256)
|
|
est = b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0") / MIB
|
|
assert est >= measured, f"{name}: MLA under-reserved {est:.0f} < {measured}"
|
|
|
|
def test_mla_not_wildly_over(self):
|
|
# 1.25 x n_embd should stay within ~1.6x of the measured 0.94x (not 2.4x like
|
|
# the regular 2.25 rate would), so a multi-GPU MLA pin keeps its context.
|
|
b = _backend(embd = 6144, mla = 256)
|
|
est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB
|
|
assert est <= 4141 * 1.7
|
|
|
|
|
|
class TestContextBufferDSV4:
|
|
"""DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention
|
|
compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache).
|
|
Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit
|
|
must see this so it does not commit the full 1M train context and OOM (spilling
|
|
to CPU at ~4 tok/s)."""
|
|
|
|
_MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx
|
|
GIB = 1024**3
|
|
|
|
def test_covers_measured_1m_buffer(self):
|
|
b = _backend(embd = 4096, arch = "deepseek4")
|
|
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
|
|
assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}"
|
|
|
|
def test_not_wildly_over_at_1m(self):
|
|
# Within ~1.3x of measured so the fit still grants a large (~256k) context.
|
|
b = _backend(embd = 4096, arch = "deepseek4")
|
|
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
|
|
assert gib <= self._MEASURED_1M_GIB * 1.3
|
|
|
|
def test_fires_for_f16_cache(self):
|
|
# The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must
|
|
# reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx.
|
|
dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes(
|
|
262144, cache_type_kv = "f16"
|
|
)
|
|
other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes(
|
|
262144, cache_type_kv = "f16"
|
|
)
|
|
assert dsv4 > 40 * other
|
|
|
|
def test_cache_type_independent(self):
|
|
# Indexer scratch is present for an f16 and a quantized cache alike.
|
|
b = _backend(embd = 4096, arch = "deepseek4")
|
|
assert b._compute_buffer_ctx_bytes(
|
|
262144, cache_type_kv = "f16"
|
|
) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
|
|
|
|
def test_flat_floor_at_small_ctx(self):
|
|
# ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB).
|
|
b = _backend(embd = 4096, arch = "deepseek4")
|
|
assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0
|
|
|
|
def test_scales_with_context_and_ubatch(self):
|
|
b = _backend(embd = 4096, arch = "deepseek4")
|
|
assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536)
|
|
assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes(
|
|
131072, n_ubatch = 256
|
|
)
|
|
|
|
def test_non_dsv4_unchanged(self):
|
|
# Regression guard: a non-deepseek4 model keeps the mask-only f16 rate.
|
|
b = _backend(embd = 4096, arch = "llama")
|
|
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000
|
|
expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY
|
|
assert per_tok == pytest.approx(expected, rel = 1e-6)
|