Investigation d'ingénierie : déplacer le transformer de Laya sur l'ANE

Environnement de recherche : Apple M3 Max (40 cœurs GPU, 128 GiB de mémoire unifiée), macOS 27.2, Core ML Tools 9.0, PyTorch 2.7.0 et NumPy 2.1.3. C’est un prototype indépendant sous experiments/ane_engineering/ ; le runtime publié est inchangé.

Résultat actuel

Un prototype multilingue fixe B=1, L=96 assigne avec succès l’encodeur complet, la tête de décision et le scorer au Neural Engine dans le plan de calcul anticipé de Core ML : 6,390 opérations non constantes préfèrent l’ANE, avec des poids de coût estimés dont la somme vaut environ 1. Les 3,809 entrées restantes sont des constantes. L’export d’origine à formes énumérées/SDPA préférait le CPU pour les 1,318 opérations assignées sous CPU_AND_NE, alors que 988 opérations individuelles listent l’ANE comme appareil pris en charge.

Le chemin de prédiction complet du prototype a mesuré 5.167 ms p50 sur 50 appels de criblage, incluant la tokenisation, la recherche d’embedding, les masques d’attention, l’inférence ANE, le calcul de tête d’action sur le CPU, la calibration et le formatage. Le corps Core ML isolé a mesuré 4.403 ms p50 sur 30 appels avec des embeddings synthétiques. Ce dernier est une mesure de composant et n’est pas une affirmation de vitesse de bout en bout. Le chargement du modèle et la compilation sont exclus des deux mesures à chaud.

Sur le sous-ensemble à longueur fixe de la référence dorée FP32 d’origine, 59/59 comparaisons de réponses concordent, incluant huit langues et des questions choice/score/noul. Le plus grand changement de probabilité calibrée est de 0.002925, et 100 appels publics répétés sont finis et renvoient des résultats arrondis identiques. La fixture d’origine de 63 questions contient trois longues entrées de 1,024 tokens et une entrée de 147 tokens à 20 options ; ces quatre évaluations sont explicitement ignorées par l’export L96. Des rubriques répétées apparaissent dans cette fixture. Ce sont des comparaisons de régression, pas 59 exemples étiquetés indépendants ni une preuve d’exactitude générale inchangée sur les tâches.

Des exports L192 et L1024 distincts placent aussi les 6,390 opérations de corps assignées sur l’ANE. L192 passe 60/60 comparaisons ; L1024 passe la fixture dorée complète 63/63. Les deux passent 100 appels publics répétés et l’erreur maximale de probabilité calibrée reste 0.002925. Le sous-ensemble à entrée longue a lui-même une erreur maximale de 0.001128. Tous les comptes d’usage de tokens évalués correspondent à la référence.

Ce sont des exécutions de criblage en série, pas des comparaisons appariées entre backends. Les temps de corps utilisent 5 appels d’échauffement et 30 mesurés ; les temps de question courte complète utilisent 10 appels d’échauffement et 50 mesurés. Le chemin complet inclut le travail de l’hôte et les contrôles d’entrée. Le criblage L96 précède les derniers contrôles de validation d’entrée supplémentaires ; la comparaison contrôlée finale utilise l’adaptateur actuel et enregistre son empreinte source. Les temps de compilation/chargement sont mesurés dans chaque processus après conversion, pas une promesse sur le tout premier démarrage système avec des caches de framework vides. Les grands graphes fixes effectuent un travail complété même pour des requêtes courtes. Un adaptateur pratique sélectionnerait des buckets de longueur séparés ; router chaque requête via L1024 supprimerait l’avantage des entrées courtes.

Une exécution distincte de la vraie fixture workload(1, long=True) confirme une requête de 1,024 tokens, plutôt qu’une requête courte complétée à cette taille. Elle mesure 91.703 ms p50 / 94.776 ms p95 sur 50 prédictions complètes après dix appels d’échauffement ; toutes les sorties arrondies restent stables. Le benchmark MLX historique sur entrée longue est de 51.98 ms p50. Ce ne sont pas des mesures appariées du même tour, mais ce criblage n’apporte aucune preuve que le graphe ANE actuel accélère les entrées longues. Le hash d’entrée, la longueur de tokens réelle, les empreintes d’expérience actuelles et les temps bruts sont conservés dans long1024-performance.json.

Preuves brutes :

