🍷 Winery Strata: fork of Niko1221/Strata (Winery Cuvée family, wine-red UI)

bbb6388 verified 29 days ago

10.3 kB

"""ref/load.py - P1.S4: GGUF -> FP32 tensors on demand, with LAZY per-expert dequantization.
The plan is explicit: "experts are dequantized lazily per (layer, expert) - do not materialize 120B
parameters". That is the whole design constraint here, and it is what `expert()` below exists for: the
routed-expert tensors are fused `[in, out, experts]`, and touching one expert must not touch the other
511.
Dequantization itself is NOT reimplemented. `gguf-py` handles the twelve standard encodings and is
already the reference this project validates its C++ against; `Q2_0` (type 42) is the one gguf-py
cannot represent at all, so it is served by this project's own decoder. That split is deliberate: it
means a bug in `tools/gguf_writer.dequantize_q2_0` cannot hide behind gguf-py and vice versa.
    from ref.load import Gguf
    g = Gguf("path.gguf")
    w = g.tensor("blk.0.attn_gate.weight")      # (out, in) fp32, whole tensor
    e = g.expert("blk.0.ffn_gate_exps.weight", 7)  # (out, in) fp32, ONE expert only
"""
from __future__ import annotations

import pathlib
import re
import sys

import numpy as np

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent / "tools"))
from gguf_reader import GGUFFile                     # noqa: E402
from gguf_writer import dequantize_q2_0              # noqa: E402

# gguf-py names differ from our reader's names in a few places; map only what differs.
_GGUF_PY_NAME = {"Q4_0": "Q4_0", "Q8_0": "Q8_0", "Q4_K": "Q4_K", "Q5_K": "Q5_K", "Q6_K": "Q6_K",
                 "Q3_K": "Q3_K", "IQ4_NL": "IQ4_NL", "IQ4_XS": "IQ4_XS", "Q5_0": "Q5_0"}

# (elements per block, bytes per block) - mirrors block_geometry() in src/artifact/gguf_reader.cpp
GEOM = {"F32": (1, 4), "F16": (1, 2), "BF16": (1, 2),
        "Q4_0": (32, 18), "Q5_0": (32, 22), "Q8_0": (32, 34), "IQ4_NL": (32, 18),
        "Q3_K": (256, 110), "Q4_K": (256, 144), "Q5_K": (256, 176), "Q6_K": (256, 210),
        "IQ4_XS": (256, 136), "Q2_0": (64, 18)}


def _dequant_flat(type_name: str, raw: np.ndarray) -> np.ndarray:
    """One dequantization dispatch, shared by tensor()/expert()/rows().
    Q2_0 is this project's own decoder because gguf-py's enum cannot represent type 42 at all;
    everything else goes through gguf-py, which is already the reference this project validates its
    C++ against.  Keeping one dispatch means a type added in one path cannot be forgotten in another.
    """
    if type_name == "Q2_0":
        return dequantize_q2_0(raw.tobytes())
    if type_name == "BF16":
        return _bf16_to_f32(raw)
    if type_name == "F32":
        return raw.view(np.float32)
    if type_name == "F16":
        return raw.view(np.float16).astype(np.float32)
    from gguf import quants, GGMLQuantizationType as Q
    return quants.dequantize(raw, Q[_GGUF_PY_NAME[type_name]]).astype(np.float32)


def _bf16_to_f32(raw: np.ndarray) -> np.ndarray:
    return (raw.view(np.uint16).astype(np.uint32) << 16).view(np.float32)


