mlx/src/solomon_mlx/prepare.py

"""Verify immutable input snapshots and create a separately checksummed BF16 model."""

import hashlib
import json
import shutil
from importlib.metadata import version
from pathlib import Path

from .artifacts import BASE_REVISION, SOLOMON_REVISION, sha256, verify_release


def verify_base(root, manifest):
    root = Path(root)
    if manifest["revision"] != BASE_REVISION:
        raise ValueError("Wrong pinned base revision")
    result = {}
    for row in manifest["files"]:
        path = root / row["name"]
        if not path.resolve().is_relative_to(root.resolve()):
            raise ValueError("Unsafe base manifest path")
        if path.stat().st_size != row["size"]:
            raise ValueError("Base size mismatch: " + row["name"])
        actual = sha256(path)
        if row["sha256"]:
            if actual != row["sha256"]:
                raise ValueError("Base checksum mismatch: " + row["name"])
        else:
            data = path.read_bytes()
            git_hash = hashlib.sha1(b"blob " + str(len(data)).encode() + b"\0" + data).hexdigest()
            if git_hash != row["blob_id"]:
                raise ValueError("Base Git blob mismatch: " + row["name"])
        result[row["name"]] = actual
    return result


def prepare(base, solomon, output, manifest):

    base, solomon, output = map(Path, (base, solomon, output))
    source_hashes = verify_release(solomon)
    base_hashes = verify_base(base, json.loads(Path(manifest).read_text()))
    if output.exists():
        raise FileExistsError("Use a new output directory; existing conversions are immutable")
    output.mkdir(parents=True)
    convert_bf16(base, output / "backbone")
    shutil.copy2(solomon / "adapter/adapter.safetensors", output / "adapter.safetensors")
    shutil.copy2(solomon / "adapter/heads.npz", output / "heads.npz")
    for name in ("LICENSE", "NOTICE", "MODIFICATIONS.md"):
        shutil.copy2(solomon / name, output / name)
    files = {
        str(p.relative_to(output)): sha256(p)
        for p in sorted(output.rglob("*"))
        if p.is_file() and ".cache" not in p.parts
    }
    binding = {
        "schema": "solomon-mlx-binding-v1",
        "profile": "quality",
        "dtype": "bfloat16",
        "adapter_dtype": "float32",
        "head_dtype": "float32",
        "recurrent_state_dtype": "float32",
        "sensitive_parameters": "FP32 normalization weights, A_log and dt_bias",
        "conversion": "upstream Qwen3.5 sanitization, norms promoted before unit offset",
        "adapter_scale": 2.0,
        "adapter_placement": "question",
        "quantization": None,
        "base_revision": BASE_REVISION,
        "solomon_revision": SOLOMON_REVISION,
        "source_hashes": source_hashes,
        "base_hashes": base_hashes,
        "dependencies": {p: version(p) for p in ("mlx", "mlx-vlm", "transformers", "numpy")},
        "files": files,
    }
    (output / "binding.json").write_text(json.dumps(binding, indent=2, sort_keys=True))
    return binding


def convert_bf16(base, output):
    """Convert one original shard at a time; never copy download caches.

    Keep normalization offsets in FP32 before adding one. Adding the unit
    offset in BF16 would irreversibly round trained normalization parameters.
    Large backbone matrices remain unquantized BF16.
    """
    import mlx.core as mx
    from mlx.utils import tree_flatten
    from mlx_vlm.models.qwen3_5.config import ModelConfig
    from mlx_vlm.models.qwen3_5.qwen3_5 import Model

    base, output = Path(base), Path(output)
    output.mkdir(parents=True, exist_ok=False)
    config = json.loads((base / "config.json").read_text())
    if config.get("model_type") != "qwen3_5" or config.get("quantization"):
        raise ValueError("Expected original unquantized Qwen3.5 architecture")
    model = Model(ModelConfig.from_dict(config))
    expected = {k: v.shape for k, v in tree_flatten(model.parameters())}
    index = {"metadata": {"total_size": 0}, "weight_map": {}}
    source_index = json.loads((base / "model.safetensors.index.json").read_text())
    for filename in sorted(set(source_index["weight_map"].values())):
        arrays = mx.load(str(base / filename))
        for key, value in arrays.items():
            sensitive = value.ndim == 1 and ("norm" in key or key.endswith(("A_log", "dt_bias")))
            arrays[key] = value.astype(mx.float32 if sensitive else mx.bfloat16)
        arrays = model.sanitize(arrays)
        arrays = model.vision_tower.sanitize(arrays)
        for key, value in arrays.items():
            if key not in expected or value.shape != expected[key]:
                raise ValueError("Converted tensor shape mismatch: " + key)
            if key in index["weight_map"]:
                raise ValueError("Duplicate converted tensor: " + key)
            index["weight_map"][key] = filename
            index["metadata"]["total_size"] += value.nbytes
        mx.save_safetensors(str(output / filename), arrays, metadata={"format": "mlx"})
        del arrays
        mx.clear_cache()
        print("Converted " + filename, flush=True)
    if set(index["weight_map"]) != set(expected):
        raise ValueError("Converted backbone is missing required parameters")
    (output / "model.safetensors.index.json").write_text(json.dumps(index, indent=2, sort_keys=True))
    for file in base.iterdir():
        if file.is_file() and file.name != "model.safetensors.index.json" and file.suffix != ".safetensors":
            shutil.copy2(file, output / file.name)