mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-10 00:08:58 +00:00
8efcc17f47
* Studio: account for DeepSeek-V4 compute buffer in context auto-fit DeepSeek-V4-Flash's lightning indexer plus compressed sparse attention reserve a large context-scaling compute buffer that _compute_buffer_ctx_bytes did not model (the KQ-mask and dequant-scratch rates both miss it, even with an f16 cache). Measured on UD-Q4_K_XL at ub 512 it is about 65.5 GiB at 1M context, which the mask estimate puts near 1.5 GiB, so the auto-fit kept the full 1M train context and llama-server OOM'd allocating the ~70 GB buffer, then spilled to CPU (~4 tok/s). Add a deepseek4-gated flat plus per-token term so the fit caps the context (about 256k on a B200) and the model stays fully on GPU. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
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 Studio'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)
|