• Plan de calcul du corps complet et temps des composants
• Validation de référence FP32 sur entrées réelles et temps de prédiction complète
• Validation L192
• Validation L1024 sur fixture complète
• Plan d’une seule couche et contrôles numériques
• Plan d’un seul MLP et contrôles numériques

MLComputePlan décrit un placement anticipé, pas une trace d’exécution matérielle. CPU_AND_NE autorise CPU et ANE ; ce n’est pas un interrupteur ANE uniquement. Dans cette expérience, toutes les opérations lourdes du corps assignées préfèrent l’ANE, mais la télémétrie matérielle à l’exécution est évaluée séparément. Le diagnostic Instruments ultérieur a enregistré une activité matérielle du Neural Engine ; sa table est globale et ne peut pas attribuer chaque événement à ce modèle. La comparaison MLX appariée, les résultats d’intégration de puissance et les limites des traces sont rapportés dans ANE_BENCHMARKS.md. Un compteur ANE à zéro d’un outil de surveillance ne peut pas établir une absence d’activité ANE sans valider ce compteur sur ce système d’exploitation/cet appareil.

Pourquoi le graphe d’origine était une mauvaise cible pour l’ANE

La baseline préserve une disposition de transformer classique B×L×C, des opérations à forme dynamique, une attention en lot sur toutes les têtes et l’opérateur SDPA de Core ML. Sous la sélection CPU/ANE, son plan contient 24 opérations SDPA sans assignation d’appareil rapportée, plus de nombreux casts, découpages, requêtes de forme, gathers et transpositions. Certaines opérations individuelles prennent en charge l’ANE, mais le graphe dans son ensemble n’est pas partitionné dessus. La prise en charge d’appareil pour des opérateurs individuels est donc une preuve insuffisante d’un chemin d’exécution ANE utile.

Le prototype réussi change plusieurs choses à la fois. C’est la preuve que la combinaison permet le placement ANE, pas une ablation achevée identifiant un unique opérateur fautif. Des contrôles SDPA à disposition d’origine et à forme fixe et d’attention explicite sont les prochaines expériences discriminantes utiles.

Les recommandations publiées d’Apple sur les Transformers préconisent des tenseurs 4D channel-first, des convolutions 1×1 pour les projections, une attention par tête et moins de copies de disposition. Ces principes ont motivé l’implémentation ; les accélérations historiques de DistilBERT chez Apple n’établissent pas une amélioration de 10× par rapport à la baseline MLX FP16 déjà rapide de ce projet. Article d’Apple sur le Transformer ANE, implémentation de référence d’Apple.

Architecture du prototype et contrat numérique

model.py est un modèle d’export distinct construit à partir des paramètres du checkpoint d’origine :

• Les activations cachées utilisent B,C,1,L. Chaque poids dense W[out,in] devient un noyau de convolution 1×1 K[out,in,0,0] sans réentraînement ni approximation de poids.
• L’attention est divisée en têtes individuelles de 64 canaux. Le tenseur de clé est transposé une fois, et deux einsums explicites calculent QK et AV en conservant la disposition 4D. Le softmax s’exécute sur l’axe des clés, la dimension 1.
• RoPE divise chaque tête en ses deux moitiés de 32 canaux. Ses cosinus, sinus et bases proviennent du modèle d’origine, y compris le theta local multilingue 160000.
• La normalisation par canal préserve l’ordre d’origine normalized * weight + bias et l’epsilon. Elle ne copie pas l’expression affine d’ordre différent ni l’écrêtage optionnel de la LayerNorm de référence d’Apple.
• La première norme d’attention de l’encodeur reste l’identité ; l’encodeur utilise un GELU erf exact, tandis que les deux FFN de la tête de décision conservent ReLU.
• L’attention complète et les masques de fenêtre glissante préservent le masquage des clés valides et la règle des requêtes complétées. Le rayon local est lu depuis local_attention // 2.
• Le gather de marqueur final est exprimé à l’aide d’un sélecteur one-hot préparé en externe et d’un einsum 4D. Le graphe conserve 32 emplacements de marqueur, y compris les emplacements inactifs ; l’hôte remplace les logits inactifs par la valeur d’origine -1e4.

Une couche d’encodeur FP32 d’origine entière et sa contrepartie BC1S différaient d’au plus 2.29e-5 dans le contrôle de disposition PyTorch. L’erreur de sortie d’une couche Core ML FP16 était plus grande, comme attendu pour ce changement de précision/backend. La sonde de corps contenait à l’origine une auto-comparaison ; cette preuve invalide a été supprimée et son contrôle de disposition PyTorch sur tout le corps est explicitement not_measured. La validation réelle du modèle complet se fait plutôt par rapport aux logits et décisions FP32 d’origine stockés.

