Release v0.4.0: selective MLX model downloads

3da88e3 verified 2 days ago

7.72 kB

"""Tensor-at-a-time safetensors loading; never load a whole checkpoint to RAM."""

from __future__ import annotations

import gc
import hashlib
import json
import struct
from collections import OrderedDict
from pathlib import Path

import mlx.core as mx
import numpy as np


class WeightStore:
    def __init__(self, path, *, cache_gb=0.0, cast=None, reader="mlx"):
        self.path = Path(path)
        files = (
            [self.path]
            if self.path.is_file()
            else sorted(self.path.glob("*.safetensors"))
        )
        if not files:
            raise FileNotFoundError(f"No safetensors weights at {self.path}")
        self.index = {}
        self.headers = {}
        self.data_start = {}
        self.cache = OrderedDict()
        self.cache_limit = int(cache_gb * 1024**3)
        self.cached_bytes = 0
        self.cast = cast
        if reader not in ("mlx", "memmap"):
            raise ValueError("reader must be mlx or memmap")
        self.reader = reader
        self.loads = 0
        self.bytes_read = 0
        self.files = files
        for file in files:
            with file.open("rb") as handle:
                raw = handle.read(8)
                if len(raw) != 8:
                    raise ValueError(f"Truncated safetensors header: {file}")
                length = struct.unpack("<Q", raw)[0]
                if length > 100_000_000:
                    raise ValueError(f"Invalid safetensors header size: {file}")
                header = json.loads(handle.read(length))
                self.data_start[file] = 8 + length
            for name, spec in header.items():
                if name == "__metadata__":
                    continue
                if name in self.index:
                    raise ValueError(f"Duplicate tensor {name} in {self.path}")
                self.index[name] = file
                self.headers[name] = spec

    def __contains__(self, key):
        return key in self.index

    def fingerprint(self):
        # Source identity for cache invalidation, not a content-integrity checksum.
        info = [
            (str(p.resolve()), p.stat().st_size, p.stat().st_mtime_ns)
            for p in self.files
        ]
        return hashlib.sha256(json.dumps(info).encode()).hexdigest()

    def get(self, key, *, rows=None, cache=True):
        if rows is None and key in self.cache:
            self.cache.move_to_end(key)
            return self.cache[key]
        if key not in self.index:
            raise KeyError(f"Missing checkpoint tensor {key} in {self.path}")
        # safetensors' Python MLX adapter routes BF16 through NumPy, which has
        # no BF16 dtype. Map its uint16 payload directly and bitcast in MLX.
        spec, file = self.headers[key], self.index[key]
        dtypes = {
            "BF16": "<u2",
            "F16": "<f2",
            "F32": "<f4",
            "U32": "<u4",
            "I32": "<i4",
            "I64": "<i8",
            "U8": "u1",
            "BOOL": "?",
        }
        if spec["dtype"] not in dtypes:
            raise ValueError(f"Unsupported tensor dtype {spec['dtype']} for {key}")
        if self.reader == "mlx" and rows is None:
            # MLX creates lazy per-tensor Load nodes. Discard the dictionary
            # immediately: retaining it would also retain every evaluated weight
            # in the shard, bypassing this store's explicit cache budget.
            value = mx.load(str(file))[key]
        else:
            # A full-vocabulary MLX Load is evaluated before a row gather.
            # Read requested rows directly from the safetensors mapping instead;
            # this preserves order/duplicates and the original BF16 bit payload.
            mapped = np.memmap(
                file,
                dtype=dtypes[spec["dtype"]],
                mode="r",
                offset=self.data_start[file] + spec["data_offsets"][0],
                shape=tuple(spec["shape"]),
            )
            value = mx.array(mapped if rows is None else mapped[rows])
            if spec["dtype"] == "BF16":
                value = value.view(mx.bfloat16)
            del mapped
        if self.cast is not None and mx.issubdtype(value.dtype, mx.floating):
            value = value.astype(self.cast)
        mx.eval(value)
        self.loads += 1
        self.bytes_read += value.nbytes
        if rows is None and cache and value.nbytes <= self.cache_limit:
            while self.cache and self.cached_bytes + value.nbytes > self.cache_limit:
                _, old = self.cache.popitem(last=False)
                self.cached_bytes -= old.nbytes
            self.cache[key] = value
            self.cached_bytes += value.nbytes
        return value

    def group(self, prefix, *, exclude=()):
        return {
            k[len(prefix) :]: self.get(k)
            for k in self.index
            if k.startswith(prefix) and not any(s in k for s in exclude)
        }

    def clear(self):
        self.cache.clear()
        self.cached_bytes = 0
        gc.collect()
        mx.clear_cache()