def _discover_shards(path) -> list[pathlib.Path]:
    """Resolve a GGUF argument to its shard list.
    The real artifact is TWO files and the split is not incidental: shard 2 holds exactly one tensor,
    `per_layer_token_embd.weight` (26.82 GiB). A reader that opens one file can therefore load the
    whole model EXCEPT the PLE table and will not fail until a forward pass reaches layer 1 - so
    shard discovery is not a convenience, it is a correctness requirement for the real model.
    Accepts a directory (all `*.gguf` in it), any single shard (siblings found by the
    `-NNNNN-of-MMMMM` suffix), or an explicit list.
    """
    if isinstance(path, (list, tuple)):
        return [pathlib.Path(p) for p in path]
    p = pathlib.Path(path)
    if p.is_dir():
        # A directory is ambiguous in practice: the project root holds BOTH the model shards (in a
        # subdirectory) and an unrelated `mmproj-*.gguf`.  Globbing `*.gguf` there silently loaded
        # the multimodal projector instead of the model.  So prefer an explicit shard set, accept a
        # single unambiguous file, and REFUSE to guess otherwise.
        shards = sorted(p.glob("*-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9].gguf"))
        if shards:
            totals = {re.match(r"^.*-(\d{5})-of-(\d{5})\.gguf$", s.name).group(2) for s in shards}
            if len(totals) != 1:
                raise ValueError(f"{p}: shard files from more than one split: {sorted(totals)}")
            return shards
        found = sorted(p.glob("*.gguf"))
        if len(found) == 1:
            return found
        if not found:
            raise FileNotFoundError(f"no .gguf files in {p}")
        raise ValueError(
            f"{p}: {len(found)} .gguf files and none is a shard set, so which is the model is "
            f"ambiguous: {[f.name for f in found]}. Pass one file explicitly.")
    m = re.match(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", p.name)
    if m:
        stem, _, total = m.groups()
        sibs = sorted(p.parent.glob(f"{stem}-*-of-{total}.gguf"))
        if len(sibs) != int(total):
            raise FileNotFoundError(
                f"{p.name}: expected {total} shards, found {len(sibs)} in {p.parent}")
        return sibs
    return [p]


class Gguf:
    """Read-only, demand-paged view of a (possibly multi-shard) GGUF.
    Nothing is dequantized until asked for.  `g` is the FIRST shard's `GGUFFile`, so metadata and the
    tensor directory are reachable as before; reads route to whichever shard owns the tensor.
    """

    def __init__(self, path):
        self.paths = _discover_shards(path)
        self.path = self.paths[0]
        self.shards: list[tuple[GGUFFile, object]] = []
        self._by_name: dict[str, tuple[int, object]] = {}
        for i, sp in enumerate(self.paths):
            gf = GGUFFile(sp)
            self.shards.append((gf, sp.open("rb")))
            for t in gf.tensors:
                if t.name in self._by_name:
                    raise ValueError(f"{sp.name}: duplicate tensor {t.name} across shards")
                self._by_name[t.name] = (i, t)
        self.g = self.shards[0][0]           # metadata / directory / alignment live in shard 1

    # ---- raw bytes
    def _raw(self, name: str, nbytes: int, byte_offset: int = 0) -> bytes:
        i, t = self._by_name[name]
        gf, fh = self.shards[i]
        fh.seek(gf.data_start + t.offset + byte_offset)
        return fh.read(nbytes)

    def info(self, name: str):
        return self._by_name[name][1]

    def row_bytes(self, t) -> int:
        el, by = GEOM[t.type_name]
        return t.shape[0] // el * by

    # ---- full tensors
    def tensor(self, name: str) -> np.ndarray:
        """Dequantize a whole tensor to FP32, shape as GGUF order reversed to (out, in[, experts])."""
        t = self.info(name)
        el, by = GEOM[t.type_name]
        if t.elements % el:
            raise ValueError(f"{name}: {t.elements} not a multiple of block {el}")
        raw = np.frombuffer(self._raw(name, t.elements // el * by), dtype=np.uint8)
        flat = _dequant_flat(t.type_name, raw)
        # GGUF dim 0 varies fastest; numpy's last axis does, so reverse and transpose to (out, in).
        return flat.reshape(tuple(reversed(t.shape)))

    # ---- the lazy row path
    def rows(self, name: str, idx) -> np.ndarray:
        """Gather specific ROWS of a 2-D tensor without decoding the rest of it.
        This exists for `per_layer_token_embd.weight`, which is `[160, 320001536]` IQ4_NL = **26.82
        GiB** - 43% of the artifact, the entire second shard - and is read 16 rows per token by one
        layer (see round 98's finding). Calling `tensor()` on it would dequantize 320 million rows to
        FP32, about 205 GB, to use 2.5 KB of them.
        A row is contiguous in GGUF (ne0 varies fastest), so one row is `row_bytes` at
        `row * row_bytes`. The type's block size must divide the row width; for this table it does
        (160 = 5 x 32 for IQ4_NL).
        """
        t = self.info(name)
        if len(t.shape) != 2:
            raise ValueError(f"{name}: rows() needs a 2-D tensor, got shape {t.shape}")
        el, by = GEOM[t.type_name]
        width = t.shape[0]
        if width % el:
            raise ValueError(f"{name}: row width {width} is not a multiple of block {el}")
        rb = self.row_bytes(t)
        idx = np.asarray(idx, dtype=np.int64).ravel()
        if idx.size and (idx.min() < 0 or idx.max() >= t.shape[1]):
            raise ValueError(f"{name}: row index out of range (0..{t.shape[1] - 1})")
        out = np.empty((idx.size, width), dtype=np.float32)
        for k, r in enumerate(idx):
            raw = np.frombuffer(self._raw(name, rb, int(r) * rb), dtype=np.uint8)
            out[k] = _dequant_flat(t.type_name, raw)
        return out

    # ---- the lazy path
    def expert(self, name: str, expert_id: int) -> np.ndarray:
        """One expert out of a fused `[in, out, experts]` tensor, WITHOUT touching the others.
        In GGUF order the expert index is the LAST axis, so expert e's matrix occupies a contiguous
        run of `in*out` elements at offset e*in*out. Only that run is read and decoded - which is the
        plan's "do not materialize 120B parameters" requirement, and the reason the engine can run at
        all on a machine where the expert set is 34 GB.
        """
        t = self.info(name)
        if len(t.shape) != 3:
            raise ValueError(f"{name}: not a fused expert tensor (shape {t.shape})")
        in_f, out_f, n_exp = t.shape
        if not 0 <= expert_id < n_exp:
            raise ValueError(f"{name}: expert {expert_id} out of range (0..{n_exp - 1})")
        el, by = GEOM[t.type_name]
        per_expert = in_f * out_f
        if per_expert % el:
            raise ValueError(f"{name}: per-expert element count {per_expert} not a multiple of {el}")
        # byte offset of expert e: whole experts before it, in blocks
        start = expert_id * (per_expert // el) * by
        raw = np.frombuffer(self._raw(name, per_expert // el * by, start), dtype=np.uint8)
        flat = _dequant_flat(t.type_name, raw)
        return flat.reshape((out_f, in_f))

    def close(self) -> None:
        for _, fh in self.shards:
            fh.close()

	"""ref/load.py - P1.S4: GGUF -> FP32 tensors on demand, with LAZY per-expert dequantization.

	The plan is explicit: "experts are dequantized lazily per (layer, expert) - do not materialize 120B
	parameters". That is the whole design constraint here, and it is what `expert()` below exists for: the
	routed-expert tensors are fused `[in, out, experts]`, and touching one expert must not touch the other
	511.

	Dequantization itself is NOT reimplemented. `gguf-py` handles the twelve standard encodings and is
	already the reference this project validates its C++ against; `Q2_0` (type 42) is the one gguf-py
	cannot represent at all, so it is served by this project's own decoder. That split is deliberate: it
	means a bug in `tools/gguf_writer.dequantize_q2_0` cannot hide behind gguf-py and vice versa.

	from ref.load import Gguf
	g = Gguf("path.gguf")
	w = g.tensor("blk.0.attn_gate.weight") # (out, in) fp32, whole tensor
	e = g.expert("blk.0.ffn_gate_exps.weight", 7) # (out, in) fp32, ONE expert only
	"""
	from __future__ import annotations

	import pathlib
	import re
	import sys

	import numpy as np

	sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent / "tools"))
	from gguf_reader import GGUFFile # noqa: E402
	from gguf_writer import dequantize_q2_0 # noqa: E402

	# gguf-py names differ from our reader's names in a few places; map only what differs.
	_GGUF_PY_NAME = {"Q4_0": "Q4_0", "Q8_0": "Q8_0", "Q4_K": "Q4_K", "Q5_K": "Q5_K", "Q6_K": "Q6_K",
	"Q3_K": "Q3_K", "IQ4_NL": "IQ4_NL", "IQ4_XS": "IQ4_XS", "Q5_0": "Q5_0"}

	# (elements per block, bytes per block) - mirrors block_geometry() in src/artifact/gguf_reader.cpp
	GEOM = {"F32": (1, 4), "F16": (1, 2), "BF16": (1, 2),
	"Q4_0": (32, 18), "Q5_0": (32, 22), "Q8_0": (32, 34), "IQ4_NL": (32, 18),
	"Q3_K": (256, 110), "Q4_K": (256, 144), "Q5_K": (256, 176), "Q6_K": (256, 210),
	"IQ4_XS": (256, 136), "Q2_0": (64, 18)}


	def _dequant_flat(type_name: str, raw: np.ndarray) -> np.ndarray:
	"""One dequantization dispatch, shared by tensor()/expert()/rows().

	Q2_0 is this project's own decoder because gguf-py's enum cannot represent type 42 at all;
	everything else goes through gguf-py, which is already the reference this project validates its
	C++ against. Keeping one dispatch means a type added in one path cannot be forgotten in another.
	"""
	if type_name == "Q2_0":
	return dequantize_q2_0(raw.tobytes())
	if type_name == "BF16":
	return _bf16_to_f32(raw)
	if type_name == "F32":
	return raw.view(np.float32)
	if type_name == "F16":
	return raw.view(np.float16).astype(np.float32)
	from gguf import quants, GGMLQuantizationType as Q
	return quants.dequantize(raw, Q[_GGUF_PY_NAME[type_name]]).astype(np.float32)


	def _bf16_to_f32(raw: np.ndarray) -> np.ndarray:
	return (raw.view(np.uint16).astype(np.uint32) << 16).view(np.float32)


	def _discover_shards(path) -> list[pathlib.Path]:
	"""Resolve a GGUF argument to its shard list.

	The real artifact is TWO files and the split is not incidental: shard 2 holds exactly one tensor,
	`per_layer_token_embd.weight` (26.82 GiB). A reader that opens one file can therefore load the
	whole model EXCEPT the PLE table and will not fail until a forward pass reaches layer 1 - so
	shard discovery is not a convenience, it is a correctness requirement for the real model.

	Accepts a directory (all `*.gguf` in it), any single shard (siblings found by the
	`-NNNNN-of-MMMMM` suffix), or an explicit list.
	"""
	if isinstance(path, (list, tuple)):
	return [pathlib.Path(p) for p in path]
	p = pathlib.Path(path)
	if p.is_dir():
	# A directory is ambiguous in practice: the project root holds BOTH the model shards (in a
	# subdirectory) and an unrelated `mmproj-.gguf`. Globbing `.gguf` there silently loaded
	# the multimodal projector instead of the model. So prefer an explicit shard set, accept a
	# single unambiguous file, and REFUSE to guess otherwise.
	shards = sorted(p.glob("*-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9].gguf"))
	if shards:
	totals = {re.match(r"^.*-(\d{5})-of-(\d{5})\.gguf$", s.name).group(2) for s in shards}
	if len(totals) != 1:
	raise ValueError(f"{p}: shard files from more than one split: {sorted(totals)}")
	return shards
	found = sorted(p.glob("*.gguf"))
	if len(found) == 1:
	return found
	if not found:
	raise FileNotFoundError(f"no .gguf files in {p}")
	raise ValueError(
	f"{p}: {len(found)} .gguf files and none is a shard set, so which is the model is "
	f"ambiguous: {[f.name for f in found]}. Pass one file explicitly.")
	m = re.match(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", p.name)
	if m:
	stem, _, total = m.groups()
	sibs = sorted(p.parent.glob(f"{stem}-*-of-{total}.gguf"))
	if len(sibs) != int(total):
	raise FileNotFoundError(
	f"{p.name}: expected {total} shards, found {len(sibs)} in {p.parent}")
	return sibs
	return [p]


	class Gguf:
	"""Read-only, demand-paged view of a (possibly multi-shard) GGUF.

	Nothing is dequantized until asked for. `g` is the FIRST shard's `GGUFFile`, so metadata and the
	tensor directory are reachable as before; reads route to whichever shard owns the tensor.
	"""

	def __init__(self, path):
	self.paths = _discover_shards(path)
	self.path = self.paths[0]
	self.shards: list[tuple[GGUFFile, object]] = []
	self._by_name: dict[str, tuple[int, object]] = {}
	for i, sp in enumerate(self.paths):
	gf = GGUFFile(sp)
	self.shards.append((gf, sp.open("rb")))
	for t in gf.tensors:
	if t.name in self._by_name:
	raise ValueError(f"{sp.name}: duplicate tensor {t.name} across shards")
	self._by_name[t.name] = (i, t)
	self.g = self.shards[0][0] # metadata / directory / alignment live in shard 1

	# ---- raw bytes
	def _raw(self, name: str, nbytes: int, byte_offset: int = 0) -> bytes:
	i, t = self._by_name[name]
	gf, fh = self.shards[i]
	fh.seek(gf.data_start + t.offset + byte_offset)
	return fh.read(nbytes)

	def info(self, name: str):
	return self._by_name[name][1]

	def row_bytes(self, t) -> int:
	el, by = GEOM[t.type_name]
	return t.shape[0] // el * by

	# ---- full tensors
	def tensor(self, name: str) -> np.ndarray:
	"""Dequantize a whole tensor to FP32, shape as GGUF order reversed to (out, in[, experts])."""
	t = self.info(name)
	el, by = GEOM[t.type_name]
	if t.elements % el:
	raise ValueError(f"{name}: {t.elements} not a multiple of block {el}")
	raw = np.frombuffer(self._raw(name, t.elements // el * by), dtype=np.uint8)
	flat = _dequant_flat(t.type_name, raw)
	# GGUF dim 0 varies fastest; numpy's last axis does, so reverse and transpose to (out, in).
	return flat.reshape(tuple(reversed(t.shape)))

	# ---- the lazy row path
	def rows(self, name: str, idx) -> np.ndarray:
	"""Gather specific ROWS of a 2-D tensor without decoding the rest of it.

	This exists for `per_layer_token_embd.weight`, which is `[160, 320001536]` IQ4_NL = **26.82
	GiB** - 43% of the artifact, the entire second shard - and is read 16 rows per token by one
	layer (see round 98's finding). Calling `tensor()` on it would dequantize 320 million rows to
	FP32, about 205 GB, to use 2.5 KB of them.

	A row is contiguous in GGUF (ne0 varies fastest), so one row is `row_bytes` at
	`row * row_bytes`. The type's block size must divide the row width; for this table it does
	(160 = 5 x 32 for IQ4_NL).
	"""
	t = self.info(name)
	if len(t.shape) != 2:
	raise ValueError(f"{name}: rows() needs a 2-D tensor, got shape {t.shape}")
	el, by = GEOM[t.type_name]
	width = t.shape[0]
	if width % el:
	raise ValueError(f"{name}: row width {width} is not a multiple of block {el}")
	rb = self.row_bytes(t)
	idx = np.asarray(idx, dtype=np.int64).ravel()
	if idx.size and (idx.min() < 0 or idx.max() >= t.shape[1]):
	raise ValueError(f"{name}: row index out of range (0..{t.shape[1] - 1})")
	out = np.empty((idx.size, width), dtype=np.float32)
	for k, r in enumerate(idx):
	raw = np.frombuffer(self._raw(name, rb, int(r) * rb), dtype=np.uint8)
	out[k] = _dequant_flat(t.type_name, raw)
	return out

	# ---- the lazy path
	def expert(self, name: str, expert_id: int) -> np.ndarray:
	"""One expert out of a fused `[in, out, experts]` tensor, WITHOUT touching the others.

	In GGUF order the expert index is the LAST axis, so expert e's matrix occupies a contiguous
	run of `inout` elements at offset ein*out. Only that run is read and decoded - which is the
	plan's "do not materialize 120B parameters" requirement, and the reason the engine can run at
	all on a machine where the expert set is 34 GB.
	"""
	t = self.info(name)
	if len(t.shape) != 3:
	raise ValueError(f"{name}: not a fused expert tensor (shape {t.shape})")
	in_f, out_f, n_exp = t.shape
	if not 0 <= expert_id < n_exp:
	raise ValueError(f"{name}: expert {expert_id} out of range (0..{n_exp - 1})")
	el, by = GEOM[t.type_name]
	per_expert = in_f * out_f
	if per_expert % el:
	raise ValueError(f"{name}: per-expert element count {per_expert} not a multiple of {el}")
	# byte offset of expert e: whole experts before it, in blocks
	start = expert_id * (per_expert // el) * by
	raw = np.frombuffer(self._raw(name, per_expert // el * by, start), dtype=np.uint8)
	flat = _dequant_flat(t.type_name, raw)
	return flat.reshape((out_f, in_f))

	def close(self) -> None:
	for _, fh in self.shards:
	fh.close()