mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 14:21:06 -10:00
DFlash2 pack with the Q4_0 selector, made on the person's computer at setup / Load from the checked Q4_0 pack and used when it verifies
The selector's two codebooks as Q4_0 rows take 174 MiB (about 88 expert slots) less graphics memory than bf16, with the same drafts: the committed text was byte-identical to the bf16 pack's in 40 of 40 requests (greedy and sampled, everything resident), draft time and acceptance equal, and ahead 1.7-2.2% where the cache is short (RESULTS-DF2MEM.md section 15). The idea and the bytes are nazeshinjite's: the community GGUF nazeshinjite/Qwen3.8-Flash-Next-DFlash2-GGUF (Apache-2.0) stores the two tables as Q4_0, and tools/dflash2_pack.py --q4_0-selector / --convert-selector writes exactly those bytes (selector_predecessor sha256 55f70049...ffd28, selector_successor 3b596669...a72d2 over the tensor bytes, equal to the GGUF's). drafting_fetch.convert_selector (called by setup, strata drafting fetch and Load, never by a start or a serve; no network; about 120 MB of memory) checks the Q4_0 pack against PACK_SHA256, converts it in-process and checks the result against CONVERTED_BYTES / CONVERTED_SHA256 (339694592 bytes, 4a42b982...3ecd); anything else is deleted and the bf16-selector pack keeps running. drafting.pack_path returns the converted pack only when it verifies (a marker beside it keeps the hash check off the start path) and the engine program carries the text 'selector codebooks: Q4_0 rows'; a missing, short or corrupt file, an engine without the text or an engine that is not known gets the Q4_0 pack. Includes f-0291-df2-comm (Q4_0 rows in the select kernel, STRATA_DF2_HEAD=native) merged into rc1.
This commit is contained in:
1 parent
f4561cbe02
commit
ebac5686f9
6 files changed
+406
-3
No files matched your search
@@ -1359,6 +1359,10 @@ Strata itself: [MIT](../LICENSE). The model files are not part of it; their lice
|
||||
image encoder (`tools/vision/`), and `gguf-py` used by the tools. See `third_party/ggml/LICENSE`.
|
||||
- Ideas from [Splash](https://github.com/incoai/splash), [ninfer](https://github.com/Neroued/ninfer) and
|
||||
[HyperQwen](https://github.com/syv-ai/HyperQwen); references in the paper.
|
||||
- The DFlash2 drafter's Q4_0 selector (174 MiB less graphics memory, same drafts and same text) is nazeshinjite's idea:
|
||||
[nazeshinjite/Qwen3.8-Flash-Next-DFlash2-GGUF](https://huggingface.co/nazeshinjite/Qwen3.8-Flash-Next-DFlash2-GGUF) (Apache-2.0) stores the drafter's
|
||||
two selector tables as Q4_0; Strata converts the checked pack on your computer (`tools/dflash2_pack.py --convert-selector`) and checks the bytes
|
||||
against the same pin.
|
||||
- The web app's font: [Outfit](https://github.com/Outfitio/Outfit-Fonts) (SIL Open Font License 1.1, see `serve/web/fonts/OFL.txt`). Its
|
||||
Monitor tab started from @code-martin's dashboard idea (PR #22).
|
||||
Math is drawn by [KaTeX](https://katex.org) (MIT, vendored in `serve/web/vendor/katex`). The chat page's wide-window layout, svg/html preview, math, Load/Unload, saved
|
||||
|
||||
@@ -1541,6 +1541,8 @@ class Runner:
|
||||
want: dict = {}
|
||||
try:
|
||||
want = self.drafting_inputs(read_config(base_path))
|
||||
if want["data_dir"] is not None and drafting.needs_conversion(**want):
|
||||
drafting_fetch.convert_selector(want["data_dir"]) # Load, not a start: the smaller selector (174 MiB), made here from the checked pack, no network
|
||||
if want["data_dir"] is None or not drafting.needs_pack(**want):
|
||||
return
|
||||
self.emit("model", "Downloading the DFlash2 drafter (1.2 GB, once)…", 0)
|
||||
|
||||
@@ -22,6 +22,8 @@ engine never switches drafter mid-run.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import mmap
|
||||
import os
|
||||
import re
|
||||
@@ -197,6 +199,15 @@ PACK_GB = 0.52 # what stays on disk (the source file is removed af
|
||||
PACK_NAME = "qwen3.8-flash-next-dflash2.q4_0.pack"
|
||||
PACK_BYTES = 522458112 # the Q4_0 pack tools/dflash2_pack.py --q4_0 makes from the pinned safetensors file, and its SHA-256
|
||||
PACK_SHA256 = "cda40d62d0eb8f37ed6e272142641a1d2f0de4a32f9e14d0416e7a71534849f1"
|
||||
# The same pack with the selector's two codebooks as Q4_0 rows (174 MiB less in graphics memory, ~88 expert slots; draft ms, acceptance and the committed text equal: RESULTS-DF2MEM.md
|
||||
# section 15). The idea and the bytes are nazeshinjite's: the community GGUF https://huggingface.co/nazeshinjite/Qwen3.8-Flash-Next-DFlash2-GGUF (Apache-2.0) stores the codebooks
|
||||
# as Q4_0, and tools/dflash2_pack.py --q4_0-selector / --convert-selector writes exactly those bytes (selector_predecessor sha256 55f70049...ffd28, selector_successor
|
||||
# 3b596669...a72d2 over the tensor bytes; checked equal to the GGUF's). We host nothing: setup / Load converts the verified Q4_0 pack above on this computer
|
||||
# (drafting_fetch.convert_selector, no network, a second or two); the engine reads either pack and says "selector codebooks: Q4_0 rows" when it loads the converted one.
|
||||
CONVERTED_NAME = "qwen3.8-flash-next-dflash2.q4_0-selq4.pack"
|
||||
CONVERTED_BYTES = 339694592
|
||||
CONVERTED_SHA256 = "4a42b98220ea7880664002c8edda1563aecbae754b762a2e2f31b2b91ee93ecd"
|
||||
_MARK_SELQ4 = b"selector codebooks: Q4_0 rows" # text of an engine that reads such a pack (src/core/dflash2.cpp)
|
||||
LEGACY_PACK_NAME = "qwen3.8-flash-next-dflash2.q8.pack" # what Strata 0.2.9 test builds packed: still used when it is the one on disk
|
||||
VRAM_GB = DRAFTER_GB["dflash2"] # graphics memory with the MTP layer not loaded: 890 MiB (RESULTS-DF2MEM.md); the engine reports "drafter 799 MiB" for the pack and buffers alone
|
||||
FOLDER = "dflash2" # under the model data folder, next to ``mtp``
|
||||
@@ -251,8 +262,78 @@ def target_path(data_dir: str | Path) -> Path:
|
||||
return Path(data_dir) / FOLDER / PACK_NAME
|
||||
|
||||
|
||||
def pack_path(data_dir: str | Path) -> Path:
|
||||
"""The pack the engine is started with: the Q4_0 pack, else a Q8 pack an earlier Strata made (no download for it), else where the Q4_0 pack will be."""
|
||||
def converted_path(data_dir: str | Path) -> Path:
|
||||
"""Where setup / Load writes the pack with the Q4_0 selector (beside the Q4_0 pack it was made from)."""
|
||||
return Path(data_dir) / FOLDER / CONVERTED_NAME
|
||||
|
||||
|
||||
def sha256_of(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open("rb") as stream:
|
||||
for block in iter(lambda: stream.read(1 << 22), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def converted_ok(data_dir: str | Path | None, *, size: int | None = None, sha256: str | None = None, name: str | None = None) -> bool:
|
||||
"""Is the converted pack here and exactly the pinned bytes? A file of another size is refused at once; the SHA-256 (about a second) is checked when the
|
||||
file is new to us, and the answer is kept in a small marker beside it (size, time, hash) so a start reads one tiny file, not 340 MB. A missing, short,
|
||||
changed or corrupt file is simply "no": the Q4_0 pack with the bf16 selector is used. Reads only; nothing is written when the marker cannot be."""
|
||||
if not data_dir:
|
||||
return False
|
||||
want_size = CONVERTED_BYTES if size is None else size
|
||||
want_sha = CONVERTED_SHA256 if sha256 is None else sha256
|
||||
path = Path(data_dir) / FOLDER / (name or CONVERTED_NAME)
|
||||
try:
|
||||
stat = path.stat()
|
||||
if not path.is_file() or stat.st_size != want_size:
|
||||
return False
|
||||
marker = path.with_name(path.name + ".ok")
|
||||
try:
|
||||
known = json.loads(marker.read_text(encoding="utf-8"))
|
||||
if known.get("size") == stat.st_size and known.get("mtime_ns") == stat.st_mtime_ns and known.get("sha256") == want_sha:
|
||||
return True
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
if sha256_of(path) != want_sha:
|
||||
return False
|
||||
try:
|
||||
marker.write_text(json.dumps({"size": stat.st_size, "mtime_ns": stat.st_mtime_ns, "sha256": want_sha}), encoding="utf-8")
|
||||
except OSError:
|
||||
pass
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
_selq4: dict[tuple[str, int, int], bool] = {}
|
||||
|
||||
|
||||
def engine_reads_selector_q4(exe: str | os.PathLike[str] | None) -> bool:
|
||||
"""Does the engine program at `exe` read a pack with Q4_0 selector codebooks? Read from the file's text (it is never started for this), once per path, size
|
||||
and time. An engine that is not known (no path) or that lacks the text gets the bf16-selector pack: an older engine would read the rows as garbage."""
|
||||
if not exe:
|
||||
return False
|
||||
try:
|
||||
path = Path(exe)
|
||||
stat = path.stat()
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
key = (str(path), stat.st_size, stat.st_mtime_ns)
|
||||
if key not in _selq4:
|
||||
try:
|
||||
with open(path, "rb") as handle, mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ) as data:
|
||||
_selq4[key] = data.find(_MARK_SELQ4) >= 0
|
||||
except (OSError, ValueError):
|
||||
_selq4[key] = False
|
||||
return _selq4[key]
|
||||
|
||||
|
||||
def pack_path(data_dir: str | Path, engine: str | os.PathLike[str] | None = None) -> Path:
|
||||
"""The pack the engine is started with: the converted pack (Q4_0 selector) when it is here, verified, and `engine` is a program that reads it; else the Q4_0 pack, else
|
||||
a Q8 pack an earlier Strata made (no download for it), else where the Q4_0 pack will be."""
|
||||
if engine is not None and converted_ok(data_dir) and engine_reads_selector_q4(engine):
|
||||
return converted_path(data_dir)
|
||||
for name in (PACK_NAME, LEGACY_PACK_NAME):
|
||||
if _here(Path(data_dir) / FOLDER / name):
|
||||
return Path(data_dir) / FOLDER / name
|
||||
@@ -537,6 +618,13 @@ def wants_one_request(mode: Any, *, klass: str = "other", model_name: str | None
|
||||
return choice_of(mode, model_name=model_name, family=family, backend=backend, archs=archs) == "dflash2"
|
||||
|
||||
|
||||
def needs_conversion(mode: Any, *, klass: str = "other", model_name: str | None = None, family: str | None = None, data_dir: str | Path | None = None,
|
||||
backend: str | None = None, archs: Any = None, **_ignored: Any) -> bool:
|
||||
"""Would this choice use DFlash2, is the Q4_0 pack here, and is the converted pack (Q4_0 selector) not made yet? (Then setup / Load makes it, locally, once.)"""
|
||||
return (choice_of(mode, model_name=model_name, family=family, backend=backend, archs=archs) == "dflash2" and bool(data_dir)
|
||||
and _here(Path(data_dir) / FOLDER / PACK_NAME) and not converted_ok(data_dir))
|
||||
|
||||
|
||||
def needs_pack(mode: Any, *, klass: str = "other", model_name: str | None = None, family: str | None = None, data_dir: str | Path | None = None,
|
||||
backend: str | None = None, archs: Any = None, **_ignored: Any) -> bool:
|
||||
"""Would this choice use the DFlash2 drafter if its file were here, and is it not here yet? (Then it is fetched before the start, once.)"""
|
||||
@@ -647,7 +735,7 @@ def apply(cfg: dict[str, Any], *, klass: str = "other", data_dir: str | Path | N
|
||||
cfg["mtp"] = True # an explicit drafter choice overrides an older "mtp": false
|
||||
else:
|
||||
base = strip_mtp_flags(strip_draft_flags(args))
|
||||
flags = list(engine_flags(pack_path(data_dir), vocab, form))
|
||||
flags = list(engine_flags(pack_path(data_dir, engine), vocab, form))
|
||||
kept_conf = _flag_value(mine, "--draft-conf")
|
||||
flags[flags.index("--draft-conf") + 1] = kept_conf if kept_conf is not None else f"{conf_for(backend, slots_of(base, vram_mib)):g}"
|
||||
if _flag_value(mine, "--spec") is None:
|
||||
|
||||
@@ -99,6 +99,47 @@ def _present(folders: list[Path], spec) -> dict[str, Path]:
|
||||
return found
|
||||
|
||||
|
||||
def convert_selector(data_dir: str | Path, *, progress: Progress | None = None, source_bytes: int | None = None, source_sha256: str | None = None,
|
||||
size: int | None = None, sha256: str | None = None) -> Path | None:
|
||||
"""Make the pack with the Q4_0 selector (drafting.CONVERTED_NAME, 174 MiB less in graphics memory) from the verified Q4_0 pack, on this computer. No network, a second
|
||||
or two on a fast machine, about 120 MB of memory. Called only by setup, `strata drafting fetch` and Load (launch), never by a start or a serve. The Q4_0 pack must be
|
||||
the pinned bytes (PACK_BYTES / PACK_SHA256) and the result the pinned bytes (CONVERTED_BYTES / CONVERTED_SHA256); anything else is deleted and None is returned:
|
||||
the Q4_0 pack with the bf16 selector keeps being used, nothing is reported as an error. The idea and the bytes are nazeshinjite's
|
||||
(https://huggingface.co/nazeshinjite/Qwen3.8-Flash-Next-DFlash2-GGUF). The keyword arguments are for tests, which convert a small synthetic pack."""
|
||||
report = progress or (lambda *_: None)
|
||||
want_size = drafting.CONVERTED_BYTES if size is None else size
|
||||
want_sha = drafting.CONVERTED_SHA256 if sha256 is None else sha256
|
||||
if drafting.converted_ok(data_dir, size=want_size, sha256=want_sha):
|
||||
return drafting.converted_path(data_dir)
|
||||
src = drafting.target_path(data_dir)
|
||||
out = drafting.converted_path(data_dir)
|
||||
tmp = out.with_suffix(".tmp")
|
||||
try:
|
||||
if not src.is_file():
|
||||
return None
|
||||
pinned_bytes = drafting.PACK_BYTES if source_bytes is None else source_bytes
|
||||
pinned_sha = drafting.PACK_SHA256 if source_sha256 is None else source_sha256
|
||||
if src.stat().st_size != pinned_bytes or transfer.sha256_file(src) != pinned_sha:
|
||||
return None # not the pack we measured: leave it as it is
|
||||
report("Preparing the DFlash2 drafter's smaller selector…", 0, 1)
|
||||
_packer().convert_selector_q4(str(src), str(tmp), say=lambda *_: None)
|
||||
if tmp.stat().st_size != want_size or transfer.sha256_file(tmp) != want_sha:
|
||||
tmp.unlink(missing_ok=True)
|
||||
return None
|
||||
os.replace(tmp, out)
|
||||
out.with_name(out.name + ".ok").unlink(missing_ok=True)
|
||||
if not drafting.converted_ok(data_dir, size=want_size, sha256=want_sha): # writes the marker
|
||||
return None
|
||||
report("The DFlash2 drafter's smaller selector is ready.", 1, 1)
|
||||
return out
|
||||
except Exception: # noqa: BLE001 - never a reason to stop a setup or a launch: the bf16-selector pack runs
|
||||
try:
|
||||
tmp.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def ensure(data_dir: str | Path, *, progress: Progress | None = None, cancel: threading.Event | None = None,
|
||||
base_url: str | None = None, files=None) -> Path:
|
||||
"""Make the packed drafter exist under ``data_dir`` and return its path. A pack that is there is kept. Resumable: an interrupted
|
||||
@@ -107,6 +148,8 @@ def ensure(data_dir: str | Path, *, progress: Progress | None = None, cancel: th
|
||||
out = drafting.target_path(data_dir)
|
||||
if drafting.pack_ready(data_dir):
|
||||
drafting.provide_vocab(data_dir)
|
||||
if files is None:
|
||||
convert_selector(data_dir, progress=report)
|
||||
report("The DFlash2 drafter is ready.", 1, 1)
|
||||
return drafting.pack_path(data_dir)
|
||||
acquisition.require("the DFlash2 drafter", "strata drafting fetch")
|
||||
@@ -157,5 +200,7 @@ def ensure(data_dir: str | Path, *, progress: Progress | None = None, cancel: th
|
||||
except OSError:
|
||||
pass
|
||||
drafting.provide_vocab(data_dir) # the draft vocabulary comes with the pack: an install that never had the MTP layer needs it
|
||||
if files is None:
|
||||
convert_selector(data_dir, progress=report) # the smaller selector (174 MiB), made here from the pack just checked; the bf16-selector pack stays beside it
|
||||
report("The DFlash2 drafter is ready.", total, total)
|
||||
return out
|
||||
@@ -0,0 +1,186 @@
|
||||
"""The DFlash2 pack with the Q4_0 selector (174 MiB less graphics memory): made on this computer from the checked Q4_0 pack, used only when it verifies and the engine reads it.
|
||||
|
||||
The idea and the bytes are nazeshinjite's (https://huggingface.co/nazeshinjite/Qwen3.8-Flash-Next-DFlash2-GGUF, Apache-2.0): the community GGUF stores the selector's two
|
||||
codebooks as Q4_0, and tools/dflash2_pack.py --q4_0-selector writes exactly those bytes (checked against the GGUF with the real checkpoint: RESULTS-DF2MEM.md section 15).
|
||||
These tests use the small synthetic checkpoint of tools/test_dflash2_pack.py: the real pins cannot be exercised without the 1.2 GB download, so the pinned numbers are
|
||||
checked for consistency and the mechanism with pins of the synthetic pack.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from strata_control import drafting, drafting_fetch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
VOCAB_REAL, RANK_REAL = 248320, 256
|
||||
|
||||
|
||||
def _load(name, rel):
|
||||
spec = importlib.util.spec_from_file_location(name, ROOT / rel)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules.setdefault(name, module)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class PinnedNumbers(unittest.TestCase):
|
||||
def test_the_converted_pack_is_the_q4_pack_less_the_codebooks_difference(self):
|
||||
"""CONVERTED_BYTES follows from PACK_BYTES by arithmetic alone: two codebooks of 248,320 x 256 go from bf16 (2 bytes a value) to Q4_0 (18 bytes per 32 values)."""
|
||||
bf16 = 2 * VOCAB_REAL * RANK_REAL * 2
|
||||
q4 = 2 * VOCAB_REAL * (RANK_REAL // 32) * 18
|
||||
self.assertEqual(drafting.PACK_BYTES - drafting.CONVERTED_BYTES, bf16 - q4)
|
||||
self.assertEqual(bf16 - q4, 182763520)
|
||||
self.assertEqual(len(drafting.CONVERTED_SHA256), 64)
|
||||
self.assertNotEqual(drafting.CONVERTED_SHA256, drafting.PACK_SHA256)
|
||||
|
||||
|
||||
class SmallPacks(unittest.TestCase):
|
||||
"""A synthetic checkpoint packed as the Q4_0 pack, then converted."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.synth = _load("t_pack_sel", "tools/test_dflash2_pack.py")
|
||||
cls.packer = _load("strata_dflash2_pack_sel", "tools/dflash2_pack.py")
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.dir = Path(self.tmp.name)
|
||||
self.synth.write_safetensors(self.dir / "m.safetensors", self.synth.tensors())
|
||||
(self.dir / "c.json").write_text(json.dumps(self.synth.CFG))
|
||||
self.data = self.dir / "data"
|
||||
self.q4 = drafting.target_path(self.data)
|
||||
self.q4.parent.mkdir(parents=True)
|
||||
self.packer.pack(str(self.dir / "m.safetensors"), str(self.dir / "c.json"), str(self.q4), say=lambda *_: None, q4=True)
|
||||
self.src_pins = {"source_bytes": self.q4.stat().st_size, "source_sha256": hashlib.sha256(self.q4.read_bytes()).hexdigest()}
|
||||
direct = self.dir / "direct.pack"
|
||||
self.packer.pack(str(self.dir / "m.safetensors"), str(self.dir / "c.json"), str(direct), say=lambda *_: None, q4=True, q4_selector=True)
|
||||
self.pins = {"size": direct.stat().st_size, "sha256": hashlib.sha256(direct.read_bytes()).hexdigest(), **self.src_pins}
|
||||
self.direct = direct
|
||||
|
||||
def convert(self, **extra):
|
||||
return drafting_fetch.convert_selector(self.data, **{**self.pins, **extra})
|
||||
|
||||
def test_the_conversion_reproduces_the_pinned_bytes(self):
|
||||
"""The pack converted from the finished Q4_0 pack is byte-identical to the one packed straight from the checkpoint with --q4_0-selector (the pinned hash)."""
|
||||
out = self.convert()
|
||||
self.assertEqual(out, drafting.converted_path(self.data))
|
||||
self.assertEqual(out.read_bytes(), self.direct.read_bytes())
|
||||
self.assertLess(out.stat().st_size, self.q4.stat().st_size)
|
||||
self.assertTrue(drafting.converted_ok(self.data, size=self.pins["size"], sha256=self.pins["sha256"]))
|
||||
self.assertTrue(self.q4.is_file(), "the Q4_0 pack stays: it is the fallback")
|
||||
|
||||
def test_every_tensor_but_the_two_codebooks_is_carried_over_unchanged(self):
|
||||
out = self.convert()
|
||||
a, b = self.q4.read_bytes(), out.read_bytes()
|
||||
|
||||
def table(blob):
|
||||
n = int.from_bytes(blob[8:12], "little")
|
||||
rows = {}
|
||||
for line in blob[12:12 + n].decode().split("\n"):
|
||||
w = line.split()
|
||||
if w and w[0] == "t":
|
||||
rows[w[1]] = (w[2], int(w[5]), int(w[6]))
|
||||
return rows, 12 + n
|
||||
ta, base_a = table(a)
|
||||
tb, base_b = table(b)
|
||||
self.assertEqual(set(ta), set(tb))
|
||||
books = {"candidate_selector.predecessor_codebook", "candidate_selector.successor_codebook"}
|
||||
for name in ta:
|
||||
ba = a[base_a + ta[name][1]: base_a + ta[name][1] + ta[name][2]]
|
||||
bb = b[base_b + tb[name][1]: base_b + tb[name][1] + tb[name][2]]
|
||||
if name in books:
|
||||
self.assertEqual((ta[name][0], tb[name][0]), ("bf16", "q4_0"))
|
||||
self.assertEqual(len(bb) * 32, len(ba) * 9, name) # 18 bytes per 32 values against 2 bytes a value
|
||||
else:
|
||||
self.assertEqual(ba, bb, name)
|
||||
|
||||
def test_a_missing_corrupt_or_short_converted_pack_falls_back_to_the_q4_pack(self):
|
||||
eng = self.dir / "engine"
|
||||
eng.write_bytes(b"\0" * 64 + drafting._MARK_SELQ4 + b"\0" * 64)
|
||||
self.assertEqual(drafting.pack_path(self.data, eng), self.q4, "not made yet")
|
||||
out = self.convert()
|
||||
with mock.patch.object(drafting, "CONVERTED_BYTES", self.pins["size"]), mock.patch.object(drafting, "CONVERTED_SHA256", self.pins["sha256"]):
|
||||
self.assertEqual(drafting.pack_path(self.data, eng), out)
|
||||
blob = bytearray(out.read_bytes())
|
||||
blob[len(blob) // 2] ^= 0xFF # one flipped byte, the same size, and the marker still describes the old file
|
||||
out.write_bytes(bytes(blob))
|
||||
self.assertEqual(drafting.pack_path(self.data, eng), self.q4, "a corrupt file (the marker no longer matches it) is not used")
|
||||
out.write_bytes(bytes(blob)[:-1000])
|
||||
self.assertEqual(drafting.pack_path(self.data, eng), self.q4, "a short file is not used")
|
||||
out.unlink()
|
||||
self.assertEqual(drafting.pack_path(self.data, eng), self.q4, "a missing file is not used")
|
||||
|
||||
def test_a_q4_pack_that_is_not_the_pinned_one_is_left_alone(self):
|
||||
self.assertIsNone(self.convert(source_sha256="0" * 64))
|
||||
self.assertFalse(drafting.converted_path(self.data).exists())
|
||||
self.assertTrue(self.q4.is_file())
|
||||
|
||||
def test_a_result_that_does_not_match_the_pin_is_deleted(self):
|
||||
self.assertIsNone(self.convert(sha256="f" * 64))
|
||||
self.assertFalse(drafting.converted_path(self.data).exists())
|
||||
self.assertFalse(drafting.converted_path(self.data).with_suffix(".tmp").exists())
|
||||
|
||||
def test_a_failing_conversion_is_never_an_error_for_the_caller(self):
|
||||
with mock.patch.object(drafting_fetch, "_packer", side_effect=RuntimeError("no numpy")):
|
||||
self.assertIsNone(self.convert())
|
||||
self.assertFalse(drafting.converted_path(self.data).exists())
|
||||
|
||||
def test_the_engine_must_read_the_pack_before_it_is_used(self):
|
||||
out = self.convert()
|
||||
old = self.dir / "old-engine"
|
||||
old.write_bytes(b"\0" * 256)
|
||||
new = self.dir / "new-engine"
|
||||
new.write_bytes(b"\0" * 64 + drafting._MARK_SELQ4 + b"\0" * 64)
|
||||
with mock.patch.object(drafting, "CONVERTED_BYTES", self.pins["size"]), mock.patch.object(drafting, "CONVERTED_SHA256", self.pins["sha256"]):
|
||||
self.assertEqual(drafting.pack_path(self.data), self.q4, "no engine named: the pack every engine reads")
|
||||
self.assertEqual(drafting.pack_path(self.data, old), self.q4, "an engine without the text would read the rows as garbage")
|
||||
self.assertEqual(drafting.pack_path(self.data, self.dir / "none"), self.q4)
|
||||
self.assertEqual(drafting.pack_path(self.data, new), out)
|
||||
|
||||
def test_auto_arguments_point_at_the_converted_pack(self):
|
||||
self.convert()
|
||||
new = self.dir / "new-engine"
|
||||
new.write_bytes(b"\0" * 64 + drafting._MARK_SELQ4 + b"\0" * 64 + b"--draft-fallback" + b"--mtp-draft-vocab FILE (or --mtp DIR")
|
||||
vocab = self.data / drafting.FOLDER / drafting.VOCAB_NAME
|
||||
vocab.write_bytes(b"\1\0\0\0" * 8)
|
||||
cfg = {"model_name": "qwen3.8-flash-next-iq3_s", "drafting": "dflash2", "args": ["--spec", "4", "--expert-cache", "1000"]}
|
||||
with mock.patch.object(drafting, "CONVERTED_BYTES", self.pins["size"]), mock.patch.object(drafting, "CONVERTED_SHA256", self.pins["sha256"]), \
|
||||
mock.patch.object(drafting, "ENGINE_SERVES_DFLASH2", True), mock.patch.object(drafting, "pack_ready", return_value=True):
|
||||
drafting.apply(cfg, klass="blackwell", data_dir=self.data, backend="cuda", archs=["sm_120"], engine=new)
|
||||
args = cfg["args"]
|
||||
self.assertEqual(args[args.index("--draft-model") + 1], str(drafting.converted_path(self.data)))
|
||||
cfg2 = {"model_name": "qwen3.8-flash-next-iq3_s", "drafting": "dflash2", "args": ["--spec", "4", "--expert-cache", "1000"]}
|
||||
drafting.apply(cfg2, klass="blackwell", data_dir=self.data, backend="cuda", archs=["sm_120"], engine=None)
|
||||
self.assertEqual(cfg2["args"][cfg2["args"].index("--draft-model") + 1], str(self.q4), "an engine that is not known gets the pack every engine reads")
|
||||
|
||||
def test_needs_conversion_only_when_dflash2_is_chosen_and_the_q4_pack_is_here(self):
|
||||
want = dict(mode="dflash2", model_name="qwen3.8-flash-next-iq3_s", data_dir=self.data, backend="cuda", archs=["sm_120"])
|
||||
with mock.patch.object(drafting, "ENGINE_SERVES_DFLASH2", True), mock.patch.object(drafting, "_here", lambda p: Path(p).is_file()):
|
||||
self.assertTrue(drafting.needs_conversion(**want))
|
||||
self.assertFalse(drafting.needs_conversion(**{**want, "mode": "mtp"}))
|
||||
self.assertFalse(drafting.needs_conversion(**{**want, "data_dir": self.dir / "empty"}))
|
||||
self.convert()
|
||||
with mock.patch.object(drafting, "CONVERTED_BYTES", self.pins["size"]), mock.patch.object(drafting, "CONVERTED_SHA256", self.pins["sha256"]):
|
||||
self.assertFalse(drafting.needs_conversion(**want), "already made")
|
||||
|
||||
def test_setup_and_load_convert_a_start_does_not(self):
|
||||
"""ensure() (strata drafting fetch, setup) converts a pack that is already here, with no network; the serving side only reads."""
|
||||
with mock.patch.object(drafting_fetch, "convert_selector") as convert, mock.patch.object(drafting, "pack_ready", return_value=True):
|
||||
drafting_fetch.ensure(self.data)
|
||||
convert.assert_called_once()
|
||||
source = (ROOT / "serve" / "server.py").read_text(encoding="utf-8")
|
||||
self.assertNotIn("convert_selector", source)
|
||||
self.assertNotIn("convert_selector", (ROOT / "strata_cli" / "launch" / "plan.py").read_text(encoding="utf-8") if (ROOT / "strata_cli" / "launch" / "plan.py").exists() else "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -2,6 +2,7 @@
|
||||
"""Pack the incoai DFlash2 drafter for Qwen3.8-Flash-Next (safetensors) into the engine's runtime file.
|
||||
|
||||
python tools/dflash2_pack.py model.safetensors config.json out.pack [--q4_0] [--q4_0-selector]
|
||||
python tools/dflash2_pack.py --convert-selector in.pack out.pack (a finished Q4_0 pack: its bf16 selector codebooks become Q4_0, nothing else changes)
|
||||
|
||||
Format (little endian): "DF2P", u32 version (1), u32 index_len, a text index (padded with spaces so the data starts on a 4096 multiple), then
|
||||
the tensors, each 256-byte aligned at its offset (from the end of the index). Index lines: "cfg <key> <value>", "t <name> <dtype> <rows>
|
||||
@@ -179,7 +180,84 @@ def pack(st_path, cfg_path, out_path, say=print, q4=False, q4_selector=False):
|
||||
return base + off
|
||||
|
||||
|
||||
def convert_selector_q4(src_path, dst_path, say=print, rows_per_step=16384):
|
||||
"""Rewrite a finished DF2P pack with the selector's two codebooks as Q4_0 rows: every other tensor's bytes are carried over unchanged, the two bf16 codebooks
|
||||
(tensors `candidate_selector.predecessor_codebook` / `successor_codebook`, [vocab x 256]) are widened exactly to float32 and quantized by quant_q4_0, which is
|
||||
ggml's quantize_row_q4_0_ref. The result is byte-identical to `pack(..., q4=True, q4_selector=True)` from the same safetensors file (and the codebooks are the
|
||||
bytes of the community GGUF nazeshinjite/Qwen3.8-Flash-Next-DFlash2-GGUF, selector_predecessor / selector_successor). Memory stays under ~300 MB: the codebooks go
|
||||
through `rows_per_step` rows at a time. Returns the size written. Needs only numpy, no network."""
|
||||
with open(src_path, "rb") as fh:
|
||||
if fh.read(4) != b"DF2P":
|
||||
raise ValueError("not a DF2P pack")
|
||||
version, index_len = struct.unpack("<II", fh.read(8))
|
||||
if version != 1:
|
||||
raise ValueError("unknown DF2P version %d" % version)
|
||||
index = fh.read(index_len).decode()
|
||||
base = 12 + index_len
|
||||
cfg_lines, tensors = [], []
|
||||
for line in index.split("\n"):
|
||||
w = line.split()
|
||||
if not w or w[0] == "end":
|
||||
continue
|
||||
if w[0] == "cfg":
|
||||
cfg_lines.append(line)
|
||||
elif w[0] == "t":
|
||||
tensors.append({"name": w[1], "dtype": w[2], "rows": int(w[3]), "cols": int(w[4]), "offset": int(w[5]), "bytes": int(w[6])})
|
||||
books = ("candidate_selector.predecessor_codebook", "candidate_selector.successor_codebook")
|
||||
for t in tensors:
|
||||
if t["name"] in books:
|
||||
if t["dtype"] != "bf16":
|
||||
raise ValueError("%s is already %s" % (t["name"], t["dtype"]))
|
||||
if t["cols"] % 32:
|
||||
raise ValueError("%s: %d columns" % (t["name"], t["cols"]))
|
||||
off = 0
|
||||
for t in tensors:
|
||||
if t["name"] in books:
|
||||
t["new_bytes"] = t["rows"] * (t["cols"] // 32) * 18
|
||||
t["new_dtype"] = "q4_0"
|
||||
else:
|
||||
t["new_bytes"], t["new_dtype"] = t["bytes"], t["dtype"]
|
||||
off = (off + 255) & ~255
|
||||
t["new_offset"] = off
|
||||
off += t["new_bytes"]
|
||||
lines = list(cfg_lines)
|
||||
for t in tensors:
|
||||
lines.append("t %s %s %d %d %d %d" % (t["name"], t["new_dtype"], t["rows"], t["cols"], t["new_offset"], t["new_bytes"]))
|
||||
lines.append("end")
|
||||
js = ("\n".join(lines) + "\n").encode()
|
||||
js += b" " * ((-(len(js) + 12)) % 4096)
|
||||
with open(dst_path, "wb") as out:
|
||||
out.write(b"DF2P" + struct.pack("<II", 1, len(js)) + js)
|
||||
new_base = out.tell()
|
||||
for t in tensors:
|
||||
out.seek(new_base + t["new_offset"])
|
||||
fh.seek(base + t["offset"])
|
||||
if t["name"] not in books:
|
||||
left = t["bytes"]
|
||||
while left:
|
||||
block = fh.read(min(left, 1 << 24))
|
||||
if not block:
|
||||
raise ValueError("%s: the pack is shorter than its index" % t["name"])
|
||||
out.write(block)
|
||||
left -= len(block)
|
||||
continue
|
||||
for r0 in range(0, t["rows"], rows_per_step):
|
||||
n = min(rows_per_step, t["rows"] - r0)
|
||||
raw = np.frombuffer(fh.read(n * t["cols"] * 2), dtype=np.uint16)
|
||||
if raw.size != n * t["cols"]:
|
||||
raise ValueError("%s: the pack is shorter than its index" % t["name"])
|
||||
w = (raw.astype(np.uint32) << 16).view(np.float32).reshape(n, t["cols"])
|
||||
out.write(quant_q4_0(w))
|
||||
out.seek(new_base + off) # the last tensor ends the file (the pack writer's size is base + off too)
|
||||
out.truncate(new_base + off)
|
||||
say("wrote %s: %d tensors, %.1f MB" % (dst_path, len(tensors), (new_base + off) / 1e6))
|
||||
return new_base + off
|
||||
|
||||
|
||||
def main():
|
||||
if len(sys.argv) == 4 and sys.argv[1] == "--convert-selector":
|
||||
convert_selector_q4(sys.argv[2], sys.argv[3])
|
||||
return
|
||||
flags = ("--q4_0", "--q4_0-selector")
|
||||
args = [a for a in sys.argv[1:] if a not in flags]
|
||||
if len(args) < 3:
|
||||
|
||||
Reference in new issue
Block a user