エンジニアリング調査：Laya のトランスフォーマーを ANE へ移す

研究環境：Apple M3 Max（40 GPU コア、128 GiB ユニファイドメモリ）、macOS 27.2、Core ML Tools 9.0、PyTorch 2.7.0、NumPy 2.1.3。これは experiments/ane_engineering/ 配下の独立したプロトタイプであり、リリース済みランタイムは変更されていません。

現在の結果

固定の B=1, L=96 の多言語プロトタイプは、Core ML の想定された計算計画において、完全なエンコーダ、意思決定ヘッド、スコアラーを Neural Engine に割り当てることに成功しました：6,390 の非定数演算が ANE を優先し、推定コスト重みの合計は約 1 になります。残りの 3,809 のエントリは定数です。元の列挙シェイプ/SDPA エクスポートは、CPU_AND_NE の下で割り当てられた 1,318 の演算すべてについて CPU を優先しました。988 の個々の演算が ANE を対応デバイスとして挙げているにもかかわらず、です。

プロトタイプの完全な予測経路は、50 回のスクリーニング呼び出しで 5.167 ms p50 を測定しました。トークン化、埋め込みルックアップ、アテンションマスク、ANE 推論、CPU アクションヘッド計算、較正、フォーマットを含みます。孤立した Core ML ボディは、合成埋め込みで 30 回の呼び出しにわたり 4.403 ms p50 を測定しました。後者はコンポーネントの計測であり、エンドツーエンドの速度の主張ではありません。モデルのロードとコンパイルは、両方のウォームタイミングから除外されています。

元の FP32 ゴールデン参照の固定長サブセットで、59/59 の回答比較が一致し、8 言語と choice/score/noul の質問を含みます。較正確率の最大変化は 0.002925 で、100 回の繰り返し公開呼び出しは有限であり、丸められた同一の結果を返します。元の 63 質問のフィクスチャには、3 つの 1,024 トークンの長入力と 1 つの 147 トークンの 20 択入力が含まれ、これら 4 つの評価は L96 エクスポートによって明示的にスキップされます。このフィクスチャには重複したルーブリックが現れます。これは回帰比較であり、59 の独立したラベル付き例や、一般的なタスク精度が変わらないことの証明ではありません。

別途の L192 と L1024 のエクスポートも、割り当てられた 6,390 のボディ演算をすべて ANE に置きます。L192 は 60/60 の比較を通過し、L1024 は完全な 63/63 のゴールデンフィクスチャを通過します。どちらも 100 回の繰り返し公開呼び出しを通過し、較正確率の最大誤差は 0.002925 のままです。長入力サブセット自体の最大誤差は 0.001128 です。評価されたすべてのトークン使用量カウントは参照と一致します。

これらは直列のスクリーニング実行であり、バックエンド間のペア比較ではありません。ボディのタイミングは 5 回のウォームアップと 30 回の計測呼び出しを使い、完全な短質問のタイミングは 10 回のウォームアップと 50 回の計測呼び出しを使います。完全な経路にはホスト側の処理と入力チェックが含まれます。L96 のスクリーニングは最後の追加の入力検証チェックより前のものであり、最終的な制御された比較は現在のアダプタを使い、そのソースフィンガープリントを記録します。コンパイル/ロード時間は、変換後の各プロセスで測定され、フレームワークキャッシュが空の初回システム起動に関する約束ではありません。大きな固定グラフは、短いリクエストであってもパディングされた処理を行います。実用的なアダプタは別々の長さバケットを選ぶでしょう。すべてのリクエストを L1024 経由にすると、短入力の優位性が失われます。

実際の workload(1, long=True) フィクスチャの別途の実行は、短いリクエストをそのサイズにパディングしたものではなく、1 つの 1,024 トークンのリクエストを確認します。10 回のウォームアップ呼び出しの後、50 回の完全な予測にわたり 91.703 ms p50 / 94.776 ms p95 を測定し、丸められた出力はすべて安定したままです。過去の MLX 長入力ベンチマークは 51.98 ms p50 です。これらは同一ラウンドのペア計測ではありませんが、このスクリーニングは、現在の ANE グラフが長入力を加速するという証拠を提供しません。入力ハッシュ、実際のトークン長、現在の実験のフィンガープリント、生のタイミングは long1024-performance.json に保持されています。

