Apple Neural Engineでトレーニング — リバースエンジニアリングされたプライベートAPI • blog.dopana

maderix ↗によるANE ↗プロジェクトは、リバースエンジニアリングされたプライベートAPIを介してApple Neural Engine（ANE）上で直接ニューラルネットワークをトレーニングする画期的な研究コンセプト実証です。CoreML、Metal、GPUは一切使用しません。

M4のANEは15.8 TFLOPS FP16（実測18.6 TOPS）を達成しますが、AppleはCoreMLを通じた推論専用に制限しています。このプロジェクトはその障壁を打破し、問題はハードウェア能力ではなくソフトウェアサポートにあることを証明しました。

ソフトウェアスタック#

CoreMLだけが唯一の経路ではありません。AppleNeuralEngine.framework内の_ANEClientクラスは、コンパイル→ロード→評価のパイプラインに直接アクセスできます：

id client = [_ANEClient sharedConnection];
id model = [_ANEModel modelAtURL:compiledURL key:@"mykey"];

[client compileModel:model options:@{
    @"kANEFModelType": @"kANEFModelMIL",
    @"kANEFNetPlistFilenameKey": @"model.mil"
} qos:21 error:&err];

[client loadModel:model options:@{} qos:21 error:&err];

objc

.mlmodelcファイルにコンパイルしてからロードする代わりに、_ANEInMemoryModelDescriptorを使用するとMILテキストをメモリ内で直接コンパイルできます。ディスクラウンドトリップは不要で、これがトレーニングを可能にする鍵です。

ANEはI/OにIOSurfaceを使用します。これはGPUテクスチャと同じ共有メモリ機構で、GPU↔ANE間のゼロコピーパイプラインを実現します。

MIL — ANEの中間言語#

ANEはONNXやprotobufを受け付けません。代わりにMIL（Model Intermediate Language）を使用します。型と形状が明示されたSSA表現です：

program(1.3) {
    func main<ios18>(
        tensor<fp16, [1, 1024, 1, 1024]> x,
        tensor<fp16, [1, 1024, 1, 1024]> w
    ) {
        tensor<fp16, [1, 1024, 1, 1024]> out =
            matmul(transpose_x = false, transpose_y = false,
                   x = x, y = w);
    } -> (out);
}

text

テンソルレイアウトはNCDHW + Interleave形式：[Batch, Channels, Depth, Height, Width]。1024×1024行列は4Dで[1, 1024, 1, 1024]になります。

トレーニングアーキテクチャ#

動的パイプラインは共有ANEカーネルを使用し、重みを空間次元にパックします。重みが変更されても再コンパイルは不要です。

MHAモデル（Stories110M）— レイヤーあたり6カーネル：

カーネル	機能
`sdpaFwd`	QKV投影 + SDPA + 出力投影
`ffnFused`	SwiGLU FFN（W1, W3, SiLU, W2）
`ffnBwdW2t` / `ffnBwdW13t`	FFN逆伝播（メモリ分割）
`sdpaBwd1` / `sdpaBwd2`	SDPA逆伝播

GQAモデル（Qwen3-0.6B）— レイヤーあたり10カーネル。グループ化クエリアテンション用に個別のwoFwd、qBwd、kvBwdを追加。

CPUが担当：RMSNorm順伝播/逆伝播、残差接続（DeepNet αスケーリング）、損失計算、dW勾配蓄積（cblas_sgemm）、Adam最適化。

パフォーマンス結果#

トレーニングスループット（M4）：

モデル	パラメータ数	ms/ステップ
Stories110M	109M	91 ms
Qwen3-0.6B	596M	412 ms

INT8 W8A8量子化 — 1.88倍の高速化：

構成	FP16	INT8	高速化
128x conv 512ch 64x64	18.6 TOPS, 14.8ms	35.1 TOPS, 7.8ms	1.88x

INT8アクティベーションはquantize/dequantizeによりタイル間のL2 SRAM帯域幅を半減します。

主要な最適化#

Channel-first CPUレイアウト — IOSurface [1,C,1,S]形式に一致、転置オーバーヘッドを完全排除
vDSPベクトル化RMSNorm — 10倍高速化（6.7ms → 0.7ms）
GCD非同期cblasオーバーラップ — dW勾配sgemmをANE評価と並列実行
Deferred cblas wait — 待機を次のステップの順伝播に繰り越し
ANE RMSNorm融合 — RMSNormをMIL演算で順伝播カーネルに組み込み
フォワードタップ — Q、K、V、アテンションスコアを結合出力で公開、CPU再計算を回避
exec()再起動 — プロセスあたり約119回のANEコンパイル制限を回避

制限事項#

SDPA因果マスキング — ANEハードウェアはattn_maskを無視。因果的注意はQ@K^T（ANE）→ マスク+softmax（CPU）→ scores@V（ANE）に分解
約119回のコンパイル制限 — コンパイラがリソースをリーク。チェックポイント付きexec()再起動で回避
FP16勾配アンダーフロー — 逆伝播行列積がfp16でアンダーフロー。グローバル損失スケーリングで修正
低い利用率 — ピークの約5-9%のみ、多くの要素単位演算はCPUにフォールバック

参考文献#

GitHub — maderix/ANE ↗
Inside the M4 Apple Neural Engine, Part 1: Reverse Engineering ↗
Inside the M4 Apple Neural Engine, Part 2: Benchmarks ↗
Inside the M4 Apple Neural Engine, Part 3: Training ↗
hollance/neural-engine ↗ — ANEコミュニティドキュメント
apple/ml-ane-transformers ↗ — AppleのANE用リファレンストランスフォーマー