Vérifier les fragments de poids MLX sur un Mac cloud

Après le transfert des poids d’un modèle sur un Mac cloud, un échec de chargement n’est pas forcément dû à un manque de mémoire unifiée. Un téléchargement interrompu, des fragments provenant de versions différentes ou un index qui pointe vers d’anciens fichiers peuvent bloquer l’expérience dès la première étape. Vérifier les fichiers sans charger les tenseurs, puis lancer MLX, permet d’éviter les téléchargements répétés et les ajustements à l’aveugle des paramètres d’inférence.

Figer d’abord le répertoire des poids

Lors des vérifications sur OsxNode, regroupez les poids, la configuration et l’index d’une même expérience dans un répertoire weights distinct, sans écraser le modèle en cours d’exécution. Prévoyez Python 3 et enregistrez le script de vérification en dehors de ce répertoire. Cet article porte uniquement sur les poids safetensors et n’exécute aucun code fourni avec le modèle.

Consignez l’identifiant de révision du modèle, la configuration de quantification, la source du téléchargement et la liste des fichiers. Ces informations doivent correspondre à la version réellement obtenue : un simple nom de modèle, dont le contenu peut évoluer, ne suffit pas. Attendez que le programme de téléchargement se termine avant de lancer les contrôles, pour ne pas prendre un fragment encore en cours d’écriture pour un fichier endommagé.

La validité de la structure des fichiers, la fidélité du transfert et la capacité du modèle à effectuer une inférence sont trois résultats de validation distincts. La réussite de l’un ne valide pas les deux autres.

Vérifier les offsets et l’index à partir des en-têtes

Les huit premiers octets d’un fichier safetensors contiennent un entier non signé en petit-boutiste qui indique la longueur, en octets, de l’en-tête JSON. Les data_offsets des tenseurs sont relatifs à la zone de données située après l’en-tête, et non au début du fichier. Confondre ces deux repères conduit à signaler à tort des dépassements de limites dans des fragments valides.

Le script ci-dessous vérifie que la zone de données est couverte sans interruption, qu’aucun nom de tenseur n’apparaît dans plusieurs fragments et que les correspondances de l’index reflètent les fichiers présents. Enregistrez-le sous verify_weights.py. La limite de 16MiB imposée à l’en-tête relève ici d’une politique de validation locale, pas d’une contrainte du format. Si un en-tête valide la dépasse, vérifiez sa provenance avant d’ajuster cette limite.

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)

Exécutez python3 verify_weights.py weights et ne poursuivez que si le code de sortie est nul. Le script suppose que l’index porte le nom standard et que les fragments se trouvent à la racine du répertoire. Toute autre organisation nécessite une adaptation explicite : ne supprimez pas l’index pour contourner le contrôle.

Valider le transfert avec les empreintes calculées à la source

Le contrôle de structure ne détecte pas les altérations des valeurs à l’intérieur des tenseurs. Générez donc une liste d’empreintes sur une source de confiance, une fois l’écriture des fichiers terminée, puis transférez ces mêmes fichiers et la liste vers la machine cible. Toutes les commandes suivantes s’exécutent dans le répertoire des poids :

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

Si un index des fragments est présent, ajoutez-le également à la liste :

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

Sur la machine cible, placez-vous dans le répertoire correspondant, puis exécutez :

shasum -a 256 -c SHA256SUMS

Tous les fichiers répertoriés doivent renvoyer OK. La liste d’empreintes doit elle-même être obtenue par un canal de confiance : si un attaquant peut remplacer à la fois les fichiers et la liste, les empreintes ne permettent pas d’authentifier leur provenance. Des poids déjà endommagés à la source réussiront aussi le contrôle du transfert s’ils sont copiés à l’identique. Il faut donc conserver le contrôle de structure.

Traiter chaque type d’échec plutôt que tout retélécharger

En cas de message Short prefix ou de longueur d’en-tête anormale, vérifiez d’abord que le processus de téléchargement est terminé, puis contrôlez la taille en octets du fragment concerné. Ne complétez pas le fichier manuellement et ne renommez pas une page d’erreur de téléchargement pour la faire passer pour un fichier de poids.

Si le message Shard index mismatch apparaît, comparez un à un les noms de fichiers de l’index avec le contenu du répertoire. Un fragment oublié lors de la copie ou des fichiers d’une ancienne version laissés sur place sont des causes fréquentes. Isolez d’abord les fichiers superflus, puis relancez le contrôle. Ne modifiez pas l’index pour l’adapter à un ensemble de poids de provenance incertaine.

Si la structure est valide mais que le contrôle des empreintes échoue, récupérez uniquement les fichiers en échec, puis relancez l’ensemble des vérifications. Si l’outil de transfert permet la reprise, assurez-vous au préalable que le fichier source distant n’a pas changé de version. Sinon, le début de l’ancien fichier risque d’être assemblé avec la fin du nouveau.

Ce n’est que si les empreintes concordent et que le chargement échoue encore qu’il faut examiner l’architecture du modèle, la configuration de quantification, la version des bibliothèques d’exécution et la pression mémoire. Cet ordre permet de distinguer les problèmes de fichiers de ceux de l’environnement d’exécution, sans devoir deviner la cause à partir du seul message d’exception.

Faire des contrôles un préalable à chaque expérience

Le script ne vérifie pas la relation mathématique entre dtype, shape et la taille en octets des tenseurs, ni la cohérence des paramètres de quantification. Il reste donc nécessaire de charger le modèle dans un environnement MLX dont la version est figée, puis de tester la génération d’une sortie courte à partir d’une entrée fixe. La seule vérification des en-têtes ne remplace pas cette étape.

Conservez au minimum dans le journal de l’expérience l’identifiant de révision, la liste d’empreintes, le code de sortie du script, la version des bibliothèques d’exécution et le résultat du premier chargement. Une simple capture d’écran montrant une réussite ne suffit pas : après une modification du répertoire, il serait difficile d’expliquer les différences.

Placez le contrôle de structure et la vérification des empreintes au début du script d’expérience, et arrêtez les tâches suivantes dès que l’une de ces étapes échoue. Lors d’un changement de révision du modèle, utilisez un nouveau répertoire et recréez la référence de validation. Avant la fin de la location, récupérez la liste d’empreintes et les journaux de contrôle. La machine suivante disposera ainsi d’entrées d’expérience vérifiables, et non d’un simple dossier étiqueté « téléchargement terminé ».

Questions fréquentes