Les régressions CPU indépendantes dans test_ane_layout.py comparent en outre un petit ConvBody complet avec le DecisionModel d’origine, en utilisant des oracles d’attention explicite et SDPA, trois types de question, des valeurs de padding modifiées, des biais de norme non nuls, plusieurs bases RoPE et un epsilon de norme non par défaut. Ces tests couvrent la disposition et la sémantique de masquage sans les confondre avec la précision matérielle FP16 du checkpoint complet. Les cinq tests de disposition ont passé localement.

Le prototype d’attention ajoute un biais de masque fini de -1e4. Cela a le comportement de masque voulu sur les activations finies validées, mais ce n’est pas une identité bit à bit avec le remplacement des scores masqués par -1e4 ou -infini pour des entrées extrêmes arbitraires. De même, l’exécution Core ML FP16 n’est pas déclarée identique bit à bit au modèle FP32 d’origine.

runtime.py fournit une seule frontière CPU→ANE→CPU pour tout le transformer, plutôt qu’une transition d’appareil par couche :

1. Le CPU tokenise chaque question, ne rassemble que les lignes d’embedding demandées et construit des masques additifs à forme fixe, des vecteurs de type et des sélecteurs de marqueur. Il ne réutilise pas les états cachés contextuels ni les K/V entre les questions.
2. Un seul appel Core ML effectue la normalisation d’embedding, les 22 couches d’encodeur, les deux couches de tête de décision et les convolutions de scoring.
3. Le CPU dérive les caractéristiques d’action à partir du softmax des logits bruts non calibrés et exécute la petite tête d’action d’origine en FP32 avec GELU erf. La calibration publique et le formatage de sortie utilisent ensuite l’implémentation existante.

Le corps exporté a un calcul FP16, tandis que la petite tête d’action hôte est en FP32. La recherche d’embedding utilise les poids d’origine stockés. Le fichier safetensors source contient 169 tenseurs FP16 et un tenseur FP32 : “FP32 reference” décrit l’exécution PyTorch d’origine, pas l’affirmation que le checkpoint d’origine est entièrement stocké en FP32. Cette frontière de précision mixte fait partie du contrat numérique du prototype. Des différences nulles de probabilité d’action sur des fixtures saturées ne prouvent pas l’identité des logits d’action.

L’adaptateur rejette les dimensions de paquet incompatibles, les ID/marqueurs hors plage, les valeurs de masque invalides, les lignes de clé d’attention vides et les longueurs d’entrée dépassant sa capacité fixe. Sa limite d’entrée par défaut est de 96 tokens et sa taille de lot de un ; plusieurs questions s’exécutent séquentiellement. Il ne tronque jamais silencieusement une requête pour tenir dans l’export plus court.

Reproduction, provenance et portails de qualité

Exécute depuis la racine du dépôt dans le .venv épinglé. Les paquets générés sont ignorés par Git ; aucun NPZ d’embedding dupliqué ni gros artefact de poids n’est requis. probe.py refuse un répertoire de sortie existant. Les nouveaux exports consignent les valeurs SHA256 des poids/config d’origine, les hashs de contenu de paquet, la forme, les versions des outils et les empreintes de source d’expérience dans un manifeste.

Les commandes de reproduction écrivent sous artifacts/ane-repro/, parce que les répertoires d’expérience committés contiennent déjà des rapports et des manifestes. Choisis un autre répertoire neuf pour répéter un export ; aucun des deux exportateurs ne réutilise silencieusement un paquet existant.

Installe les dépendances de conversion, de développement et de recherche sur la compression avec pip install -e '.[convert,dev,research]' (ou les extras uv sync correspondants). L’extra research épingle kmeans1d==0.4.0 ; cette dépendance optionnelle est utilisée pour les expériences K-means regroupé FP16 et consignée dans les nouveaux manifestes.

Par défaut, la conversion résout laya-multilingual vers le checkpoint Hugging Face épinglé du dépôt et le télécharge si nécessaire. Passe --source /path/to/checkpoint pour utiliser des fichiers locaux existants. La validation prend par défaut le répertoire source consigné dans le manifeste du paquet. Le runtime ANEAgent lui-même n’accepte que des fichiers locaux ; il vérifie les hashs des poids/config d’origine, la forme fixe et le hash de contenu du paquet avant le chargement. La validation rejette aussi une référence dorée dont le hash de poids source diffère. Le SHA256 mesuré des poids sources multilingues est 9d628fd971b700382ac6f65920a86f149777b2e748e0c955fb3b19695aa8f204.