生の証拠：

• 完全ボディの計算計画とコンポーネントのタイミング
• 実入力の FP32 参照検証と完全予測のタイミング
• L192 の検証
• L1024 の完全フィクスチャ検証
• 単層の計画と数値チェック
• 単一 MLP の計画と数値チェック

MLComputePlan は想定された配置を記述するものであり、ハードウェア実行トレースではありません。CPU_AND_NE は CPU と ANE を許可し、ANE 専用のスイッチではありません。この実験では、割り当てられたヘビーなボディ演算がすべて ANE を優先しますが、実行時のハードウェアテレメトリは別途評価されます。後続の Instruments 診断は Neural Engine のハードウェア活動を記録しましたが、そのテーブルはグローバルであり、すべてのイベントをこのモデルに帰属させることはできません。ペアの MLX 比較、電力積分の結果、トレースの限界は ANE_BENCHMARKS.md で報告されています。監視ツールによる値ゼロの ANE カウンタは、そのカウンタをこの OS/デバイスで検証しない限り、ANE 活動の不在を証明することはできません。

なぜ元のグラフは ANE の目標として不適だったか

ベースラインは、従来の B×L×C トランスフォーマーレイアウト、動的シェイプ演算、ヘッド全体をバッチ化したアテンション、Core ML の SDPA 演算子を保っています。CPU/ANE 選択の下で、その計画にはデバイス割り当てが報告されていない 24 の SDPA 演算と、多数の cast、slice、shape query、gather、transpose が含まれます。個々の演算のいくつかは ANE をサポートしますが、グラフ全体はそこに分割配置されていません。したがって、個々の演算子のデバイス対応は、有用な ANE 実行経路の十分な証拠にはなりません。

成功したプロトタイプは複数の変更を同時に行います。それは、その組み合わせが ANE 配置を可能にする証拠であり、単一の犯人演算子を特定する完了したアブレーションではありません。固定シェイプの元レイアウト SDPA と明示的アテンションの対照は、次に行う価値のある識別実験です。

Apple が公開した Transformer のガイダンスは、チャネル優先の 4D テンソル、射影のための 1×1 畳み込み、ヘッド単位のアテンション、レイアウトコピーの削減を推奨しています。これらの原則が実装の動機になりました。Apple の過去の DistilBERT の高速化は、このプロジェクトのすでに高速な MLX FP16 ベースラインに対する 10 倍の改善を証明するものではありません。Apple の ANE Transformer の記事、Apple の参照実装。

プロトタイプのアーキテクチャと数値契約

model.py は、元のチェックポイントパラメータから構築された別途のエクスポートモデルです：

• 隠れアクティベーションは B,C,1,L を使います。各密重み W[out,in] は、再学習や重み近似なしに 1×1 畳み込みカーネル K[out,in,0,0] になります。
• アテンションは個々の 64 チャネルのヘッドに分割されます。キーテンソルは一度だけ転置され、2 つの明示的な einsum が 4D レイアウトを保ったまま QK と AV を計算します。Softmax はキー軸、次元 1 にわたって実行されます。
• RoPE は各ヘッドを 32 チャネルの 2 つの半分に分割します。その余弦、正弦、基底は元のモデルから来ており、多言語のローカル theta 160000 を含みます。
• チャネル正規化は、元の normalized * weight + bias の順序と epsilon を保持します。Apple の参照 LayerNorm にある、順序の異なるアフィン式や任意のクリッピングはコピーしません。
• エンコーダの最初のアテンション正規化は恒等のままです。エンコーダは正確な erf GELU を使い、2 つの意思決定ヘッド FFN は ReLU を保ちます。
• 完全アテンションとスライディングウィンドウマスクは、有効キーのマスキングとパディング済みクエリの規則を保持します。ローカル半径は local_attention // 2 から読み取られます。
• 最後のマーカー gather は、外部で用意した one-hot セレクタと 4D einsum を使って表現されます。グラフは非アクティブなスロットを含めて 32 のマーカースロットを保ち、ホストは非アクティブなロジットを元の -1e4 の値に置き換えます。

