mlx-dfloat 0.1.0__py3-none-any.whl

mlx_dfloat/mflux/flux1/transformer.py

@@ -0,0 +1,240 @@
+"""mflux's FLUX.1 Transformer with the block seam composed in front of its two per-block hooks."""
+
+import json
+from collections.abc import Mapping
+from functools import cache
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+from mlx.utils import tree_flatten
+
+from mlx_dfloat._safetensors import TensorInfo, read_header
+from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError
+from mlx_dfloat.format import DF11Checkpoint
+from mlx_dfloat.integrate import seam
+from mlx_dfloat.integrate.coverage import check_extras_cover, extras_plan, read_extra
+from mlx_dfloat.integrate.names import NameMap, Shapes
+from mlx_dfloat.integrate.placeholders import PLACEHOLDER, install_placeholders
+from mlx_dfloat.integrate.providers import WeightProvider
+from mlx_dfloat.mflux import require_mflux
+from mlx_dfloat.mflux.flux1.names import (
+    DOUBLE_PREFIX,
+    DROPPED_EXTRAS,
+    SINGLE_PREFIX,
+    check_flux_groups,
+    flux_name_map,
+)
+
+MAX_BUILD_ACTIVE_BYTES = 2 * 1024**3
+
+
+class SeamMixin:
+    """Overrides mflux's ``_apply_joint_transformer_block`` / ``_apply_single_transformer_block``.
+
+    Compose it in front of mflux's ``Transformer`` (``seam_transformer_class``) or a fake with the same hooks.
+    Call ``attach`` before the first step. A step is ``out = transformer(...)``; with ``verify_in_call`` the
+    decode status words are checked before the call returns, otherwise call ``verify_step()`` after the
+    step's final eval.
+    """
+
+    _seam: seam.SeamState
+
+    def attach(
+        self,
+        provider: WeightProvider,
+        shapes: Shapes,
+        *,
+        eval_policy: str = "per-block",
+        tracer: seam.TraceHook | None = None,
+        verify_in_call: bool = False,
+    ) -> None:
+        """Bind the provider, the block shapes (from ``install_placeholders``) and the eval policy."""
+        self._seam = seam.attach_state(
+            provider, shapes, eval_policy=eval_policy, tracer=tracer, verify_in_call=verify_in_call
+        )
+
+    def _state(self) -> seam.SeamState:
+        try:
+            return self._seam
+        except AttributeError as exc:
+            raise DFloatIntegrationError(
+                "call attach(provider, shapes) before running a step"
+            ) from exc
+
+    def detach(self) -> None:
+        """Forget the provider and the shapes; the next step needs a new ``attach``."""
+        if hasattr(self, "_seam"):
+            self._seam.provider.reset()
+            del self._seam
+
+    def verify_step(self) -> None:
+        """Run the provider's deferred checks; call it after the step's final ``mx.eval``."""
+        seam.verify_step(self._state())
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Run mflux's step; a failed step leaves no stale seam state."""
+        state = self._state()
+        seam.begin_step(state)
+        try:
+            out = super().__call__(*args, **kwargs)  # type: ignore[misc]
+            return seam.end_step(state, out)
+        except BaseException:
+            seam.abort_step(state)
+            raise
+
+    def _apply_joint_transformer_block(self, idx: int, block: Any, **kwargs: Any) -> Any:
+        return seam.run_block(
+            self._state(),
+            f"{DOUBLE_PREFIX}.{idx}",
+            block,
+            lambda: super(SeamMixin, self)._apply_joint_transformer_block(  # type: ignore[misc]
+                idx=idx, block=block, **kwargs
+            ),
+        )
+
+    def _apply_single_transformer_block(self, idx: int, block: Any, **kwargs: Any) -> Any:
+        return seam.run_block(
+            self._state(),
+            f"{SINGLE_PREFIX}.{idx}",
+            block,
+            lambda: super(SeamMixin, self)._apply_single_transformer_block(  # type: ignore[misc]
+                idx=idx, block=block, **kwargs
+            ),
+        )
+
+
+@cache
+def seam_transformer_class() -> type:
+    """``SeamMixin`` composed in front of mflux's ``Transformer`` (imports mflux).
+
+    Raises:
+        DFloatDependencyError: The optional ``mflux`` extra is not installed.
+    """
+    require_mflux()
+    from mflux.models.flux.model.flux_transformer.transformer import Transformer
+
+    return type("SeamTransformer", (SeamMixin, Transformer), {})
+
+
+def base_transformer_index(
+    root: Path, *, index_file: str = "diffusion_pytorch_model.safetensors.index.json"
+) -> dict[str, tuple[Path, TensorInfo]]:
+    """Every tensor of a sharded BF16 transformer directory: name -> (shard path, tensor info), from the index.
+
+    Raises:
+        DFloatFormatError: No usable index (missing, malformed JSON, or a ``weight_map`` that is not
+            a non-empty object of shard-name strings), or a shard name that is not a plain file name
+            in ``root``.
+    """
+    try:
+        weight_map = json.loads((root / index_file).read_text(encoding="utf-8"))["weight_map"]
+    except (OSError, ValueError, KeyError, TypeError) as exc:
+        raise DFloatFormatError(f"{root / index_file}: no usable weight index: {exc}") from exc
+    if (
+        not isinstance(weight_map, dict)
+        or not weight_map
+        or not all(isinstance(v, str) for v in weight_map.values())
+    ):
+        raise DFloatFormatError(
+            f"{root / index_file}: weight_map must be a non-empty object of shard-name strings"
+        )
+    index: dict[str, tuple[Path, TensorInfo]] = {}
+    for shard_name in sorted(set(weight_map.values())):
+        if not isinstance(shard_name, str) or Path(shard_name).name != shard_name:
+            raise DFloatFormatError(
+                f"{root / index_file}: shard {shard_name!r} escapes the directory"
+            )
+        shard = root / shard_name
+        index.update({name: (shard, info) for name, info in read_header(shard).items()})
+    missing = sorted(set(weight_map) - set(index))
+    if missing:
+        raise DFloatFormatError(
+            f"{root / index_file}: tensors named by the index but absent from the shards: {missing[:5]}"
+        )
+    return index
+
+
+def base_extras(
+    index: Mapping[str, tuple[Path, TensorInfo]], ckpt: DF11Checkpoint
+) -> dict[str, tuple[Path, TensorInfo]]:
+    """The index minus the checkpoint's block matrices: what ``build_transformer`` loads as extras from a BF16 base."""
+    matrices = {m for group in ckpt.groups.values() for m in group.matrix_names}
+    return {name: entry for name, entry in index.items() if name not in matrices}
+
+
+def build_transformer(
+    model_config: Any,
+    ckpt: DF11Checkpoint,
+    *,
+    name_map: NameMap | None = None,
+    n_double: int | None = None,
+    n_single: int | None = None,
+    extras: Mapping[str, tuple[Path, TensorInfo]] | None = None,
+) -> tuple[Any, Shapes]:
+    """Construct mflux's ``Transformer`` (seamed) with placeholders for the matrices and the extras loaded.
+
+    Block counts come from the checkpoint's groups, or ``n_double``/``n_single`` for a reduced-depth build
+    (each between zero and the checkpoint's own count). Every block matrix is a placeholder and every other
+    parameter comes from ``extras`` (a BF16 base's index, via ``base_extras``) or, when ``extras`` is None,
+    the checkpoint's own extras; the coverage of both is asserted, and so is the MLX active memory the
+    build added (under ``MAX_BUILD_ACTIVE_BYTES``).
+
+    Raises:
+        DFloatFormatError: The groups are not FLUX.1's, an extra is not BF16 or has the wrong shape.
+        DFloatDependencyError: The optional ``mflux`` extra is not installed.
+        DFloatIntegrationError: A depth override is negative or exceeds the checkpoint's own block
+            count, an uncovered parameter, an extra without a target, or too much active memory added
+            by the build.
+    """
+    names = flux_name_map() if name_map is None else name_map
+    ckpt_double, ckpt_single = check_flux_groups(ckpt)
+    n_double = ckpt_double if n_double is None else n_double
+    n_single = ckpt_single if n_single is None else n_single
+    if n_double < 0 or n_single < 0:
+        raise DFloatIntegrationError(
+            f"asked for {n_double} double / {n_single} single blocks; a depth cannot be negative"
+        )
+    if n_double > ckpt_double or n_single > ckpt_single:
+        raise DFloatIntegrationError(
+            f"asked for {n_double} double / {n_single} single blocks; the checkpoint has "
+            f"{ckpt_double} / {ckpt_single}"
+        )
+    before = int(mx.get_active_memory())
+    transformer = seam_transformer_class()(
+        model_config, num_transformer_blocks=n_double, num_single_transformer_blocks=n_single
+    )
+    shapes = install_placeholders(
+        [
+            (DOUBLE_PREFIX, transformer.transformer_blocks),
+            (SINGLE_PREFIX, transformer.single_transformer_blocks),
+        ],
+        names,
+    )
+    matrix_paths = {f"{block}.{attr}.weight" for block, per in shapes.items() for attr in per}
+    params = dict(tree_flatten(transformer.parameters()))
+    plan = extras_plan(
+        ckpt,
+        names,
+        counts={DOUBLE_PREFIX: n_double, SINGLE_PREFIX: n_single},
+        dropped=DROPPED_EXTRAS,
+        extras=extras,
+    )
+    check_extras_cover(params, (name for name, _path, _info in plan), matrix_paths)
+    weights: list[tuple[str, mx.array]] = []
+    for name, path, info in plan:
+        array = read_extra(path, info)
+        if tuple(array.shape) != tuple(params[name].shape):
+            raise DFloatFormatError(
+                f"{name}: extra has shape {tuple(array.shape)}, parameter {tuple(params[name].shape)}"
+            )
+        weights.append((name, array))
+    transformer.load_weights(weights, strict=False)
+    mx.eval([array for _name, array in weights], PLACEHOLDER)
+    added = int(mx.get_active_memory()) - before
+    if added >= MAX_BUILD_ACTIVE_BYTES:
+        raise DFloatIntegrationError(
+            f"extras added {added / 1024**3:.2f} GiB of active memory "
+            f"(limit {MAX_BUILD_ACTIVE_BYTES / 1024**3:g} GiB)"
+        )
+    return transformer, shapes

