RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
soup/tests/test_v07202.py

1807 lines
70 KiB
Python
Raw Permalink Blame History

"""v0.72.2 — NF4 layer streaming.
The first slot that gives a 4 GB card something with no alternative: an 8B base
quantised to NF4 is ~3.6 GB of host RAM, page-lockable under this box's measured
7.12 GB ceiling, so 8B fine-tuning on a laptop becomes real.
Gate 1 (bit-exactness vs **resident NF4**, not resident bf16) passed 6/6 before a
line of ``src/`` was written — see ``.claude/v0721-gate-results.md``. These tests
pin the properties that gate established, plus the silent-failure classes it
surfaced:
* **PEFT dispatches a different LoRA math path** unless the model carries the
``is_loaded_in_4bit`` markers ``from_pretrained`` stamps. Without them PEFT
picks the generic ``lora.layer.Linear``, which still *runs* against a
``Linear4bit`` base but casts and accumulates differently — measured as a
9.375e-01 logit divergence with byte-identical weights AND adapters. No crash,
no warning, a healthy-looking loss curve.
* **A shard cache keyed without the quantisation** would stream bf16 bytes into
an NF4 model (or the reverse) and mis-train silently.
* ``Params4bit`` carries a ``quant_state`` and cannot be byte-copied into a
plain buffer (plan P3) — the views must be *rebuilt* over the pooled buffer.
"""
import os
import sys
import pytest
pytestmark = pytest.mark.filterwarnings("ignore::UserWarning")
def _cuda() -> bool:
try:
import torch
except ImportError: # pragma: no cover - torch is a [train] extra
return False
return torch.cuda.is_available()
requires_cuda = pytest.mark.skipif(not _cuda(), reason="needs a CUDA device")
def _mps_is_the_accelerator() -> bool:
"""True on an Apple-Silicon runner with no CUDA.
``TrainingArguments`` picks ``mps`` as its device there, while this suite
builds the streamed model on ``cpu``; a real training step then moves the
batch to MPS and hits "Placeholder storage has not been allocated on MPS
device". NF4 streaming is measured on CUDA and CPU only — bitsandbytes'
4-bit kernels are not supported on MPS at all — so the step is skipped
rather than making an unverified claim about it. Mirrors the identical
guard in tests/test_v07200.py.
"""
try:
import torch
if torch.cuda.is_available():
return False
backend = getattr(torch.backends, "mps", None)
return bool(backend is not None and backend.is_available())
except Exception:
return False
skip_on_mps = pytest.mark.skipif(
_mps_is_the_accelerator(),
reason="MPS is untested for NF4 streaming (measured on CUDA + CPU only)",
)
# ==========================================================================
# fixtures
# ==========================================================================
def _write_safetensors(path, tensors):
from safetensors.torch import save_file
save_file({k: v.clone() for k, v in tensors.items()}, path)
return path
#: The two decoder weights the fake checkpoint quantises. Layernorms stay
#: unquantised, exactly as ``replace_with_bnb_linear`` leaves them — so every
#: shard here exercises the MIXED case.
QUANT_SUFFIXES = frozenset({"self_attn.q_proj.weight", "mlp.down_proj.weight"})
def _fake_weights_dir(tmp_path, n_layers=3):
"""A tiny llama-shaped checkpoint. Shapes are block-size friendly."""
import torch
torch.manual_seed(0)
blob = {}
for idx in range(n_layers):
pre = f"model.layers.{idx}."
blob[pre + "self_attn.q_proj.weight"] = torch.randn(64, 64, dtype=torch.float32)
blob[pre + "mlp.down_proj.weight"] = torch.randn(64, 128, dtype=torch.float32)
blob[pre + "input_layernorm.weight"] = torch.randn(64, dtype=torch.float32)
blob["model.embed_tokens.weight"] = torch.randn(32, 64, dtype=torch.float32)
blob["model.norm.weight"] = torch.randn(64, dtype=torch.float32)
src = tmp_path / "weights"
src.mkdir()
_write_safetensors(str(src / "model.safetensors"), blob)
return str(src)
# ==========================================================================
# the sharder
# ==========================================================================
class TestNF4Sharding:
def test_packed_absmax_and_specs_are_written(self, tmp_path):
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
ABSMAX_SUFFIX,
NESTED_ABSMAX_SUFFIX,
NESTED_OFFSET_SUFFIX,
QUANT_NF4,
layer_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
index = shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
assert index.quant == QUANT_NF4
assert index.double_quant is True
assert set(index.quant_specs) == set(QUANT_SUFFIXES)
blob = load_file(layer_shard_path(out, 0))
for key in QUANT_SUFFIXES:
assert blob[key].dtype.__str__() == "torch.uint8", key
assert key + ABSMAX_SUFFIX in blob
assert key + NESTED_ABSMAX_SUFFIX in blob
assert key + NESTED_OFFSET_SUFFIX in blob
# an unquantised weight in the SAME layer is untouched
assert str(blob["input_layernorm.weight"].dtype) == "torch.float32"
assert "input_layernorm.weight" + ABSMAX_SUFFIX not in blob
@pytest.mark.parametrize("double_quant", [True, False])
def test_rebuilt_state_matches_the_one_bitsandbytes_built(self, tmp_path, double_quant):
"""The load-bearing property: a ``QuantState`` reassembled from the
streamed tensors must dequantise BIT-IDENTICALLY to the one bnb kept
alongside the packed bytes. An absolute tolerance against the original
weight would only measure NF4's (lossy) quantisation error and would
pass for a subtly wrong reconstruction.
Parametrized over ``double_quant`` so ``rebuild_quant_state``'s
``state2=None`` branch is exercised NUMERICALLY, not just by checking
which sidecar keys got written."""
import torch
from bitsandbytes.functional import dequantize_4bit, quantize_4bit
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
NF4_BLOCKSIZE,
QUANT_NF4,
extras_shard_path,
layer_shard_path,
shard_checkpoint,
)
from soup_cli.utils.layer_stream_runtime import rebuild_quant_state
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
index = shard_checkpoint(
src,
out,
dtype="float32",
quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES,
double_quant=double_quant,
)
blob = load_file(layer_shard_path(out, 0))
codes = load_file(extras_shard_path(out))
original = load_file(os.path.join(src, "model.safetensors"))
key = "mlp.down_proj.weight"
spec = index.quant_specs[key]
assert spec.nested is double_quant
want_tensor = original["model.layers.0." + key]
reference_packed, reference_state = quantize_4bit(
want_tensor,
blocksize=NF4_BLOCKSIZE,
compress_statistics=double_quant,
quant_type="nf4",
)
mine = dequantize_4bit(blob[key], rebuild_quant_state(key, blob, spec, codes))
theirs = dequantize_4bit(reference_packed, reference_state)
assert torch.equal(mine, theirs), (mine - theirs).abs().max().item()
def test_code_table_divergence_is_refused(self, tmp_path):
"""One shared resident copy of the codebook is only safe because it is
constant across weights. The guard is unreachable with today's bnb, so
exercise it directly rather than leaving a `raise` nobody has run."""
import torch
from soup_cli.utils.layer_shard import _CodeTables
tables = _CodeTables()
tables.observe("a", torch.zeros(16), None)
with pytest.raises(ValueError, match="code table differs"):
tables.observe("b", torch.ones(16), None)
def test_nested_code_table_divergence_is_refused(self, tmp_path):
import torch
from soup_cli.utils.layer_shard import _CodeTables
tables = _CodeTables()
tables.observe("a", torch.zeros(16), torch.zeros(256))
with pytest.raises(ValueError, match="nested NF4 code table differs"):
tables.observe("b", torch.zeros(16), torch.ones(256))
def test_reconstruction_is_the_right_weight(self, tmp_path):
"""A faithful rebuild of the WRONG tensor would satisfy the test above.
NF4 is lossy, so the standard is relative: much closer to its own
weight than to a sibling of identical shape."""
from bitsandbytes.functional import dequantize_4bit
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
QUANT_NF4,
extras_shard_path,
layer_shard_path,
shard_checkpoint,
)
from soup_cli.utils.layer_stream_runtime import rebuild_quant_state
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
index = shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
blob = load_file(layer_shard_path(out, 0))
codes = load_file(extras_shard_path(out))
original = load_file(os.path.join(src, "model.safetensors"))
key = "mlp.down_proj.weight"
back = dequantize_4bit(
blob[key], rebuild_quant_state(key, blob, index.quant_specs[key], codes)
)
mine = original["model.layers.0." + key].to(back.dtype)
sibling = original["model.layers.1." + key].to(back.dtype)
def rel(other):
return ((back - other).norm() / other.norm()).item()
assert rel(mine) < 0.15, rel(mine)
assert rel(sibling) > 1.0, rel(sibling)
def test_quantisation_is_byte_deterministic(self, tmp_path):
"""Offline sharding only reproduces a resident load if this holds."""
import torch
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
QUANT_NF4,
layer_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
first = str(tmp_path / "a")
second = str(tmp_path / "b")
for out in (first, second):
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
left = load_file(layer_shard_path(first, 0))
right = load_file(layer_shard_path(second, 0))
assert set(left) == set(right)
for key in left:
assert torch.equal(left[key], right[key]), key
def test_code_tables_are_shared_in_extras(self, tmp_path):
"""One resident copy of the 16-entry NF4 code table, not one per weight."""
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
NF4_CODE_KEY,
NF4_NESTED_CODE_KEY,
QUANT_NF4,
extras_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
extras = load_file(extras_shard_path(out))
assert extras[NF4_CODE_KEY].shape == (16,)
assert extras[NF4_NESTED_CODE_KEY].shape == (256,)
def test_zero_dim_offset_survives_the_round_trip(self, tmp_path):
"""``quant_state.offset`` is a 0-dim scalar; a shape-assuming writer
would flatten it and the dequant would land somewhere else entirely."""
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
NESTED_OFFSET_SUFFIX,
QUANT_NF4,
layer_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
blob = load_file(layer_shard_path(out, 0))
offset = blob["mlp.down_proj.weight" + NESTED_OFFSET_SUFFIX]
assert offset.shape == ()
def test_no_double_quant_omits_the_nested_tensors(self, tmp_path):
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
ABSMAX_SUFFIX,
NESTED_ABSMAX_SUFFIX,
QUANT_NF4,
layer_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
index = shard_checkpoint(
src,
out,
dtype="float32",
quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES,
double_quant=False,
)
assert index.double_quant is False
assert index.quant_specs["mlp.down_proj.weight"].nested is False
blob = load_file(layer_shard_path(out, 0))
assert "mlp.down_proj.weight" + ABSMAX_SUFFIX in blob
assert "mlp.down_proj.weight" + NESTED_ABSMAX_SUFFIX not in blob
class TestQuantCacheInvalidation:
"""A cache keyed without the quantisation streams the WRONG BYTES."""
def test_bf16_cache_is_not_reused_for_an_nf4_request(self, tmp_path):
from soup_cli.utils.layer_shard import (
QUANT_NF4,
QUANT_NONE,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
plain = shard_checkpoint(src, out, dtype="float32")
assert plain.quant == QUANT_NONE
nf4 = shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
assert nf4.quant == QUANT_NF4
assert nf4.quant_specs
def test_nf4_cache_is_not_reused_for_a_bf16_request(self, tmp_path):
from soup_cli.utils.layer_shard import QUANT_NF4, QUANT_NONE, shard_checkpoint
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
plain = shard_checkpoint(src, out, dtype="float32")
assert plain.quant == QUANT_NONE
assert not plain.quant_specs
def test_double_quant_change_invalidates(self, tmp_path):
from soup_cli.utils.layer_shard import QUANT_NF4, shard_checkpoint
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
again = shard_checkpoint(
src,
out,
dtype="float32",
quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES,
double_quant=False,
)
assert again.double_quant is False
def test_quant_device_change_invalidates(self, tmp_path):
"""CPU and CUDA agree on the packed nibbles but not on every float32
nested statistic, so a CPU-quantised cache reused for a CUDA run would
break bit-exactness against a resident load. dtype happens to co-vary
with device today, which would mask this — hence an explicit key."""
from soup_cli.utils.layer_shard import QUANT_NF4, shard_checkpoint
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
first = shard_checkpoint(
src,
out,
dtype="float32",
quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES,
quant_device="cpu",
)
assert first.quant_device == "cpu"
if not _cuda():
pytest.skip("needs a CUDA device to prove the invalidation")
second = shard_checkpoint(
src,
out,
dtype="float32",
quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES,
quant_device="cuda",
)
assert second.quant_device == "cuda"
def test_device_ordinal_does_not_invalidate(self, tmp_path):
"""cuda:0 and cuda:1 quantise identically — only the KIND is keyed, or
every multi-GPU box would reshard on a different ordinal."""
from soup_cli.utils.layer_shard import (
QUANT_NF4,
layer_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES, quant_device="cpu",
)
stamp = os.stat(layer_shard_path(out, 0)).st_mtime_ns
again = shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES, quant_device="cpu:0",
)
assert again.quant_device == "cpu"
assert os.stat(layer_shard_path(out, 0)).st_mtime_ns == stamp
def test_identical_request_hits_the_cache(self, tmp_path):
"""The control: without this, the three tests above prove nothing —
a sharder that never caches passes all of them."""
from soup_cli.utils.layer_shard import (
QUANT_NF4,
layer_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
stamp = os.stat(layer_shard_path(out, 0)).st_mtime_ns
shard_checkpoint(
src, out, dtype="float32", quant=QUANT_NF4, quant_suffixes=QUANT_SUFFIXES
)
assert os.stat(layer_shard_path(out, 0)).st_mtime_ns == stamp
class TestShardQuantGuards:
def test_unknown_quant_is_refused_naming_the_allowlist(self, tmp_path):
from soup_cli.utils.layer_shard import shard_checkpoint
src = _fake_weights_dir(tmp_path)
with pytest.raises(ValueError, match="nf4"):
shard_checkpoint(src, str(tmp_path / "o"), dtype="float32", quant="int8")
def test_nf4_without_suffixes_is_refused(self, tmp_path):
"""Silently writing unquantised bytes under an ``nf4`` label would
stream full-precision weights into ``Linear4bit`` modules."""
from soup_cli.utils.layer_shard import QUANT_NF4, shard_checkpoint
src = _fake_weights_dir(tmp_path)
with pytest.raises(ValueError, match="quant_suffixes"):
shard_checkpoint(src, str(tmp_path / "o"), dtype="float32", quant=QUANT_NF4)
def test_suffixes_without_nf4_is_refused(self, tmp_path):
from soup_cli.utils.layer_shard import shard_checkpoint
src = _fake_weights_dir(tmp_path)
with pytest.raises(ValueError, match="quant_suffixes"):
shard_checkpoint(
src, str(tmp_path / "o"), dtype="float32", quant_suffixes=QUANT_SUFFIXES
)
def test_a_layer_with_different_shapes_is_refused(self, tmp_path):
"""Layer 0's ``quant_state`` (its absmax blocking) is reused for every
layer, so a differently-shaped sibling would be reconstructed against
the wrong statistics. The pre-NF4 guard compared names only."""
import torch
from soup_cli.utils.layer_shard import QUANT_NF4, shard_checkpoint
torch.manual_seed(0)
blob = {}
for idx, width in enumerate((128, 256)): # layer 1 is deliberately wider
pre = f"model.layers.{idx}."
blob[pre + "self_attn.q_proj.weight"] = torch.randn(64, 64)
blob[pre + "mlp.down_proj.weight"] = torch.randn(64, width)
blob[pre + "input_layernorm.weight"] = torch.randn(64)
blob["model.embed_tokens.weight"] = torch.randn(32, 64)
blob["model.norm.weight"] = torch.randn(64)
src = tmp_path / "ragged"
src.mkdir()
_write_safetensors(str(src / "model.safetensors"), blob)
src = str(src)
with pytest.raises(ValueError, match="different tensor shapes"):
shard_checkpoint(
src,
str(tmp_path / "o"),
dtype="float32",
quant=QUANT_NF4,
quant_suffixes=QUANT_SUFFIXES,
)
def test_a_suffix_absent_from_the_checkpoint_is_refused(self, tmp_path):
"""A typo'd or stale suffix set means those weights ship unquantised."""
from soup_cli.utils.layer_shard import QUANT_NF4, shard_checkpoint
src = _fake_weights_dir(tmp_path)
with pytest.raises(ValueError, match="self_attn.nonexistent.weight"):
shard_checkpoint(
src,
str(tmp_path / "o"),
dtype="float32",
quant=QUANT_NF4,
quant_suffixes=frozenset({"self_attn.nonexistent.weight"}),
)
class TestHostileIndexIsRefused:
"""``index.json`` is read back off disk and its ``shape`` / ``blocksize``
flow into bitsandbytes' native dequantise kernels, which do NOT bounds-check
against the real packed tensor. A corrupted or tampered index must fail as a
clean Python exception, never as an out-of-bounds read in C."""
def _spec(self, **over):
from soup_cli.utils.layer_shard import NF4WeightSpec
payload = {
"shape": [64, 128],
"dtype": "float32",
"blocksize": 64,
"quant_type": "nf4",
"nested": False,
"nested_blocksize": 0,
}
payload.update(over)
return NF4WeightSpec.from_json(payload)
def test_wellformed_spec_round_trips(self):
"""Control — the guards below must not reject honest input."""
assert self._spec().shape == (64, 128)
@pytest.mark.parametrize(
"field,value,needle",
[
("shape", [0, 8], "positive"),
("shape", [-4, 8], "positive"),
("shape", [], "positive"),
("shape", [2**20, 2**20], "too large"),
("blocksize", 0, "blocksize"),
("blocksize", -64, "blocksize"),
("quant_type", "evil", "quant_type"),
("dtype", "complex128", "dtype"),
],
)
def test_malformed_field_is_refused(self, field, value, needle):
with pytest.raises(ValueError, match=needle):
self._spec(**{field: value})
def test_nested_blocksize_required_when_nested(self):
with pytest.raises(ValueError, match="nested_blocksize"):
self._spec(nested=True, nested_blocksize=0)
def test_shape_claiming_more_than_the_packed_bytes_is_refused(self, tmp_path):
"""The actual attack shape: an honest shard, an index that overstates
the tensor. The kernel would read past the buffer."""
import json
from soup_cli.utils.layer_shard import QUANT_NF4, read_shard_index, shard_checkpoint
from soup_cli.utils.layer_stream_runtime import build_streamed_model
weights, _, _ = _tiny_llama_dir(tmp_path, n_layers=2)
shards = str(tmp_path / "shards")
index = shard_checkpoint(
weights,
shards,
dtype="float32",
arch="llama",
quant=QUANT_NF4,
quant_suffixes=frozenset({"mlp.down_proj.weight"}),
quant_device="cpu",
)
path = os.path.join(shards, "index.json")
with open(path, encoding="utf-8") as fh:
payload = json.load(fh)
# claim 8x the real element count
real = payload["quant_specs"]["mlp.down_proj.weight"]["shape"]
payload["quant_specs"]["mlp.down_proj.weight"]["shape"] = [real[0] * 8, real[1]]
with open(path, "w", encoding="utf-8") as fh:
json.dump(payload, fh)
with pytest.raises(ValueError, match="inconsistent|packed"):
build_streamed_model(
model_id=weights,
shard_dir=shards,
index=read_shard_index(shards),
lora_config=_tiny_lora(),
device="cpu",
dtype="float32",
buffers=2,
pin=False,
quant=QUANT_NF4,
)
assert index.quant == QUANT_NF4 # the honest index was fine
def test_an_honest_shard_passes_the_consistency_check(self, tmp_path):
"""The control for the hostile case above: the guard is one-sided
(bitsandbytes pads a non-block-aligned tensor up to the next block), so
an accidentally-strict `>=` must not reject real shards. The 3-element
layernorm-sized weight below is deliberately NOT block-aligned."""
from soup_cli.utils.layer_shard import ABSMAX_SUFFIX, NF4WeightSpec
from soup_cli.utils.layer_stream_runtime import validate_quant_shape
# 100 elements at blocksize 64 -> bnb pads to 128 -> 64 packed bytes,
# 2 absmax blocks. The index honestly claims 100.
spec = NF4WeightSpec((10, 10), "float32", 64, "nf4", False, 0)
shard_spec = {
"w": ((64, 1), "uint8"),
"w" + ABSMAX_SUFFIX: ((2,), "float32"),
}
validate_quant_shape("w", spec, shard_spec) # must not raise
# exactly-aligned is the common case and must also pass
exact = NF4WeightSpec((64, 64), "float32", 64, "nf4", False, 0)
validate_quant_shape(
"w",
exact,
{"w": ((2048, 1), "uint8"), "w" + ABSMAX_SUFFIX: ((64,), "float32")},
)
def test_quant_specs_must_agree_with_the_quant_field(self, tmp_path):
"""``quant='none'`` matches the cache key, so a tampered index carrying
specs anyway would be accepted and then reconstruct NF4 against a
non-4bit skeleton."""
from soup_cli.utils.layer_shard import QUANT_NONE, NF4WeightSpec, ShardIndex
from soup_cli.utils.layer_stream_runtime import install_streaming
bogus = ShardIndex(
n_layers=1,
layer_keys=(),
extra_keys=(),
dtype="float32",
total_params=0,
arch="llama",
soup_version="x",
quant=QUANT_NONE,
quant_specs={"w": NF4WeightSpec((4, 4), "float32", 64, "nf4", False, 0)},
)
with pytest.raises(ValueError, match="quant_specs"):
install_streaming(object(), shard_dir="/nope", index=bogus, device="cpu")
class TestUnquantisedPathUnchanged:
"""v0.72.0's bf16 gates stay valid only if that path is byte-untouched."""
def test_plain_shard_has_no_quant_metadata(self, tmp_path):
from safetensors.torch import load_file
from soup_cli.utils.layer_shard import (
NF4_CODE_KEY,
QUANT_NONE,
extras_shard_path,
shard_checkpoint,
)
src = _fake_weights_dir(tmp_path)
out = str(tmp_path / "shards")
index = shard_checkpoint(src, out, dtype="float32")
assert index.quant == QUANT_NONE
assert index.quant_specs == {}
assert NF4_CODE_KEY not in load_file(extras_shard_path(out))
# ==========================================================================
# the meta skeleton — the gate's single most important finding
# ==========================================================================
def _tiny_llama_dir(tmp_path, n_layers=2, tie=True):
"""A real (tiny) Llama checkpoint on disk — no network.
``hidden_size`` is 64 and NOT 32, and that is load-bearing for the NF4
tests. bitsandbytes' CPU 4-bit forward calls
``_convert_weight_packed_for_cpu``, which reshapes absmax to
``[rows, blocks_per_row]``. At hidden 32 a weight has 32*32/64 = 16 absmax
blocks for 32 rows, so ``blocks_per_row`` floors to **zero** and it raises
``shape '[32, 0]' is invalid for input of size 16``. At 64 there are 64
blocks for 64 rows and it is fine.
A CUDA build never calls that function, so this is invisible on a GPU
machine and fails on every CPU-only CI runner. Do not shrink this back.
"""
import torch
from safetensors.torch import save_file
from transformers import LlamaConfig, LlamaForCausalLM
torch.manual_seed(7)
config = LlamaConfig(
vocab_size=64,
hidden_size=64,
intermediate_size=64,
num_hidden_layers=n_layers,
num_attention_heads=4,
num_key_value_heads=2,
tie_word_embeddings=tie,
max_position_embeddings=128,
)
model = LlamaForCausalLM(config).to(torch.float32).eval()
weights = tmp_path / "model"
weights.mkdir(parents=True, exist_ok=True)
state = {k: v.contiguous() for k, v in model.state_dict().items()}
if tie:
state.pop("lm_head.weight", None)
save_file(state, str(weights / "model.safetensors"))
config.save_pretrained(str(weights))
return str(weights), model, config
def _write_tiny_tokenizer(directory):
"""A real, offline PreTrainedTokenizerFast so ``setup()`` can tokenize
without touching the HF cache or the network."""
import json
from tokenizers import Tokenizer, models, pre_tokenizers
vocab = {"<unk>": 0, "<s>": 1, "</s>": 2, "<pad>": 3}
for word in ("hello", "world", "hi", "yo", "the", "cat", "sat", "on", "mat"):
vocab[word] = len(vocab)
tokenizer = Tokenizer(models.WordLevel(vocab=vocab, unk_token="<unk>"))
tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
tokenizer.save(os.path.join(directory, "tokenizer.json"))
with open(os.path.join(directory, "tokenizer_config.json"), "w", encoding="utf-8") as fh:
json.dump(
{
"tokenizer_class": "PreTrainedTokenizerFast",
"unk_token": "<unk>",
"bos_token": "<s>",
"eos_token": "</s>",
"pad_token": "<pad>",
"model_max_length": 128,
"clean_up_tokenization_spaces": False,
},
fh,
)
def _tiny_lora():
from peft import LoraConfig, TaskType
return LoraConfig(
r=4,
lora_alpha=8,
lora_dropout=0.0,
bias="none",
target_modules=["q_proj", "v_proj"],
task_type=TaskType.CAUSAL_LM,
)
def _lora_class_name(model):
layer = model.base_model.model.model.layers[0].self_attn.q_proj
return f"{type(layer).__module__}.{type(layer).__name__}"
class TestPeftDispatchesTheBnbLoraPath:
"""With byte-identical weights AND byte-identical adapters, streamed and
resident logits differed by **9.375e-01** because PEFT silently chose a
different LoRA implementation. ``from_pretrained`` stamps
``is_loaded_in_4bit``; a ``meta`` skeleton has no such marker, so PEFT fell
back to the generic ``lora.layer.Linear``, which still RUNS against a
``Linear4bit`` base but casts and accumulates differently. No crash, no
warning, a healthy-looking loss curve."""
def test_streamed_skeleton_gets_the_bnb_aware_lora_layer(self, tmp_path):
from peft import get_peft_model
from soup_cli.utils.layer_stream_runtime import build_meta_skeleton
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="bfloat16", quant="nf4")
for param in model.parameters():
param.requires_grad = False
assert _lora_class_name(get_peft_model(model, _tiny_lora())) == (
"peft.tuners.lora.bnb.Linear4bit"
)
def test_control_stripping_the_marker_flips_peft_to_the_generic_layer(self, tmp_path):
"""Without this control the test above proves nothing — it would pass
for any implementation that happens to produce a Linear4bit base."""
from peft import get_peft_model
from soup_cli.utils.layer_stream_runtime import build_meta_skeleton
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="bfloat16", quant="nf4")
del model.is_loaded_in_4bit # exactly what a meta skeleton lacks
for param in model.parameters():
param.requires_grad = False
assert _lora_class_name(get_peft_model(model, _tiny_lora())) == (
"peft.tuners.lora.layer.Linear"
)
def test_markers_match_what_from_pretrained_stamps(self, tmp_path):
from transformers.utils.quantization_config import QuantizationMethod
from soup_cli.utils.layer_stream_runtime import build_meta_skeleton
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="bfloat16", quant="nf4")
assert model.is_loaded_in_4bit is True
assert model.is_quantized is True
assert model.quantization_method == QuantizationMethod.BITS_AND_BYTES
def test_hf_quantizer_is_present_and_trainable(self, tmp_path):
"""``Trainer.__init__`` reads ``is_quantized and not
hf_quantizer.is_trainable`` as "cannot fine-tune this", then formats the
error from ``model.hf_quantizer.quantization_config.quant_method``. With
``is_quantized`` stamped but no quantizer, EVERY streaming run dies at
trainer construction with an AttributeError."""
from soup_cli.utils.layer_stream_runtime import build_meta_skeleton
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="bfloat16", quant="nf4")
assert model.hf_quantizer is not None
assert model.hf_quantizer.is_trainable is True
assert model.hf_quantizer.quantization_config.quant_method is not None
def test_streamed_quant_config_matches_the_resident_one(self):
"""The release's central claim is that streamed NF4 is bit-exact against
RESIDENT NF4, and that holds only while both sides quantise with the
same settings. `quant_menu.build_quantization_config_for_loader` is what
every resident 4-bit load in this repo uses; if the two ever drift, the
parity tests would keep passing against a reference nobody ships."""
from soup_cli.config.loader import load_config_from_string
from soup_cli.utils.layer_stream_runtime import build_nf4_config
from soup_cli.utils.quant_menu import build_quantization_config_for_loader
cfg = load_config_from_string(
"base: m\ntask: sft\ndata:\n train: d.jsonl\n"
"training:\n quantization: 4bit\n"
)
resident = build_quantization_config_for_loader(tcfg=cfg.training, base="m")
streamed = build_nf4_config("bfloat16")
assert streamed.bnb_4bit_quant_type == resident.bnb_4bit_quant_type
assert (
streamed.bnb_4bit_use_double_quant is resident.bnb_4bit_use_double_quant
)
assert streamed.load_in_4bit is resident.load_in_4bit
def test_unquantised_skeleton_is_untouched(self, tmp_path):
"""v0.72.0's bit-exact bf16 gates only stay valid if that path is
byte-identical to what it was."""
import bitsandbytes as bnb
from soup_cli.utils.layer_stream_runtime import build_meta_skeleton
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="float32")
assert not any(
isinstance(param, bnb.nn.Params4bit) for param in model.parameters()
)
assert getattr(model, "is_loaded_in_4bit", False) is False
class TestQuantisedLayerSuffixes:
def test_finds_exactly_the_decoder_linears(self, tmp_path):
from soup_cli.utils.layer_stream_runtime import (
build_meta_skeleton,
quantised_layer_suffixes,
)
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="bfloat16", quant="nf4")
assert quantised_layer_suffixes(model) == frozenset(
{
"self_attn.q_proj.weight",
"self_attn.k_proj.weight",
"self_attn.v_proj.weight",
"self_attn.o_proj.weight",
"mlp.gate_proj.weight",
"mlp.up_proj.weight",
"mlp.down_proj.weight",
}
)
def test_layernorms_stay_out(self, tmp_path):
from soup_cli.utils.layer_stream_runtime import (
build_meta_skeleton,
quantised_layer_suffixes,
)
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="bfloat16", quant="nf4")
found = quantised_layer_suffixes(model)
assert not any("layernorm" in name for name in found)
def test_lm_head_is_never_quantised(self, tmp_path):
"""lm_head lives in extras and stays at the base dtype — quantising it
would both break the tied-embedding restore and cost accuracy where it
hurts most."""
import bitsandbytes as bnb
from soup_cli.utils.layer_stream_runtime import build_meta_skeleton
weights, _, _ = _tiny_llama_dir(tmp_path, tie=False)
model = build_meta_skeleton(weights, dtype="bfloat16", quant="nf4")
assert not isinstance(model.lm_head.weight, bnb.nn.Params4bit)
def test_unquantised_skeleton_reports_no_suffixes(self, tmp_path):
from soup_cli.utils.layer_stream_runtime import (
build_meta_skeleton,
quantised_layer_suffixes,
)
weights, _, _ = _tiny_llama_dir(tmp_path)
model = build_meta_skeleton(weights, dtype="float32")
assert quantised_layer_suffixes(model) == frozenset()
# ==========================================================================
# the streaming runtime
# ==========================================================================
def _nf4_stream(tmp_path, n_layers=2, tie=True, buffers=2, device="cpu", dtype="float32"):
"""Shard a tiny Llama as NF4 and build the streamed model over it.
``quant_device`` is pinned to the training device on purpose: CPU and CUDA
agree on the packed nibbles and absmax, but float32 double-quant nested
statistics differ by a reduction order, so a CUDA-quantised shard set
compared against a CPU-resident reference would fail for a reason that has
nothing to do with streaming.
"""
from soup_cli.utils.layer_shard import QUANT_NF4, shard_checkpoint
from soup_cli.utils.layer_stream_runtime import (
build_meta_skeleton,
build_streamed_model,
quantised_layer_suffixes,
)
weights, resident, _ = _tiny_llama_dir(tmp_path, n_layers=n_layers, tie=tie)
probe = build_meta_skeleton(weights, dtype=dtype, quant=QUANT_NF4)
suffixes = quantised_layer_suffixes(probe)
del probe
shards = str(tmp_path / "shards")
index = shard_checkpoint(
weights,
shards,
dtype=dtype,
arch="llama",
quant=QUANT_NF4,
quant_suffixes=suffixes,
quant_device=device,
)
model, runtime = build_streamed_model(
model_id=weights,
shard_dir=shards,
index=index,
lora_config=_tiny_lora(),
device=device,
dtype=dtype,
buffers=buffers,
pin=False,
seed=3,
quant=QUANT_NF4,
)
return model, runtime, weights, index, shards
def _resident_nf4(weights, dtype="float32", device="cpu"):
import torch
from peft import get_peft_model
from transformers import AutoModelForCausalLM
from soup_cli.utils.layer_stream_runtime import build_nf4_config
model = AutoModelForCausalLM.from_pretrained(
weights,
quantization_config=build_nf4_config(dtype),
dtype=getattr(torch, dtype),
device_map={"": device},
)
model.config.use_cache = False
for param in model.parameters():
param.requires_grad = False
return get_peft_model(model, _tiny_lora())
def _randomise_lora_b(model, seed=7):
"""PEFT initialises ``lora_B = 0``, so at step 0 the adapter contributes
NOTHING and every "did the adapter path run?" assertion passes vacuously —
including a bit-exactness check against a model whose adapter is also
doing nothing. The gate hit exactly this. Make B load-bearing first."""
import torch
generator = torch.Generator().manual_seed(seed)
with torch.no_grad():
for name, param in model.named_parameters():
if "lora_B" in name:
param.copy_(
torch.randn(param.shape, generator=generator).to(
param.device, param.dtype
)
* 0.02
)
def _sync_adapters(dst, src):
"""Copy LoRA weights src -> dst across the ``.inner.`` wrapper difference.
Returns the number copied; 0 means the comparison would be vacuous.
"""
import torch
def norm(key):
return key.replace(".inner.", ".")
source = {
norm(k): v.detach().clone() for k, v in src.state_dict().items() if "lora_" in k
}
copied = 0
with torch.no_grad():
for key, tensor in dst.state_dict().items():
if "lora_" not in key:
continue
match = source.get(norm(key))
if match is not None:
tensor.copy_(match.to(tensor.device, tensor.dtype))
copied += 1
return copied
class TestNF4SpecFromShard:
def test_per_tensor_dtype_comes_from_the_header(self, tmp_path):
"""An NF4 shard is deliberately MIXED-dtype: packed nibbles are uint8,
absmax is uint8 under double quant, the nested absmax and offset are
float32, and layernorms stay at the base dtype. Forcing one dtype
across the pool would reinterpret the packed bytes as floats."""
from soup_cli.utils.layer_shard import (
ABSMAX_SUFFIX,
NESTED_ABSMAX_SUFFIX,
NESTED_OFFSET_SUFFIX,
)
from soup_cli.utils.layer_stream_runtime import RamSource
_, _, _, _, shards = _nf4_stream(tmp_path)
spec = RamSource.spec_from_shard(shards)
key = "mlp.down_proj.weight"
assert spec[key][1] == "uint8"
assert spec[key + ABSMAX_SUFFIX][1] == "uint8"
assert spec[key + NESTED_ABSMAX_SUFFIX][1] == "float32"
assert spec[key + NESTED_OFFSET_SUFFIX][1] == "float32"
assert spec["input_layernorm.weight"][1] == "float32"
def test_zero_dim_offset_keeps_its_shape(self, tmp_path):
from soup_cli.utils.layer_shard import NESTED_OFFSET_SUFFIX
from soup_cli.utils.layer_stream_runtime import RamSource
_, _, _, _, shards = _nf4_stream(tmp_path)
spec = RamSource.spec_from_shard(shards)
assert spec["mlp.down_proj.weight" + NESTED_OFFSET_SUFFIX][0] == ()
class TestNF4StreamedModel:
def test_decoder_weights_never_leave_meta(self, tmp_path):
"""If they materialise, it is not streaming — it is a resident load
with extra steps, and the whole VRAM claim evaporates."""
model, _, _, _, _ = _nf4_stream(tmp_path)
meta = [n for n, p in model.named_parameters() if p.is_meta]
assert meta, "no decoder weights left on meta"
assert all(".layers." in name for name in meta)
def test_adapters_are_real_not_meta(self, tmp_path):
"""PEFT creates adapters on the base layer's device, which here is
``meta`` — the optimizer accepts them and the run trains nothing."""
model, _, _, _, _ = _nf4_stream(tmp_path)
trainable = [
(n, p) for n, p in model.named_parameters() if p.requires_grad and "lora_" in n
]
assert trainable
assert not any(p.is_meta for _, p in trainable)
def test_pool_streams_the_quant_sidecars(self, tmp_path):
from soup_cli.utils.layer_shard import ABSMAX_SUFFIX
_, runtime, _, _, _ = _nf4_stream(tmp_path)
pooled = set(runtime.pool.buffers[0])
assert "mlp.down_proj.weight" in pooled
assert "mlp.down_proj.weight" + ABSMAX_SUFFIX in pooled
def test_rebuilt_view_shares_storage_with_the_pooled_buffer(self, tmp_path):
"""``Params4bit`` is rebuilt on EVERY forward and every recompute. If
the constructor ever copies — a ``.clone()``, a ``.to()``, a dtype
promotion — streaming quietly allocates a whole layer per call and the
bounded-VRAM claim is gone, with nothing failing to show it."""
from soup_cli.utils.layer_stream_runtime import decoder_owner, rebuild_params4bit
model, runtime, _, index, _ = _nf4_stream(tmp_path)
layer = decoder_owner(model).layers[0]
buffers = runtime.pool.buffers[0]
key = "mlp.down_proj.weight"
view = rebuild_params4bit(key, buffers, index.quant_specs[key], layer.codes)
assert view.data_ptr() == buffers[key].data_ptr()
def test_store_is_far_smaller_than_bf16(self, tmp_path):
"""The point of the whole slot: ~0.52 bytes/param instead of 2."""
from soup_cli.utils.layer_shard import shard_checkpoint
from soup_cli.utils.layer_stream_runtime import RamSource
_, runtime, weights, _, _ = _nf4_stream(tmp_path)
plain_dir = str(tmp_path / "plain")
shard_checkpoint(weights, plain_dir, dtype="float32", arch="llama")
plain = RamSource(plain_dir, 2, RamSource.spec_from_shard(plain_dir), pin=False)
assert runtime.source.nbytes < plain.nbytes
class TestNF4BitExactVsResident:
"""Gate 1 check 1 + check 2, as CI tests. The reference is RESIDENT NF4,
not resident bf16 — the latter differs by quantisation error and would hide
a real bug inside it."""
def test_logits_are_bit_exact(self, tmp_path):
import torch
model, _, weights, _, _ = _nf4_stream(tmp_path)
resident = _resident_nf4(weights)
_randomise_lora_b(resident)
assert _sync_adapters(model, resident) > 0, "vacuous: no adapters copied"
ids = torch.randint(0, 64, (1, 16), generator=torch.Generator().manual_seed(1))
batch = {"input_ids": ids, "attention_mask": torch.ones_like(ids)}
model.eval()
resident.eval()
with torch.no_grad():
mine = model(**batch).logits.float()
theirs = resident(**batch).logits.float()
diff = (mine - theirs).abs().max().item()
assert diff == 0.0, diff
def test_layer_zero_adapter_receives_gradient(self, tmp_path):
"""plan P2: a ``detach()``/``no_grad()`` anywhere in the base forward
severs the graph. Lower adapters then never train while the loss still
falls, because the upper ones still learn."""
import torch
model, _, weights, _, _ = _nf4_stream(tmp_path)
resident = _resident_nf4(weights)
_randomise_lora_b(resident)
_sync_adapters(model, resident)
ids = torch.randint(0, 64, (1, 16), generator=torch.Generator().manual_seed(2))
model.train()
model(input_ids=ids, attention_mask=torch.ones_like(ids), labels=ids).loss.backward()
grads = {}
for name, param in model.named_parameters():
if not (param.requires_grad and "lora_" in name and ".layers." in name):
continue
idx = int(name.split(".layers.")[1].split(".")[0])
value = 0.0 if param.grad is None else param.grad.abs().max().item()
grads[idx] = max(grads.get(idx, 0.0), value)
assert grads[0] > 0.0, grads
assert all(value > 0.0 for value in grads.values()), grads
def test_loss_curves_match_resident(self, tmp_path):
import torch
model, _, weights, _, _ = _nf4_stream(tmp_path)
resident = _resident_nf4(weights)
_randomise_lora_b(resident)
_sync_adapters(model, resident)
ids = torch.randint(0, 64, (1, 16), generator=torch.Generator().manual_seed(3))
batch = {"input_ids": ids, "attention_mask": torch.ones_like(ids), "labels": ids}
def run(target):
torch.manual_seed(0)
opt = torch.optim.AdamW(
[p for p in target.parameters() if p.requires_grad], lr=1e-3
)
out = []
target.train()
for _ in range(5):
loss = target(**batch).loss
loss.backward()
opt.step()
opt.zero_grad(set_to_none=True)
out.append(loss.item())
return out
assert run(model) == run(resident)
class TestNF4StaleCacheIsRefused:
def test_missing_code_table_names_the_fix(self, tmp_path):
"""A v0.72.0 bf16 cache has no code tables. Rebuilding a QuantState
against ``None`` would raise somewhere deep inside bitsandbytes."""
import pytest as _pytest
from soup_cli.utils.layer_shard import NF4WeightSpec
from soup_cli.utils.layer_stream_runtime import rebuild_quant_state
spec = NF4WeightSpec(
shape=(4, 4),
dtype="float32",
blocksize=64,
quant_type="nf4",
nested=False,
nested_blocksize=0,
)
with _pytest.raises(ValueError, match="reshard"):
rebuild_quant_state("w", {}, spec, {})
class TestAdapterRoundTripUnderNF4:
"""v0.72.1's property must survive the NF4 rewrite: an adapter trained on
the streamed path has to load into a NORMAL model."""
def test_saved_keys_carry_no_inner_segment(self, tmp_path):
from peft import get_peft_model_state_dict
model, _, _, _, _ = _nf4_stream(tmp_path)
keys = list(get_peft_model_state_dict(model))
assert keys
assert not [key for key in keys if ".inner." in key], keys[:4]
def test_adapter_reloads_into_a_plain_nf4_model(self, tmp_path):
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM
from soup_cli.utils.layer_stream_runtime import build_nf4_config
model, _, weights, _, _ = _nf4_stream(tmp_path)
_randomise_lora_b(model)
out = str(tmp_path / "adapter")
model.save_pretrained(out)
base = AutoModelForCausalLM.from_pretrained(
weights,
quantization_config=build_nf4_config("float32"),
dtype=torch.float32,
device_map={"": "cpu"},
)
reloaded = PeftModel.from_pretrained(base, out)
landed = {
name: param
for name, param in reloaded.named_parameters()
if "lora_B" in name
}
assert landed
assert any(param.abs().max().item() > 0 for param in landed.values()), (
"every lora_B is zero — the adapter was dropped on reload, which "
"raises no exception"
)
# ==========================================================================
# schema
# ==========================================================================
def _stream_yaml(**overrides):
base = {
"quantization": "4bit",
"extra": "",
}
base.update(overrides)
return f"""
base: meta-llama/Llama-3.1-8B
task: sft
backend: transformers
modality: text
data:
train: data.jsonl
format: alpaca
training:
batch_size: 1
gradient_accumulation_steps: 1
quantization: {base["quantization"]}
stream_layers: true
{base["extra"]} lora:
r: 16
"""
class TestSchemaAcceptsNF4:
def test_four_bit_streaming_now_parses(self):
from soup_cli.config.loader import load_config_from_string
cfg = load_config_from_string(_stream_yaml(quantization="4bit"))
assert cfg.training.stream_layers is True
assert cfg.training.quantization == "4bit"
def test_none_still_parses(self):
"""Control — the v0.72.0 bf16 path must not have been traded away."""
from soup_cli.config.loader import load_config_from_string
cfg = load_config_from_string(_stream_yaml(quantization="none"))
assert cfg.training.quantization == "none"
@pytest.mark.parametrize("quant", ["8bit", "gptq", "bitnet_1.58"])
def test_other_quantisations_are_refused_naming_the_supported_set(self, quant):
from soup_cli.config.loader import load_config_from_string
with pytest.raises(Exception, match="4bit"):
load_config_from_string(_stream_yaml(quantization=quant))
def test_refusal_names_stream_layers_not_something_else(self):
"""A pre-existing validator could reject 8bit for an unrelated reason
and this suite would never notice."""
from soup_cli.config.loader import load_config_from_string
with pytest.raises(Exception, match="stream_layers"):
load_config_from_string(_stream_yaml(quantization="8bit"))
# ==========================================================================
# the planner's NF4 arithmetic
# ==========================================================================
class TestNF4StoreEstimate:
def test_plain_dtype_is_unscaled(self):
from soup_cli.utils.layer_stream import estimate_stream_store_bytes
assert estimate_stream_store_bytes(1_000, dtype="bfloat16", quant="none") == 1_000
def test_nf4_is_about_a_quarter_of_bf16(self):
from soup_cli.utils.layer_stream import estimate_stream_store_bytes
got = estimate_stream_store_bytes(2_000_000, dtype="bfloat16", quant="nf4")
# 0.516 bytes/param vs 2 -> 0.258
assert 0.25 < got / 2_000_000 < 0.27
def test_single_quant_is_larger_than_double_quant(self):
from soup_cli.utils.layer_stream import estimate_stream_store_bytes
double = estimate_stream_store_bytes(
2_000_000, dtype="bfloat16", quant="nf4", double_quant=True
)
single = estimate_stream_store_bytes(
2_000_000, dtype="bfloat16", quant="nf4", double_quant=False
)
assert single > double
def test_an_8b_bf16_checkpoint_fits_this_box_under_nf4(self):
"""THE reason the estimate exists. The pre-flight probe compares the
base against free RAM *before* sharding. Measuring an 8B checkpoint at
its 16.1 GB bf16 on-disk size would refuse the run outright — i.e.
refuse precisely the headline this slot exists to deliver."""
from soup_cli.utils.layer_stream import (
RAM_TIER_HEADROOM,
estimate_stream_store_bytes,
)
on_disk = 16_060_000_000 # Llama-3.1-8B bf16 safetensors
free_ram = 16_900_000_000 # this box
assert on_disk >= free_ram * RAM_TIER_HEADROOM # the naive probe refuses
nf4 = estimate_stream_store_bytes(on_disk, dtype="bfloat16", quant="nf4")
assert nf4 < free_ram * RAM_TIER_HEADROOM # the NF4-aware one does not
def test_unknown_quant_is_refused(self):
from soup_cli.utils.layer_stream import estimate_stream_store_bytes
with pytest.raises(ValueError, match="nf4"):
estimate_stream_store_bytes(1_000, dtype="bfloat16", quant="int8")
class TestPreflightUsesTheStreamedSize:
"""The estimate exists to stop the pre-flight refusing an 8B run because
its bf16 files are 16 GB on disk. Unit-testing the arithmetic alone would
not catch a regression that simply stops passing ``quant=`` through."""
def _run(
self, tmp_path, monkeypatch, *, quantization, free_ram, on_disk,
stream_source="auto", disk_kind="nvme",
):
from soup_cli.config.loader import load_config_from_string
from soup_cli.trainer.sft import SFTTrainerWrapper
weights, _, _ = _tiny_llama_dir(tmp_path, n_layers=2)
_write_tiny_tokenizer(weights)
monkeypatch.chdir(tmp_path)
monkeypatch.setenv("SOUP_LAYER_STREAM_CACHE_DIR", str(tmp_path / "cache"))
monkeypatch.setattr(
"soup_cli.utils.spectrum_scan.resolve_model_weights", lambda *_a, **_k: weights
)
monkeypatch.setattr(
"soup_cli.utils.layer_stream.free_ram_bytes", lambda: free_ram
)
monkeypatch.setattr(
"soup_cli.utils.layer_shard.source_weight_bytes", lambda *_a, **_k: on_disk
)
# Pinned, not probed: the real media type differs between the dev box
# (NVMe) and a CI runner (often "unknown"), and an environment-dependent
# tier decision would make these assertions flaky rather than wrong.
monkeypatch.setattr(
"soup_cli.utils.layer_stream.detect_disk_kind", lambda *_a, **_k: disk_kind
)
cfg = load_config_from_string(
f"""
base: {weights}
task: sft
backend: transformers
modality: text
data:
train: data.jsonl
format: alpaca
training:
batch_size: 1
gradient_accumulation_steps: 1
quantization: {quantization}
stream_layers: true
stream_source: {stream_source}
lora:
r: 4
target_modules: [q_proj, v_proj]
"""
)
wrapper = SFTTrainerWrapper(cfg)
wrapper.device = "cpu"
wrapper._setup_streaming_transformers(cfg, cfg.training)
return wrapper
# 10 GB free * 0.7 headroom = a 7 GB budget. A 16 GB bf16 checkpoint blows
# it; the same weights as NF4 (~4.1 GB) do not.
FREE_RAM = 10_000_000_000
ON_DISK = 16_000_000_000
def test_nf4_run_is_allowed_where_the_raw_size_would_refuse(
self, tmp_path, monkeypatch
):
wrapper = self._run(
tmp_path,
monkeypatch,
quantization="4bit",
free_ram=self.FREE_RAM,
on_disk=self.ON_DISK,
stream_source="ram",
)
assert wrapper.model.is_loaded_in_4bit is True
def test_control_the_same_size_unquantised_is_still_refused(
self, tmp_path, monkeypatch
):
"""Without this control the test above proves nothing — a pre-flight
that never refuses anything would satisfy it.
Scoped to `stream_source: ram` from v0.72.3 on, because that is where
the early size probe is decisive: under the new `auto` default a base
too large for RAM becomes a tier decision rather than an error. The
property under test is unchanged — the probe applies `quant=`.
"""
with pytest.raises(ValueError, match="stream_source='ram'"):
self._run(
tmp_path,
monkeypatch,
quantization="none",
free_ram=self.FREE_RAM,
on_disk=self.ON_DISK,
stream_source="ram",
)
def test_nf4_still_refuses_when_even_quantised_it_will_not_fit(
self, tmp_path, monkeypatch
):
"""And the NF4 branch must not become a blanket bypass.
Asserted through `stream_source: ram`, because under v0.72.3's default
`auto` a base too large for RAM is a tier decision rather than an error.
The property being pinned is unchanged: the estimate is applied to the
NF4 size, it is not skipped."""
with pytest.raises(ValueError, match="once quantised to NF4"):
self._run(
tmp_path,
monkeypatch,
quantization="4bit",
free_ram=self.FREE_RAM,
on_disk=self.ON_DISK * 20,
stream_source="ram",
)
class TestReportedParameterCount:
"""Found by the step-6 smoke, not by any unit test: a real
``training.stream_layers`` run with ``quantization: 4bit`` printed
"878,154,048 total" for SmolLM2-135M (true count 134,515,008), while the
RESIDENT NF4 path printed 134,975,808. PEFT special-cases ``Params4bit`` as
``numel * 2 * quant_storage.itemsize`` — correct for a resident one, whose
numel is the packed count, but not for a ``meta`` placeholder still carrying
the logical shape. At 8B that misprint would read ~52B."""
def test_runtime_carries_the_true_source_count(self, tmp_path):
_, runtime, _, index, _ = _nf4_stream(tmp_path)
assert runtime.total_params == index.total_params
assert runtime.stats()["total_params"] == index.total_params
def test_reported_total_is_the_real_parameter_count(self, tmp_path, monkeypatch):
"""End-to-end through the trainer, against the model's own config."""
from soup_cli.config.loader import load_config_from_string
from soup_cli.trainer.sft import SFTTrainerWrapper
weights, resident, config = _tiny_llama_dir(tmp_path, n_layers=2)
_write_tiny_tokenizer(weights)
monkeypatch.chdir(tmp_path)
monkeypatch.setenv("SOUP_LAYER_STREAM_CACHE_DIR", str(tmp_path / "cache"))
monkeypatch.setattr(
"soup_cli.utils.spectrum_scan.resolve_model_weights", lambda *_a, **_k: weights
)
cfg = load_config_from_string(
f"""
base: {weights}
task: sft
backend: transformers
modality: text
data:
train: data.jsonl
format: alpaca
training:
batch_size: 1
gradient_accumulation_steps: 1
quantization: 4bit
stream_layers: true
lora:
r: 4
target_modules: [q_proj, v_proj]
"""
)
wrapper = SFTTrainerWrapper(cfg)
wrapper.device = "cpu"
wrapper._setup_streaming_transformers(cfg, cfg.training)
# tie_word_embeddings=True, so the checkpoint omits lm_head.weight and
# the sharder's count matches the model's distinct parameters.
real = sum(p.numel() for p in resident.parameters())
reported = wrapper._stream_runtime.total_params
assert reported == real, (reported, real)
# and the PEFT number it replaces really is wrong, or this guards nothing
_trainable, peft_total = wrapper.model.get_nb_trainable_parameters()
assert peft_total > real * 2
class TestDeviceMapValue:
"""``accelerate.prepare_model`` does ``torch.device(value).index`` for a
4-bit model; a bare ``"cuda"`` yields None and the next line raises a
TypeError naming nothing actionable. Pure string handling — no GPU needed."""
@pytest.mark.parametrize(
"device,expected", [("cuda:1", 1), ("cuda:0", 0), ("cpu", "cpu"), ("meta", "meta")]
)
def test_indexed_devices_become_ints(self, device, expected):
from soup_cli.utils.layer_stream_runtime import _device_map_value
assert _device_map_value(device) == expected
def test_bare_cuda_gets_an_index(self):
from soup_cli.utils.layer_stream_runtime import _device_map_value
value = _device_map_value("cuda")
assert isinstance(value, int)
# ==========================================================================
# trainer wiring
# ==========================================================================
class TestTrainerWiring:
def test_streaming_setup_threads_nf4_end_to_end(self, tmp_path, monkeypatch):
"""The integration gap v0.72.0 shipped a CRITICAL through: no test built
the model the way ``soup train`` actually builds it."""
import torch
from soup_cli.config.loader import load_config_from_string
from soup_cli.trainer.sft import SFTTrainerWrapper
weights, _, _ = _tiny_llama_dir(tmp_path, n_layers=2)
_write_tiny_tokenizer(weights)
monkeypatch.chdir(tmp_path)
monkeypatch.setenv("SOUP_LAYER_STREAM_CACHE_DIR", str(tmp_path / "cache"))
monkeypatch.setattr(
"soup_cli.utils.spectrum_scan.resolve_model_weights", lambda *_a, **_k: weights
)
cfg = load_config_from_string(
f"""
base: {weights}
task: sft
backend: transformers
modality: text
data:
train: data.jsonl
format: alpaca
training:
batch_size: 1
gradient_accumulation_steps: 1
quantization: 4bit
stream_layers: true
lora:
r: 4
target_modules: [q_proj, v_proj]
"""
)
wrapper = SFTTrainerWrapper(cfg)
wrapper.device = "cpu"
wrapper._setup_streaming_transformers(cfg, cfg.training)
assert wrapper.model.is_loaded_in_4bit is True
assert any(p.is_meta for _, p in wrapper.model.named_parameters())
ids = torch.randint(0, 64, (1, 8))
out = wrapper.model(input_ids=ids, attention_mask=torch.ones_like(ids), labels=ids)
assert torch.isfinite(out.loss)
class TestNF4EndToEndSetup:
"""The v0.72.0 CRITICAL (``Trainer.__init__`` moving meta weights) hid in
exactly this gap: unit tests stopped at ``model(input_ids=...)`` while the
real path is ``wrapper.setup(dataset)`` -> TRL SFTTrainer -> ``train()``.
NF4 changes model construction, so it needs its own pass through it."""
def _wrapper(self, tmp_path, monkeypatch):
import yaml
from soup_cli.config.loader import load_config_from_string
from soup_cli.trainer.sft import SFTTrainerWrapper
weights, _, _ = _tiny_llama_dir(tmp_path, n_layers=2)
_write_tiny_tokenizer(weights)
monkeypatch.setenv("SOUP_LAYER_STREAM_CACHE_DIR", str(tmp_path / "cache"))
monkeypatch.chdir(tmp_path)
cfg = load_config_from_string(
yaml.safe_dump(
{
"base": weights,
"task": "sft",
"backend": "transformers",
"modality": "text",
"data": {
"train": "train.jsonl",
"max_length": 64,
"chat_template": "chatml",
},
"training": {
"batch_size": 1,
"gradient_accumulation_steps": 1,
"quantization": "4bit",
"stream_layers": True,
"epochs": 1,
"logging_steps": 1,
"save_steps": 1000,
"gradient_checkpointing": True,
"lora": {"r": 4, "alpha": 8, "target_modules": ["q_proj", "v_proj"]},
},
"output": str(tmp_path / "out"),
}
)
)
dataset = {
"train": [
{
"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "hello world"},
]
}
for _ in range(4)
]
}
device = "cuda" if _cuda() else "cpu"
return SFTTrainerWrapper(cfg, device=device), dataset
def test_setup_builds_a_real_trl_trainer_under_nf4(self, tmp_path, monkeypatch):
wrapper, dataset = self._wrapper(tmp_path, monkeypatch)
wrapper.setup(dataset)
assert wrapper.trainer is not None
assert wrapper.model.is_loaded_in_4bit is True
@skip_on_mps
def test_one_nf4_training_step_actually_runs(self, tmp_path, monkeypatch):
wrapper, dataset = self._wrapper(tmp_path, monkeypatch)
wrapper.setup(dataset)
wrapper.trainer.args.max_steps = 1
wrapper.trainer.train()
assert wrapper._stream_runtime.pool.loads > 0, "no layer was ever streamed"
@skip_on_mps
def test_the_saved_adapter_is_canonical(self, tmp_path, monkeypatch):
"""v0.72.1's fix has to survive NF4 all the way through TRL's save."""
from safetensors.torch import load_file
wrapper, dataset = self._wrapper(tmp_path, monkeypatch)
wrapper.setup(dataset)
wrapper.trainer.args.max_steps = 1
wrapper.trainer.train()
out = str(tmp_path / "saved")
wrapper.model.save_pretrained(out)
keys = list(load_file(os.path.join(out, "adapter_model.safetensors")))
assert keys
assert not [key for key in keys if ".inner." in key], keys[:4]
@requires_cuda
class TestNF4ParityOnCuda:
"""The CPU tests above prove the mechanism; this proves it on the device
the feature actually ships for, where bitsandbytes uses different kernels."""
def test_logits_bit_exact_vs_resident_nf4_on_cuda(self, tmp_path):
import torch
model, _, weights, _, _ = _nf4_stream(
tmp_path, device="cuda", dtype="bfloat16"
)
resident = _resident_nf4(weights, dtype="bfloat16", device=0)
_randomise_lora_b(resident)
assert _sync_adapters(model, resident) > 0, "vacuous: no adapters copied"
ids = torch.randint(
0, 64, (1, 16), generator=torch.Generator().manual_seed(11)
).cuda()
batch = {"input_ids": ids, "attention_mask": torch.ones_like(ids)}
model.eval()
resident.eval()
with torch.no_grad():
diff = (
(model(**batch).logits.float() - resident(**batch).logits.float())
.abs()
.max()
.item()
)
assert diff == 0.0, diff
def test_same_seed_twice_is_identical_on_cuda(self, tmp_path):
"""A prefetch race shows up here and nowhere else — the CPU path has no
separate stream to race with."""
import torch
ids = torch.randint(
0, 64, (1, 16), generator=torch.Generator().manual_seed(12)
).cuda()
def run():
model, _, _, _, _ = _nf4_stream(tmp_path, device="cuda", dtype="bfloat16")
torch.manual_seed(0)
opt = torch.optim.AdamW(
[p for p in model.parameters() if p.requires_grad], lr=1e-3
)
out = []
model.train()
for _ in range(3):
loss = model(
input_ids=ids, attention_mask=torch.ones_like(ids), labels=ids
).loss
loss.backward()
opt.step()
opt.zero_grad(set_to_none=True)
out.append(loss.item())
return out
assert run() == run()
# ==========================================================================
# import hygiene
# ==========================================================================
class TestNoTopLevelTorch:
@pytest.mark.parametrize(
"module", ["layer_stream", "layer_shard", "layer_stream_runtime"]
)
def test_module_has_no_top_level_training_import(self, module):
"""An AST guard proves a SYNTACTIC property only — the runtime
authority is tests/test_cli_startup_is_light.py. Kept as a cheap first
line, exactly as CLAUDE.md prescribes."""
import ast
import pathlib
import soup_cli
path = pathlib.Path(soup_cli.__file__).parent / "utils" / f"{module}.py"
tree = ast.parse(path.read_text(encoding="utf-8"))
banned = {"torch", "bitsandbytes", "transformers", "peft", "safetensors", "accelerate"}
for node in tree.body:
if isinstance(node, ast.Import):
for alias in node.names:
assert alias.name.split(".")[0] not in banned, alias.name
elif isinstance(node, ast.ImportFrom) and node.module:
assert node.module.split(".")[0] not in banned, node.module
assert sys.version_info >= (3, 10)
Reference in New Issue View Git Blame Copy Permalink