class Linear:
    def __init__(
        self, weight, bias=None, scales=None, qbiases=None, bits=8, group_size=64
    ):
        self.weight = weight
        self.bias = bias
        self.scales = scales
        self.qbiases = qbiases
        self.bits = bits
        self.group_size = group_size

    @classmethod
    def load(cls, store, prefix, *, bits=None, group_size=64, grouped_heads=None):
        # Converted MLX quantized tensors use .scales/.biases, independently of bias.
        if prefix + ".scales" in store:
            meta = json.loads((store.path / "quantization.json").read_text())
            return cls(
                store.get(prefix + ".weight"),
                store.get(prefix + ".bias") if prefix + ".bias" in store else None,
                store.get(prefix + ".scales"),
                store.get(prefix + ".biases"),
                meta["bits"],
                meta["group_size"],
            )
        w = store.get(prefix + ".weight")
        b = store.get(prefix + ".bias") if prefix + ".bias" in store else None
        # Original H3 QKV is [head, qkv, dim, input], unlike MLX [qkv, head, dim, input].
        if grouped_heads is not None:
            heads, dim = grouped_heads
            w = w.reshape(heads, 3, dim, -1).transpose(1, 0, 2, 3).reshape(w.shape)
            if b is not None:
                b = b.reshape(heads, 3, dim).transpose(1, 0, 2).reshape(b.shape)
        if bits is not None and w.dtype != mx.float32:
            q, s, z = mx.quantize(w, bits=bits, group_size=group_size)
            mx.eval(q, s, z)
            return cls(q, b, s, z, bits, group_size)
        mx.eval(w)
        return cls(w, b)

    @property
    def dtype(self):
        return self.scales.dtype if self.scales is not None else self.weight.dtype

    @property
    def nbytes(self):
        return sum(
            x.nbytes
            for x in (self.weight, self.bias, self.scales, self.qbiases)
            if x is not None
        )

    def __call__(self, x):
        x = x.astype(self.dtype)
        if self.scales is None:
            # One BF16 rounding after the bias epilogue, as in torch F.linear.
            return (
                x @ self.weight.T
                if self.bias is None
                else mx.addmm(self.bias, x, self.weight.T)
            )
        else:
            out = mx.quantized_matmul(
                x,
                self.weight,
                self.scales,
                self.qbiases,
                transpose=True,
                group_size=self.group_size,
                bits=self.bits,
            )
        return out if self.bias is None else out + self.bias

	"""Tensor-at-a-time safetensors loading; never load a whole checkpoint to RAM."""

	from __future__ import annotations

	import gc
	import hashlib
	import json
	import struct
	from collections import OrderedDict
	from pathlib import Path

	import mlx.core as mx
	import numpy as np


	class WeightStore:
	def __init__(self, path, *, cache_gb=0.0, cast=None, reader="mlx"):
	self.path = Path(path)
	files = (
	[self.path]
	if self.path.is_file()
	else sorted(self.path.glob("*.safetensors"))
	)
	if not files:
	raise FileNotFoundError(f"No safetensors weights at {self.path}")
	self.index = {}
	self.headers = {}
	self.data_start = {}
	self.cache = OrderedDict()
	self.cache_limit = int(cache_gb * 1024**3)
	self.cached_bytes = 0
	self.cast = cast
	if reader not in ("mlx", "memmap"):
	raise ValueError("reader must be mlx or memmap")
	self.reader = reader
	self.loads = 0
	self.bytes_read = 0
	self.files = files
	for file in files:
	with file.open("rb") as handle:
	raw = handle.read(8)
	if len(raw) != 8:
	raise ValueError(f"Truncated safetensors header: {file}")
	length = struct.unpack("<Q", raw)[0]
	if length > 100_000_000:
	raise ValueError(f"Invalid safetensors header size: {file}")
	header = json.loads(handle.read(length))
	self.data_start[file] = 8 + length
	for name, spec in header.items():
	if name == "__metadata__":
	continue
	if name in self.index:
	raise ValueError(f"Duplicate tensor {name} in {self.path}")
	self.index[name] = file
	self.headers[name] = spec

	def __contains__(self, key):
	return key in self.index

	def fingerprint(self):
	# Source identity for cache invalidation, not a content-integrity checksum.
	info = [
	(str(p.resolve()), p.stat().st_size, p.stat().st_mtime_ns)
	for p in self.files
	]
	return hashlib.sha256(json.dumps(info).encode()).hexdigest()

	def get(self, key, *, rows=None, cache=True):
	if rows is None and key in self.cache:
	self.cache.move_to_end(key)
	return self.cache[key]
	if key not in self.index:
	raise KeyError(f"Missing checkpoint tensor {key} in {self.path}")
	# safetensors' Python MLX adapter routes BF16 through NumPy, which has
	# no BF16 dtype. Map its uint16 payload directly and bitcast in MLX.
	spec, file = self.headers[key], self.index[key]
	dtypes = {
	"BF16": "<u2",
	"F16": "<f2",
	"F32": "<f4",
	"U32": "<u4",
	"I32": "<i4",
	"I64": "<i8",
	"U8": "u1",
	"BOOL": "?",
	}
	if spec["dtype"] not in dtypes:
	raise ValueError(f"Unsupported tensor dtype {spec['dtype']} for {key}")
	if self.reader == "mlx" and rows is None:
	# MLX creates lazy per-tensor Load nodes. Discard the dictionary
	# immediately: retaining it would also retain every evaluated weight
	# in the shard, bypassing this store's explicit cache budget.
	value = mx.load(str(file))[key]
	else:
	# A full-vocabulary MLX Load is evaluated before a row gather.
	# Read requested rows directly from the safetensors mapping instead;
	# this preserves order/duplicates and the original BF16 bit payload.
	mapped = np.memmap(
	file,
	dtype=dtypes[spec["dtype"]],
	mode="r",
	offset=self.data_start[file] + spec["data_offsets"][0],
	shape=tuple(spec["shape"]),
	)
	value = mx.array(mapped if rows is None else mapped[rows])
	if spec["dtype"] == "BF16":
	value = value.view(mx.bfloat16)
	del mapped
	if self.cast is not None and mx.issubdtype(value.dtype, mx.floating):
	value = value.astype(self.cast)
	mx.eval(value)
	self.loads += 1
	self.bytes_read += value.nbytes
	if rows is None and cache and value.nbytes <= self.cache_limit:
	while self.cache and self.cached_bytes + value.nbytes > self.cache_limit:
	_, old = self.cache.popitem(last=False)
	self.cached_bytes -= old.nbytes
	self.cache[key] = value
	self.cached_bytes += value.nbytes
	return value

	def group(self, prefix, *, exclude=()):
	return {
	k[len(prefix) :]: self.get(k)
	for k in self.index
	if k.startswith(prefix) and not any(s in k for s in exclude)
	}

	def clear(self):
	self.cache.clear()
	self.cached_bytes = 0
	gc.collect()
	mx.clear_cache()


	class Linear:
	def __init__(
	self, weight, bias=None, scales=None, qbiases=None, bits=8, group_size=64
	):
	self.weight = weight
	self.bias = bias
	self.scales = scales
	self.qbiases = qbiases
	self.bits = bits
	self.group_size = group_size

	@classmethod
	def load(cls, store, prefix, *, bits=None, group_size=64, grouped_heads=None):
	# Converted MLX quantized tensors use .scales/.biases, independently of bias.
	if prefix + ".scales" in store:
	meta = json.loads((store.path / "quantization.json").read_text())
	return cls(
	store.get(prefix + ".weight"),
	store.get(prefix + ".bias") if prefix + ".bias" in store else None,
	store.get(prefix + ".scales"),
	store.get(prefix + ".biases"),
	meta["bits"],
	meta["group_size"],
	)
	w = store.get(prefix + ".weight")
	b = store.get(prefix + ".bias") if prefix + ".bias" in store else None
	# Original H3 QKV is [head, qkv, dim, input], unlike MLX [qkv, head, dim, input].
	if grouped_heads is not None:
	heads, dim = grouped_heads
	w = w.reshape(heads, 3, dim, -1).transpose(1, 0, 2, 3).reshape(w.shape)
	if b is not None:
	b = b.reshape(heads, 3, dim).transpose(1, 0, 2).reshape(b.shape)
	if bits is not None and w.dtype != mx.float32:
	q, s, z = mx.quantize(w, bits=bits, group_size=group_size)
	mx.eval(q, s, z)
	return cls(q, b, s, z, bits, group_size)
	mx.eval(w)
	return cls(w, b)

	@property
	def dtype(self):
	return self.scales.dtype if self.scales is not None else self.weight.dtype

	@property
	def nbytes(self):
	return sum(
	x.nbytes
	for x in (self.weight, self.bias, self.scales, self.qbiases)
	if x is not None
	)

	def __call__(self, x):
	x = x.astype(self.dtype)
	if self.scales is None:
	# One BF16 rounding after the bias epilogue, as in torch F.linear.
	return (
	x @ self.weight.T
	if self.bias is None
	else mx.addmm(self.bias, x, self.weight.T)
	)
	else:
	out = mx.quantized_matmul(
	x,
	self.weight,
	self.scales,
	self.qbiases,
	transpose=True,
	group_size=self.group_size,
	bits=self.bits,
	)
	return out if self.bias is None else out + self.bias