# Small placement probes, then the complete model.
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind mlp --length 96 --output artifacts/ane-repro/mlp96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind layer --length 96 --output artifacts/ane-repro/layer96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 96 --output artifacts/ane-repro/body96
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --length 96 --repeats 100 \
  --output artifacts/ane-repro/validation96.json

Le validateur exige que toutes les décisions argmax évaluées concordent, des erreurs de probabilité calibrée et d’action <=0.02, des sorties finies, un usage de tokens inchangé et des sorties publiques répétées identiques. Un candidat échoué écrit passed: false et se termine en échec. Les cas ignorés restent non validés même si le sous-ensemble à forme fixe passe. Le rapport L96 initial avait ces portails appliqués explicitement après mesure ; il est marqué en conséquence, sans changer ses temps enregistrés.

Exports fixes plus longs et leurs commandes complètes de validation de la référence dorée :

.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 192 --output artifacts/ane-repro/body192
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body192/model.mlpackage --length 192 \
  --output artifacts/ane-repro/validation192.json
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 1024 --output artifacts/ane-repro/body1024
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 \
  --output artifacts/ane-repro/validation1024.json
.venv/bin/python -m experiments.ane_engineering.benchmark \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 --long \
  --output artifacts/ane-repro/long1024-performance.json

Ces commandes reproduisent les exports plus longs vérifiés indépendamment listés ci-dessus. Leur placement et leur fidélité numérique ont été vérifiés séparément du graphe L96 ; compléter les requêtes courtes à 1024 tokens n’est pas la politique de production proposée.

Criblage de compression et l’objectif 10×

palettize.py prépare des variantes de palette indépendantes portant uniquement sur les poids, avec des tables de correspondance uniformes 8 bits comme premier candidat de criblage peu coûteux. Seuls les poids de convolution de plus de 2048 éléments sont sélectionnés ; les constantes RoPE, la normalisation, l’arithmétique d’activation et les poids d’action de l’hôte restent inchangés. Les canaux de sortie regroupés utilisent des tables de correspondance séparées. K-means est un mode plus coûteux. Il utilise l’implémentation kmeans1d installée pour ces groupes FP16. Core ML Tools 9.0 parallélise les groupes indépendants via un starmap de pool de processus quand num_kmeans_workers > 1 ; les expériences utilisent huit workers pour le K-means hors ligne et un thread de bibliothèque mathématique par worker. Le nombre de workers change le débit d’export, pas l’objectif de codebook visé. Les candidats à 6/4 bits sont des artefacts approximatifs distincts, pas des implémentations exactes.

Les tailles de compression se réfèrent au paquet de corps de transformer exporté. La table d’embedding d’origine de 196.608 millions d’entrées reste sur l’hôte, avec seulement les lignes demandées recherchées pour chaque requête, et les petits poids d’action de l’hôte restent inchangés. Un paquet de corps réduit d’environ deux fois n’est pas une réduction de deux fois du checkpoint entier, de la mémoire du runtime ou de l’énergie par requête.

.venv/bin/python -m experiments.ane_engineering.palettize \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --bits 8 --mode uniform --group-size 32 \
  --output artifacts/ane-repro/body96-w8
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96-w8/model.mlpackage --length 96 \
  --output artifacts/ane-repro/validation96-w8.json

# Independent K-means candidates; inspect each validation exit status.
for bits in 8 6 4; do
  VECLIB_MAXIMUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 \
    .venv/bin/python -m experiments.ane_engineering.palettize \
    --package artifacts/ane-repro/body96/model.mlpackage \
    --bits "$bits" --mode kmeans --group-size 32 --workers 8 \
    --output "artifacts/ane-repro/body96-w${bits}km"
  .venv/bin/python -m experiments.ane_engineering.validate \
    --package "artifacts/ane-repro/body96-w${bits}km/model.mlpackage" --length 96 \
    --output "artifacts/ane-repro/validation96-w${bits}km.json"
done

