クラウド Mac で MLX 重みシャードの整合性を検証する

モデルの重みをクラウド Mac に転送したあと、読み込みに失敗したからといって、ユニファイドメモリ不足とは限りません。ダウンロードの中断、別バージョンのシャードの混入、古いファイルを参照するインデックスなどが原因で、実験の最初の段階でつまずくこともあります。まずテンソルを読み込まずにファイルの問題を切り分けてから MLX を実行すれば、何度もダウンロードし直したり、やみくもに推論パラメータを調整したりせずに済みます。

検証対象の重みディレクトリを固定する

OsxNode 上で検証する際は、1 回の実験で使う重み、設定、インデックスを専用の weights ディレクトリにまとめ、実行中のモデルを上書きしないようにします。Python 3 を用意し、検証スクリプトはそのディレクトリの外に保存してください。本記事で検証するのは safetensors の重みのみで、モデルに付属するコードは実行しません。

モデルのリビジョン識別子、量子化設定、ダウンロード元、ファイル一覧を記録します。これらは実際に取得したバージョンに基づく情報でなければなりません。参照先が変わり得るモデル名だけでは不十分です。また、書き込み中のシャードを破損ファイルと誤判定しないよう、ダウンローダーが終了してから検証を始めます。

ファイル構造が正しいこと、転送内容が一致すること、モデルで推論できることは、それぞれ別の検証結果です。どれか 1 つが合格しても、残り 2 つの合格を保証するものではありません。

ヘッダーからオフセットとインデックスを検証する

safetensors の先頭 8 バイトはリトルエンディアンの符号なし整数で、JSON ヘッダーのバイト長を表します。テンソルの data_offsets は、ファイル先頭ではなく、ヘッダー直後のデータ領域を基準とした位置です。この 2 つの基準を混同すると、正常なシャードを範囲外と誤判定してしまいます。

次のスクリプトは、データ領域が隙間なく連続して覆われているか、シャード間でテンソル名が重複していないか、インデックスの対応関係が実際のファイルと一致しているかを検証します。verify_weights.py として保存してください。ヘッダー長の上限を 16MiB としているのは、この検証で採用するローカルな方針であり、フォーマット自体の上限ではありません。正当な大きなヘッダーに遭遇した場合は、出所を確認してから上限を調整します。

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)

python3 verify_weights.py weights を実行し、終了コードがゼロの場合にのみ次へ進みます。このスクリプトは、標準のインデックスファイル名を使い、シャードをディレクトリの直下に置く構成を前提としています。別の構成では明示的な対応が必要です。検証を回避するためにインデックスを削除してはいけません。

転送元のハッシュで転送結果を検証する

構造の検証では、テンソル内部の数値の破損は検出できません。信頼できる転送元で、ファイルの書き込みが完了してからハッシュ一覧を生成し、同じ一式のファイルと一覧を転送先のマシンに送ります。以下のコマンドは、いずれも重みディレクトリ内で実行します。

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

シャードのインデックスがある場合は、それも一覧に追加します。

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

転送先のマシンで対応するディレクトリに移動し、次を実行します。

shasum -a 256 -c SHA256SUMS

一覧に含まれるすべてのファイルで OK が返る必要があります。ハッシュ一覧そのものも、信頼できる経路で取得しなければなりません。攻撃者がファイルと一覧の両方を差し替えられる場合、ハッシュでは出所を認証できません。また、転送元ですでに破損していた重みも、そのまま転送されれば照合に合格するため、構造の検証も省略せずに行います。

すべてを再ダウンロードせず、失敗の種類に応じて対処する

Short prefix やヘッダー長の異常が出た場合は、まずダウンロードプロセスが終了していることを確認し、該当するシャードのバイト数を照合します。ファイルの不足分を手作業で埋めたり、ダウンロード時のエラーページを重みファイルに改名したりしてはいけません。

Shard index mismatch が出た場合は、インデックス内のファイル名とディレクトリ内のファイル一覧を 1 つずつ比較します。よくある原因は、コピー時のシャードの取りこぼしや、旧バージョンのファイルの残存です。まず余分なファイルを隔離し、再度検証してください。出所の不明な重み一式に合わせてインデックスを書き換えてはいけません。

構造の検証に合格してもハッシュが一致しない場合は、不一致のファイルだけを再取得し、その後に検証一式をやり直します。転送ツールが中断箇所からの再開に対応していても、まずリモート側の元ファイルが別バージョンに差し替えられていないか確認してください。そうしないと、旧ファイルの前半と新ファイルの後半がつなぎ合わされる可能性があります。

ハッシュが一致しているのに読み込みに失敗する場合に初めて、モデルアーキテクチャ、量子化設定、ランタイムのバージョン、メモリ負荷の調査に進みます。この順序なら、ファイルの不具合と実行環境の不具合を切り分けられ、例外メッセージだけを頼りに原因を推測する必要がありません。

実験の開始前に検証を必須化する

このスクリプトは、dtype、shape、テンソルのバイト数の数学的な整合性を検証していません。量子化パラメータの組み合わせが適切かどうかも判定しません。最後には、バージョンを固定した MLX 環境で実際にモデルを読み込み、固定した入力から短い出力を生成するテストを行う必要があります。ヘッダーの検証だけで、この手順を代替することはできません。

実験記録には、少なくともリビジョン識別子、ハッシュ一覧、スクリプトの終了コード、ランタイムのバージョン、初回読み込みの結果を残します。成功時のスクリーンショットだけでは、次にディレクトリの内容が変わったとき、差異を説明するのが難しくなります。

実験スクリプトの先頭に構造検証とハッシュ照合を組み込み、どちらかが失敗した時点で後続の処理を停止させます。モデルのリビジョンを変更するときは新しいディレクトリを使い、検証の基準を作り直します。レンタル期間が終わる前に、一覧と検証記録も手元に保存してください。そうすれば、次のマシンに引き継ぐのは、単に「ダウンロード済み」と呼ばれるフォルダではなく、照合可能な実験入力になります。

よくある質問