元の FP32 エンコーダ層全体とその BC1S の対応物は、PyTorch レイアウトチェックで最大 2.29e-5 だけ異なりました。FP16 Core ML の層出力誤差はより大きく、この精度/バックエンドの変更に対して予想どおりでした。ボディプローブは当初、自己比較を含んでいました。その無効な証拠は削除され、その全身の PyTorch レイアウトチェックは明示的に not_measured です。実際の完全モデル検証は、代わりに保存された元の FP32 ロジットと意思決定に対して行われます。

test_ane_layout.py の独立した CPU 回帰はさらに、完全な小さな ConvBody を元の DecisionModel と比較し、明示的アテンションと SDPA アテンションのオラクル、3 つの質問タイプ、変更されたパディング値、非ゼロの正規化バイアス、複数の RoPE 基底、非既定の正規化 epsilon を使います。これらのテストは、完全チェックポイントの FP16 ハードウェア精度と混同することなく、レイアウトとマスキングのセマンティクスをカバーします。5 つのレイアウトテストはすべてローカルで通過しました。

アテンションプロトタイプは、有限の -1e4 マスクバイアスを加えます。これは限られた検証済みアクティベーションに対して意図したマスク挙動を持ちますが、任意の極端な入力に対してマスクされたスコアを -1e4 や -infinity に置き換えることとビット単位で同一ではありません。同様に、Core ML FP16 実行が元の FP32 モデルとビット単位で同一であるとは主張しません。

runtime.py は、層ごとのデバイス遷移ではなく、トランスフォーマー全体に対して 1 つの CPU→ANE→CPU 境界を提供します：

1. CPU は各質問をトークン化し、要求された埋め込み行だけを集め、固定シェイプの加算マスク、タイプベクトル、マーカーセレクタを構築します。質問間で文脈的な隠れ状態や K/V を再利用しません。
2. 1 回の Core ML 呼び出しが、埋め込み正規化、22 のエンコーダ層すべて、両方の意思決定ヘッド層、スコアリング畳み込みを実行します。
3. CPU は 未較正の生ロジット softmax からアクション特徴量を導出し、小さな元のアクションヘッドを FP32 で erf GELU とともに実行します。その後、公開の較正と出力フォーマットが既存の実装を使います。

エクスポートされたボディは FP16 演算で、小さなホストアクションヘッドは FP32 です。埋め込みルックアップは元の保存済み重みを使います。ソースの safetensors ファイルは 169 の FP16 テンソルと 1 つの FP32 テンソルを含みます。「FP32 参照」は元の PyTorch 実行を記述するもので、元のチェックポイントが完全に FP32 で保存されているという主張ではありません。この混合精度の境界は、プロトタイプの数値契約の一部です。飽和したフィクスチャでのアクション確率差ゼロは、アクションロジットの同一性を証明しません。

アダプタは、互換性のないパッケージ次元、範囲外の ID/マーカー、無効なマスク値、空のアテンションキー行、固定容量を超える入力長を拒否します。既定の入力上限は 96 トークンで、バッチサイズは 1 です。複数の質問は逐次実行されます。短いエクスポートに収めるためにリクエストを黙って切り詰めることは決してありません。

再現、出所、品質ゲート

固定された .venv で、リポジトリルートから実行します。生成されたパッケージは Git に無視されます。重複した埋め込み NPZ や大きな重み成果物は必要ありません。probe.py は既存の出力ディレクトリを拒否します。新しいエクスポートは、元の重み/設定の SHA256 値、パッケージ内容ハッシュ、シェイプ、ツールのバージョン、実験ソースのフィンガープリントをマニフェストに記録します。