Les deux criblages W8 uniformes conservent tous les 59 décisions argmax de la fixture et passent 100 appels répétés, mais échouent au portail de probabilité : la taille de groupe 32 atteint une erreur de 0.023612 et la taille de groupe 4 de 0.033858. Le criblage W8 K-means/group32 passe avec une erreur maximale de 0.014393 et 59/59 décisions. Les probabilités d’action saturées masquent des différences de logits d’action allant jusqu’à 16.24 pour ce candidat K-means ; passer cette petite fixture de régression n’établit pas une calibration générale ni une exactitude de tâche préservées. Les candidats compressés restent des modèles approximatifs identifiés séparément.

Le candidat fixe W6 K-means/group32 conserve aussi 59/59 décisions argmax et des sorties répétées stables, mais échoue avec une erreur maximale de probabilité de 0.052243. Son erreur maximale de logits d’action est de 76.62. Réduire les poids à six bits ne satisfait donc pas le portail d’acceptation inchangé, même si sa latence de criblage sur requête courte reste proche de FP16.

W4 K-means/group32 conserve aussi 59/59 décisions, mais l’erreur maximale de probabilité monte à 0.200221 et l’erreur maximale de logits d’action à 605.30. Il échoue au même portail. L’absence de changements d’argmax dans les cinq criblages compressés montre pourquoi les décisions saturées de cette fixture, seules, constituent un test d’acceptation inadéquat.

Toutes les variantes compressées conservent 6,390 opérations assignées à un appareil préféré pour l’ANE, 59/59 accords d’argmax de fixture et 100 appels répétés stables. Les entrées de plan restantes incluent des constantes et des expressions de reconstruction de LUT de poids ; les métadonnées de placement seules ne prouvent pas quelle quantité de données compressées circule depuis la DRAM pendant une requête. Les trois exports K-means prennent respectivement 186.50, 56.13 et 23.90 secondes avec huit workers hors ligne. La configuration et les processus workers se terminent avant chaque mesure d’inférence.

Ces criblages en série de 50 appels n’établissent pas de rapports d’accélération. Le criblage FP16 précède les derniers contrôles de validation d’entrée, et les criblages ne sont pas entrelacés. W8 K-means est le seul finaliste compressé pour la comparaison intra-session plus forte dans ANE_BENCHMARKS.md. Les variantes rejetées sont conservées comme preuve de la frontière de précision, pas comme déploiements recommandés. La compression n’a été validée que pour ce sous-ensemble multilingue L96 ; le résultat L1024 à 63 questions ci-dessus concerne l’export FP16 distinct.

La compression de palette Core ML reconstruit des poids flottants à partir de tables de correspondance indexées ; des tenseurs stockés plus petits n’établissent pas à eux seuls une inférence plus rapide ou une énergie plus faible. Chaque variante a besoin des mêmes portails de précision, d’un plan de calcul neuf et d’une comparaison appariée de vitesse/énergie de bout en bout. Documentation de palettisation Core ML.

Le placement ANE réussi initial établit un chemin d’optimisation crédible, pas un résultat de 10×. La comparaison doit utiliser MLX FP16, y compris sa variante compilée là où elle est plus rapide, et rapporter une portée de tâche égale. La puissance doit être intégrée sur des requêtes complètes. L’énergie système brute et toute estimation moins l’inactivité doivent être rapportées avec leurs limites de mesure. La revue mathématique indépendante est ANE_MATH.md.

Ce qu’établit l’implémentation manuelle

Le travail manuel utile ici est une réécriture complète et validée indépendamment du graphe de calcul vers une disposition que Core ML peut mapper sur l’ANE. Cela change le placement d’exécution tout en conservant les paramètres entraînés. C’est nettement plus efficace que de changer compute_units sur le graphe d’origine. Cela ne supprime pas les 24 blocs séquentiels d’attention/MLP ni leur travail de projection dense.

Le prochain candidat exact pour les requêtes longues est une attention qui ne visite réellement que des fenêtres locales, implémentée avec des tuiles fixes de requête/clé et les règles d’origine de padding et de RoPE. Le graphe actuel calcule encore une matrice de scores dense et applique un masque local. Un graphe tuilé pourrait réduire ce travail, mais davantage de tranches, de frontières et de petites contractions peuvent compromettre l’ordonnancement ANE ; son placement, sa précision et son bénéfice de bout en bout restent non mesurés. Une compression de poids supplémentaire nécessite une calibration ou une récupération de qualité après les échecs ci-dessus. La distillation ou moins de couches introduirait un nouveau modèle et nécessiterait une évaluation plus large de la qualité des tâches. Aucune de ces directions non implémentées n’apporte aujourd’hui la preuve d’un gain de 10×.