Проверка целостности шардов весов MLX на облачном Mac

Если после переноса весов модели на облачный Mac модель не загружается, причина не обязательно в нехватке объединённой памяти. Прерванная загрузка, шарды от разных версий или индекс, ссылающийся на старые файлы, могут остановить эксперимент ещё до его начала. Сначала исключите проблемы с файлами с помощью проверки без загрузки тензоров, а уже затем запускайте MLX. Это позволит избежать повторного скачивания всех весов и настройки параметров инференса вслепую.

Подготовьте отдельный каталог весов

Для проверки на OsxNode поместите веса, конфигурацию и индекс одного эксперимента в отдельный каталог weights. Не перезаписывайте файлы модели, которая уже работает. Подготовьте Python 3, а скрипт проверки сохраните за пределами этого каталога. В этой статье проверяются только веса в формате safetensors; код, поставляемый с моделью, не выполняется.

Зафиксируйте идентификатор ревизии модели, параметры квантования, источник загрузки и список файлов. Эти сведения должны соответствовать фактически полученной версии: одного имени модели, за которым со временем может скрываться другое содержимое, недостаточно. Начинайте проверку только после завершения процесса скачивания, чтобы не принять ещё записываемый шард за повреждённый файл.

Корректная структура файлов, неизменность содержимого при передаче и способность модели выполнять инференс — три разных результата проверки. Успех в одном из них не подтверждает два остальных.

Проверьте смещения и индекс по заголовкам файлов

Первые восемь байт safetensors — беззнаковое целое число с порядком байтов little-endian, задающее длину JSON-заголовка в байтах. Значения data_offsets для тензоров отсчитываются от начала области данных после заголовка, а не от начала файла. Если перепутать эти точки отсчёта, корректные шарды будут ошибочно отмечены как содержащие выход за границы.

Следующий скрипт проверяет, что диапазоны тензоров непрерывно покрывают область данных, имена тензоров не повторяются в разных шардах, а отображение в индексе соответствует фактическим файлам. Сохраните его как verify_weights.py. Ограничение размера заголовка в 16 МиБ — локальное правило проверки, а не ограничение формата. Если встретится корректный заголовок большего размера, сначала проверьте источник и только затем изменяйте лимит.

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 поочерёдно сравните имена файлов в индексе со списком файлов в каталоге. Частые причины — пропущенный при копировании шард или оставшиеся файлы старой версии. Сначала перенесите лишние файлы в отдельное место и повторите проверку. Не подгоняйте индекс под набор весов неизвестного происхождения.

Если структура корректна, но хеши не совпадают, повторно получите только файлы, не прошедшие проверку, а затем запустите весь набор проверок заново. Если инструмент передачи поддерживает докачку, сначала убедитесь, что версия исходного файла на удалённой стороне не изменилась. Иначе можно получить файл, первая часть которого относится к старой версии, а вторая — к новой.

Только если хеши совпадают, а модель всё равно не загружается, переходите к проверке архитектуры модели, параметров квантования, версий библиотек и нагрузки на память. Такой порядок отделяет проблемы файлов от проблем среды выполнения и избавляет от необходимости угадывать причину по тексту исключения.

Сделайте проверки обязательным этапом запуска эксперимента

Скрипт не проверяет математическое соответствие между dtype, shape и числом байтов тензора, а также совместимость параметров квантования. Поэтому в завершение всё равно нужно загрузить модель в окружении MLX с зафиксированными версиями и выполнить тест с фиксированным входом и коротким выводом. Проверка одних заголовков этого не заменяет.

В журнале эксперимента сохраняйте как минимум идентификатор ревизии, список контрольных сумм, код завершения скрипта, версии библиотек и результат первой загрузки модели. Одного скриншота успешного запуска недостаточно: после следующего изменения содержимого каталога будет трудно объяснить различия.

Добавьте проверку структуры и сверку хешей в начало скрипта эксперимента и останавливайте дальнейшие задачи при сбое любого из этих этапов. При смене ревизии модели используйте новый каталог и заново формируйте эталонный набор для сверки. До окончания аренды заберите список контрольных сумм и журнал проверок. Тогда на следующую машину попадут проверяемые входные данные эксперимента, а не просто папка с пометкой «скачивание завершено».

Часто задаваемые вопросы