再現コマンドは artifacts/ane-repro/ 配下に書き込みます。コミットされた実験ディレクトリにはすでにレポートとマニフェストが含まれているためです。エクスポートを繰り返す際は別の新しいディレクトリを選んでください。どちらのエクスポーターも既存のパッケージを黙って再利用しません。

変換、開発、圧縮研究の依存は pip install -e '.[convert,dev,research]'（または対応する uv sync の extras）でインストールします。research 追加分は kmeans1d==0.4.0 を固定します。この任意依存は FP16 のグループ化 K-means 実験に使われ、新しいマニフェストに記録されます。

既定では、変換は laya-multilingual をリポジトリの固定された Hugging Face チェックポイントに解決し、必要に応じてダウンロードします。既存のローカルファイルを使うには --source /path/to/checkpoint を渡します。検証は既定で、パッケージマニフェストに記録されたソースディレクトリを使います。ANEAgent ランタイム自体はローカルファイルのみを受け付け、ロード前に元の重み/設定のハッシュ、固定シェイプ、パッケージ内容ハッシュを検証します。検証はまた、ソース重みのハッシュが異なるゴールデン参照を拒否します。測定された多言語のソース重み SHA256 は 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

バリデータは、評価されたすべての argmax 意思決定が一致すること、較正確率とアクション確率の誤差が <=0.02 であること、有限の出力、変わらないトークン使用量、同一の繰り返し公開出力を要求します。失敗した候補は passed: false を書き込み、失敗として終了します。スキップされたケースは、固定シェイプのサブセットが通っても未検証のままです。最初の L96 レポートは、測定後にこれらのゲートを明示的に適用しており、記録されたタイミングは変えずにそのように印が付けられています。

より長い固定エクスポートと、その完全なゴールデン参照検証コマンド：

.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

これらのコマンドは、上記に挙げた独立してチェックされたより長いエクスポートを再現します。それらの配置と数値忠実性は、L96 グラフとは別にチェックされました。短いリクエストを 1024 トークンにパディングすることは、提案された本番ポリシーではありません。

圧縮のスクリーニングと 10 倍目標

palettize.py は独立した重みのみのパレット変種を用意し、8 ビットの一様ルックアップテーブルを安価な最初のスクリーニング候補とします。2048 要素より大きい畳み込み重みだけが選択され、RoPE 定数、正規化、アクティベーション演算、ホストアクション重みは変更されません。グループ化された出力チャネルは別々のルックアップテーブルを使います。K-means はより高価なモードです。これらの FP16 グループに対して、インストールされた kmeans1d 実装を使います。Core ML Tools 9.0 は num_kmeans_workers > 1 のとき、プロセスプールの starmap を介して独立したグループを並列化します。実験は、オフライン K-means に 8 ワーカー、ワーカーあたり 1 つの数値ライブラリスレッドを使います。ワーカー数はエクスポートのスループットを変えますが、意図されたコードブックの目的関数は変えません。低ビットの 6/4 ビット候補は別の近似成果物であり、正確な実装ではありません。

圧縮サイズは エクスポートされたトランスフォーマーボディパッケージを指します。元の 1 億 9606 万エントリの埋め込みテーブルはホストにとどまり、リクエストごとに要求された行だけがルックアップされ、小さなホストアクション重みは変更されません。ボディパッケージがおよそ 2 倍縮むことは、チェックポイント全体、ランタイムメモリ、リクエストあたりエネルギーの 2 倍削減ではありません。

.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

2 つの一様 W8 スクリーニングはどちらも 59 のフィクスチャ argmax 意思決定をすべて保持し、100 回の繰り返し呼び出しを通過しますが、確率ゲートに失敗します：グループサイズ 32 は 0.023612 の誤差に達し、グループサイズ 4 は 0.033858 に達します。W8 K-means/group32 のスクリーニングは、最大誤差 0.014393、59/59 の意思決定で通過します。飽和したアクション確率は、その K-means 候補について最大 16.24 のアクションロジット差を隠します。この小さな回帰フィクスチャを通過しても、一般的な較正やタスク精度の保持を証明するものではありません。圧縮候補は、別個に識別された近似モデルのままです。

