Check MLX weight shard integrity on a cloud Mac

If a model fails to load after you transfer its weights to a cloud Mac, insufficient unified memory is not necessarily the cause. An interrupted download, shards mixed in from another revision, or an index pointing to old files can stop an experiment before it begins. Checking the files without loading tensors before running MLX helps avoid repeated downloads and blind changes to inference parameters.

Start with a fixed weights directory

When validating weights on OsxNode, put the weights, configuration, and index for each experiment in a separate weights directory. Do not overwrite a model that is currently running. Have Python 3 available, and save the verification script outside that directory. This article checks only safetensors weights and does not execute any code bundled with the model.

Record the model revision identifier, quantization configuration, download source, and file list. These details must reflect the version you actually obtained; a model name whose contents may change is not enough. Wait until the downloader exits before starting the checks so that shards still being written are not mistaken for damaged files.

A valid file structure, an unchanged transfer, and successful model inference are three separate validation results. Passing any one does not establish the other two.

Check offsets and the index using file headers

The first eight bytes of a safetensors file are an unsigned little-endian integer giving the length of the JSON header in bytes. Each tensor's data_offsets are relative to the data region after the header, not the beginning of the file. Confusing these two reference points will incorrectly flag valid shards as out of bounds.

The script below checks that the data region is covered without gaps, that tensor names are not duplicated across shards, and that the index mapping matches the actual files. Save it as verify_weights.py. The 16MiB header limit is a local validation policy, not a format limit. If you encounter a valid larger header, verify its source before adjusting the limit.

import json
import struct
import sys
from pathlib import Path

def require(ok, message):
    if not ok:
        raise ValueError(message)

def unique_object(pairs):
    result = {}
    for key, value in pairs:
        require(key not in result, f"Duplicate JSON key: {key}")
        result[key] = value
    return result

def read_json(raw):
    return json.loads(raw, object_pairs_hook=unique_object)

def verify(root):
    require(root.is_dir(), "Missing model directory")
    files = sorted(root.glob("*.safetensors"))
    require(bool(files), "No safetensors files")
    actual = {}

    for path in files:
        require(not path.is_symlink(), "Symlink shard rejected")
        size = path.stat().st_size
        with path.open("rb") as stream:
            raw = stream.read(8)
            require(len(raw) == 8, f"Short prefix: {path.name}")
            length = struct.unpack("<Q", raw)[0]
            require(
                0 < length <= min(16 * 1024 * 1024, size - 8),
                f"Invalid header length: {path.name}"
            )
            header = read_json(stream.read(length))

        require(isinstance(header, dict), "Header must be an object")
        data_size = size - 8 - length
        ranges = []
        for name, info in header.items():
            if name == "__metadata__":
                continue
            require(isinstance(info, dict), f"Invalid tensor: {name}")
            offsets = info.get("data_offsets")
            require(
                isinstance(offsets, list) and len(offsets) == 2
                and all(type(v) is int for v in offsets),
                f"Invalid offsets: {name}"
            )
            start, end = offsets
            require(0 <= start <= end <= data_size,
                    f"Out of bounds: {name}")
            require(name not in actual, f"Duplicate tensor: {name}")
            actual[name] = path.name
            ranges.append((start, end))

        cursor = 0
        for start, end in sorted(ranges):
            require(start == cursor, f"Gap or overlap: {path.name}")
            cursor = end
        require(cursor == data_size, f"Trailing data: {path.name}")

    index = root / "model.safetensors.index.json"
    if index.exists():
        document = read_json(index.read_text(encoding="utf-8"))
        mapping = document.get("weight_map")
        require(isinstance(mapping, dict), "Missing weight_map")
        require(mapping == actual, "Shard index mismatch")
    else:
        require(len(files) == 1, "Multiple shards require an index")

    print(f"PASS: {len(files)} files, {len(actual)} tensors")

try:
    require(len(sys.argv) == 2, "Usage: verify_weights.py DIRECTORY")
    verify(Path(sys.argv[1]))
except Exception as error:
    print(f"FAIL: {error}", file=sys.stderr)
    sys.exit(1)

Run python3 verify_weights.py weights and continue only if the exit code is zero. The script assumes the standard index filename and expects shards at the top level of the directory. Other layouts require explicit adaptations; do not simply delete the index to bypass the check.

Verify the transfer against source-side hashes

Structural checks cannot detect corrupted numerical values inside tensors. Generate a checksum manifest on the trusted source machine after all files have finished writing, then transfer that same set of files and the manifest to the destination machine. Run all of the following commands inside the weights directory:

shasum -a 256 ./*.safetensors ./config.json > SHA256SUMS

If a shard index exists, add it to the manifest as well:

shasum -a 256 ./model.safetensors.index.json >> SHA256SUMS

On the destination machine, enter the corresponding directory and run:

shasum -a 256 -c SHA256SUMS

Every listed file should return OK. The manifest itself must come through a trusted channel. If an attacker can replace both the files and the manifest, hashes do not authenticate their source. Weights already damaged at the source will also pass the transfer check unchanged, so keep the structural checks as well.

Handle each failure type instead of downloading everything again

For Short prefix or an invalid header length, first confirm that the download process has exited, then check the affected shard's byte count. Do not manually pad the file or rename a downloaded error page as a weights file.

For Shard index mismatch, compare each filename in the index with the directory listing. Common causes include shards omitted during copying and files left over from an older revision. Isolate extra files first, then rerun the check. Do not edit the index to fit a set of weights of unknown origin.

If the structural check passes but the hash check fails, retrieve only the failed files again, then rerun the full set of checks. If the transfer tool supports resuming, first confirm that the remote source file has not been replaced with another version. Otherwise, the first half of an old file may be joined to the second half of a new one.

Only when the hashes pass but loading still fails should you investigate the model architecture, quantization configuration, runtime library versions, and memory pressure. This order separates file problems from runtime environment problems without relying on error wording to guess the cause.

Make verification a prerequisite for experiments

The script does not validate the mathematical relationship between dtype, shape, and a tensor's byte count, nor does it determine whether the quantization parameters are compatible. You still need a final model-loading test in an MLX environment with pinned versions, followed by a short-output test using a fixed input. Header validation alone cannot replace this step.

Keep at least the revision identifier, checksum manifest, script exit code, runtime library versions, and initial loading result in your experiment records. Do not save only a screenshot of a successful run; it will be difficult to explain differences after the directory contents change.

Run the structural and hash checks at the start of the experiment script, and stop all subsequent work if either fails. Use a new directory and generate a new baseline whenever you change model revisions. Take the manifest and verification records with you before the rental period ends. That way, the next machine receives verifiable experiment inputs, not just a folder labeled “download complete.”

Frequently asked questions