mlx_dfloat/reference.py

@@ -0,0 +1,241 @@
+"""Pure-NumPy, bit-exact DFloat11 decoder: the oracle every GPU kernel is tested against.
+
+It emulates upstream's CUDA algorithm in lockstep: each of 512 threads per 4096-byte block owns one
+64-bit chunk of the exponent stream and decodes the codes that start in it (from its 5-bit gap); an
+exclusive scan turns per-thread counts into output indices. Structural checks run as it decodes.
+"""
+
+import os
+
+import numpy as np
+import numpy.typing as npt
+
+from mlx_dfloat.errors import DFloatFormatError, DFloatResourceError
+from mlx_dfloat.format import (
+    BYTES_PER_THREAD,
+    LUT_POINTER_MIN,
+    THREADS_PER_BLOCK,
+    DF11Group,
+    GroupArrays,
+    validate_group_arrays,
+)
+
+_BITS_PER_THREAD = 8 * BYTES_PER_THREAD
+
+IntArray = npt.NDArray[np.int64]
+BoolArray = npt.NDArray[np.bool_]
+U16Array = npt.NDArray[np.uint16]
+
+
+def thread_gaps(gaps: npt.NDArray[np.uint8], n_threads: int) -> IntArray:
+    """Unpack the MSB-first 5-bit start offsets, one per thread."""
+    need = -(-5 * n_threads // 8)
+    bits = np.unpackbits(np.asarray(gaps[:need], dtype=np.uint8))[: 5 * n_threads].reshape(
+        n_threads, 5
+    )
+    packed = (
+        (bits[:, 0] << 4) | (bits[:, 1] << 3) | (bits[:, 2] << 2) | (bits[:, 3] << 1) | bits[:, 4]
+    )
+    return packed.astype(np.int64)
+
+
+def estimate_decode_bytes(arrays: GroupArrays) -> int:
+    """Upper estimate of the decoder's working memory for one group."""
+    return 12 * arrays.n_threads * BYTES_PER_THREAD + 3 * arrays.n_elements
+
+
+def default_memory_budget() -> int:
+    """Half of physical RAM."""
+    return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") // 2
+
+
+def max_code_length(luts: npt.NDArray[np.uint8]) -> int:
+    """Longest code in a group's codebook (the maximum of the lengths row)."""
+    return int(np.asarray(luts)[-1].max())
+
+
+def _peek32(stream: npt.NDArray[np.uint8], pos: IntArray) -> npt.NDArray[np.uint32]:
+    byte = pos >> 3
+    shift = (pos & 7).astype(np.uint64)
+    word = np.zeros(pos.size, dtype=np.uint64)
+    for k in range(5):
+        word = (word << np.uint64(8)) | stream[byte + k].astype(np.uint64)
+    return ((word >> (np.uint64(8) - shift)) & np.uint64(0xFFFFFFFF)).astype(np.uint32)
+
+
+def _lookup(
+    luts: npt.NDArray[np.uint8], words: npt.NDArray[np.uint32]
+) -> tuple[IntArray, IntArray, BoolArray]:
+    n_luts = luts.shape[0]
+    sym = luts[0, words >> 24].astype(np.int64)
+    ok = np.ones(words.size, dtype=np.bool_)
+    for level in (1, 2, 3):
+        ptr = sym >= LUT_POINTER_MIN
+        if not ptr.any():
+            break
+        row = 256 - sym[ptr]
+        in_range = (row >= 1) & (row <= n_luts - 2)
+        byte = (words[ptr] >> np.uint32(24 - 8 * level)) & np.uint32(0xFF)
+        sym[ptr] = np.where(in_range, luts[np.where(in_range, row, 0), byte], LUT_POINTER_MIN)
+        ok[ptr] &= in_range
+    ok &= sym < LUT_POINTER_MIN
+    length = np.where(ok, luts[n_luts - 1, np.minimum(sym, 255)], 0).astype(np.int64)
+    ok &= length > 0
+    return sym, length, ok
+
+
+def _count(
+    stream: npt.NDArray[np.uint8], luts: npt.NDArray[np.uint8], starts: IntArray, ends: IntArray
+) -> tuple[IntArray, BoolArray, IntArray]:
+    counts = np.zeros(starts.size, dtype=np.int64)
+    invalid = np.zeros(starts.size, dtype=np.bool_)
+    pos = starts.copy()
+    active = np.flatnonzero(pos < ends)
+    while active.size:
+        _, length, ok = _lookup(luts, _peek32(stream, pos[active]))
+        invalid[active[~ok]] = True
+        good = active[ok]
+        counts[good] += 1
+        pos[good] += length[ok]
+        active = good[pos[good] < ends[good]]
+    return counts, invalid, pos
+
+
+def _write(
+    stream: npt.NDArray[np.uint8],
+    luts: npt.NDArray[np.uint8],
+    sign_mantissa: npt.NDArray[np.uint8],
+    starts: IntArray,
+    first: IntArray,
+    stop: IntArray,
+) -> tuple[U16Array, int]:
+    n = sign_mantissa.size
+    out = np.zeros(n, dtype=np.uint16)
+    pos = starts.copy()
+    idx = first.copy()
+    active = np.flatnonzero(idx < stop)
+    while active.size:
+        sym, length, _ = _lookup(luts, _peek32(stream, pos[active]))  # validated in pass 1
+        k = idx[active]
+        sm = sign_mantissa[k].astype(np.uint16)
+        out[k] = ((sm & 0x80) << 8) | (sym.astype(np.uint16) << 7) | (sm & 0x7F)
+        idx[active] += 1
+        pos[active] += length
+        active = active[idx[active] < stop[active]]
+    last_thread = int(np.searchsorted(first, n - 1, side="right") - 1)
+    return out, int(pos[last_thread])
+
+
+def decode_group(
+    arrays: GroupArrays,
+    *,
+    name: str = "<group>",
+    max_memory_bytes: int | None = None,
+    check_stream_end: bool = True,
+) -> U16Array:
+    """Decode one compressed group to its BF16 bit patterns (flat, concatenation order).
+
+    ``check_stream_end=False`` skips only the EOF byte-count check; it exists for fixtures that are
+    deliberately truncated mid-stream (the committed upstream slice), never for real checkpoints.
+
+    Raises:
+        DFloatFormatError: The arrays are structurally invalid or internally inconsistent.
+        DFloatResourceError: The estimated working memory exceeds ``max_memory_bytes`` (default:
+            half of physical RAM).
+    """
+    validate_group_arrays(arrays, name=name)
+    budget = default_memory_budget() if max_memory_bytes is None else max_memory_bytes
+    need = estimate_decode_bytes(arrays)
+    if need > budget:
+        raise DFloatResourceError(
+            f"{name}: decoding needs about {need} bytes, over the {budget}-byte budget"
+        )
+    n, n_bytes, n_threads = arrays.n_elements, arrays.n_bytes, arrays.n_threads
+    n_bits = 8 * n_bytes
+    last_real = (n_bits - 1) // _BITS_PER_THREAD
+    g = np.arange(last_real + 1, dtype=np.int64)
+    starts = g * _BITS_PER_THREAD + thread_gaps(arrays.gaps, n_threads)[: last_real + 1]
+    ends = (g + 1) * _BITS_PER_THREAD
+    stream = np.zeros(n_threads * BYTES_PER_THREAD + 8, dtype=np.uint8)
+    stream[:n_bytes] = arrays.encoded_exponent
+    luts = np.ascontiguousarray(arrays.luts)
+
+    counts, invalid, ended = _count(stream, luts, starts, ends)
+    bad = np.flatnonzero(invalid[:-1])
+    if bad.size:
+        t = int(bad[0])
+        raise DFloatFormatError(
+            f"{name}: invalid code in thread {t} (block {t // THREADS_PER_BLOCK})"
+        )
+    inside = ended[:-1] < n_bits
+    broken = np.flatnonzero(inside & (ended[:-1] != starts[1:]))
+    if broken.size:
+        t = int(broken[0])
+        raise DFloatFormatError(
+            f"{name}: thread {t} ends at bit {int(ended[t])} but thread {t + 1} starts at bit "
+            f"{int(starts[t + 1])} (corrupt gaps?)"
+        )
+    first = np.zeros_like(counts)
+    np.cumsum(counts[:-1], out=first[1:])
+    positions = arrays.output_positions.astype(np.int64)
+    block_starts = np.arange(positions.size - 1, dtype=np.int64) * THREADS_PER_BLOCK
+    mismatch = np.flatnonzero(first[block_starts] != positions[:-1])
+    if mismatch.size:
+        b = int(mismatch[0])
+        raise DFloatFormatError(
+            f"{name}: block {b} starts at element {int(first[block_starts[b]])}, output_positions says "
+            f"{int(positions[b])}"
+        )
+    if positions.size - 1 == arrays.n_blocks - 1:
+        # The last block holds at least one byte, so its first thread is always a real one
+        # (tail <= last_real) and first[tail] is in range.
+        tail = (arrays.n_blocks - 1) * THREADS_PER_BLOCK
+        if first[tail] < n:
+            raise DFloatFormatError(
+                f"{name}: block {arrays.n_blocks - 1} starts codes but output_positions has no entry for it"
+            )
+    total = int(first[-1] + counts[-1])
+    if total < n:
+        hint = " (invalid code in the last thread)" if invalid[-1] else ""
+        raise DFloatFormatError(f"{name}: stream holds {total} codes, expected {n}{hint}")
+    stop = np.minimum(first + counts, n)
+    out, end_bit = _write(stream, luts, np.asarray(arrays.sign_mantissa), starts, first, stop)
+    if check_stream_end and -(-end_bit // 8) != n_bytes:
+        raise DFloatFormatError(
+            f"{name}: data ends at bit {end_bit}, so the stream should be {-(-end_bit // 8)} bytes, not {n_bytes}"
+        )
+    return out
+
+
+def split_matrices(flat: U16Array, split_positions: npt.NDArray[np.int64]) -> list[U16Array]:
+    """Split a decoded group into its matrices (flat), at ``split_positions``."""
+    return list(np.split(flat, split_positions.astype(np.int64)))
+
+
+def decode_matrices(
+    group: DF11Group, *, max_memory_bytes: int | None = None
+) -> dict[str, U16Array]:
+    """Decode a group and return its matrices keyed by weight name (flat BF16 bit patterns)."""
+    arrays = group.load()
+    parts = split_matrices(
+        decode_group(arrays, name=group.name, max_memory_bytes=max_memory_bytes),
+        arrays.split_positions,
+    )
+    return dict(zip(group.matrix_names, parts, strict=True))
+
+
+def max_elements_per_block(output_positions: npt.NDArray[np.uint32]) -> int:
+    """Largest number of elements whose codes start in one 4096-byte block (threadgroup sizing)."""
+    return int(np.diff(output_positions.astype(np.int64)).max())
+
+
+__all__ = [
+    "decode_group",
+    "decode_matrices",
+    "default_memory_budget",
+    "estimate_decode_bytes",
+    "max_code_length",
+    "max_elements_per_block",
+    "split_matrices",
+    "thread_gaps",
+]