固定の W6 K-means/group32 候補も 59/59 の argmax 意思決定と安定した繰り返し出力を保ちますが、最大確率誤差 0.052243 で失敗します。その最大アクションロジット誤差は 76.62 です。したがって、重みを 6 ビットに減らしても変更されていない受け入れゲートを満たしません。短リクエストのスクリーニングレイテンシは FP16 に近いままですが。

W4 K-means/group32 も 59/59 の意思決定を保ちますが、最大確率誤差は 0.200221 に、最大アクションロジット誤差は 605.30 に上昇します。同じゲートに失敗します。5 つの圧縮スクリーニングすべてで argmax の変化がないことは、このフィクスチャの飽和した意思決定だけでは不十分な受け入れテストである理由を示しています。

すべての圧縮変種は、6,390 の NE 優先のデバイス割り当て演算、59/59 のフィクスチャ argmax 一致、100 回の安定した繰り返し呼び出しを保ちます。残りの計画エントリには、定数と重み LUT 再構築式が含まれます。配置メタデータだけでは、リクエスト中にどれだけの圧縮データが DRAM から移動するかを証明できません。3 つの K-means エクスポートは、8 つのオフラインワーカーでそれぞれ 186.50、56.13、23.90 秒かかります。セットアップとワーカープロセスは、すべての推論計測の前に終了します。

これらの 50 回呼び出しの直列スクリーニングは、高速化比率を証明するものではありません。FP16 のスクリーニングは最後の入力検証チェックより前のもので、スクリーニングは交互に配置されていません。W8 K-means は、ANE_BENCHMARKS.md のより強力な同一セッション比較のための唯一の圧縮ファイナリストです。却下された変種は、精度の境界の証拠として保持されるもので、推奨されるデプロイではありません。圧縮はこの多言語 L96 サブセットについてのみ検証されています。上記の 63 質問 L1024 の結果は、別途の FP16 エクスポートに関するものです。

Core ML のパレット圧縮は、インデックス付きルックアップテーブルから浮動小数点重みを再構築します。保存テンソルが小さいことだけでは、より高速な推論やより低いエネルギーを証明しません。すべての変種には、同じ精度ゲート、新しい計算計画、ペアのエンドツーエンド速度/エネルギー比較が必要です。Core ML パレット化のドキュメント。

最初の ANE 配置の成功は、信頼できる最適化経路を確立するものであり、10 倍の結果ではありません。比較は MLX FP16 を使い、より高速な場合はそのコンパイル済み変種も含め、等しいタスク範囲を報告しなければなりません。電力は完全なリクエストにわたって積分しなければなりません。総システムエネルギーと、あればアイドル差し引き後の推定の両方を、その測定限界とともに報告すべきです。独立した数学的レビューは ANE_MATH.md です。

手動実装が確立するもの

ここでの有用な手作業は、Core ML が ANE にマッピングできるレイアウトへ計算グラフを完全に、独立して検証しながら書き直したことです。これは学習済みパラメータを保ちながら実行配置を変えます。元のグラフで compute_units を変えるよりはるかに効果的です。24 の逐次的なアテンション/MLP ブロックやその密な射影の処理を取り除くものではありません。

長いリクエストに対する次の正確な候補は、固定のクエリ/キータイルと元のパディングおよび RoPE 規則を使って実装される、実際にローカルウィンドウだけを訪れるアテンションです。現在のグラフは依然として密なスコア行列を計算し、局所マスクを適用します。タイル化されたグラフはその処理を減らせる可能性がありますが、より多くのスライス、境界、小さな縮約が ANE のスケジューリングを損なうかもしれません。その配置、精密度、エンドツーエンドの利得は未測定のままです。さらなる重み圧縮には、上記の失敗の後の較正または品質回復が必要です。蒸留や層数の削減は新しいモデルを導入し、より広範なタスク品質評価を必要とします。これらの未実装の方向はいずれも、今日 10 倍の利得の証拠を提供しません。