ANE — Apple Neural Engine 上でのニューラルネットワークトレーニング

GitHub: https://github.com/maderix/ANE

🔍 概要

ANE は、Apple Silicon チップ（M1/M2/M4 シリーズ）に搭載されている専用 AI アクセラレータである Apple の Neural Engine (ANE) 上で、完全な ニューラルネットワークトレーニング（フォワードパス + バックワードパス） を直接実行できることを初めて公に実証する研究プロジェクトです。

Apple は公式には CoreML フレームワークを通じて ANE を 推論専用 に制限しています。このプロジェクトは、Apple の プライベート API（_ANEClient、_ANECompiler、_ANEInMemoryModelDescriptor）を リバースエンジニアリング することで、その制限を回避し、バックプロパゲーションを含むカスタム計算グラフを ANE ハードウェア上でネイティブに実行します。

これはプロダクション用のフレームワークでは ありません。これは 概念実証 であり、ベンチマークリファレンスであり、その制限はハードウェアの制約ではなくソフトウェアによって課せられていることを証明するものです。

🎯 なぜこれが重要なのか

Apple Silicon チップ（特に M4）には、15.8 TFLOPS と評価される ANE が搭載されています。これは、Apple が推論にのみ使用を限定している、膨大な量の専用 AI コンピューティング能力です。CoreML を実行するたびに、推論のみが得られます。トレーニングは常に GPU または CPU にフォールバックします。

このプロジェクトは、「ANE でトレーニングすることは可能なのか？」という疑問に答えます。

答えははいです。そして、このリポジトリがその証拠です。

⚙️ 仕組み

このプロジェクトは、以下の手順で完全な Transformer レイヤートレーニングループ を実装しています。

MIL プログラム生成 — Objective-C で実行時に Apple の Model Intermediate Language (MIL) プログラムを構築し、畳み込み（線形レイヤー用）、行列乗算（アテンション用）、ソフトマックス、要素ごとの演算を定義します。
インメモリコンパイル — _ANEInMemoryModelDescriptor を使用して、MIL テキストと重みブロブを直接 ANE プログラムにコンパイルし、.mlmodelc ファイルをディスクに書き込む必要がありません。
IOSurface I/O — 入出力テンソルは、ANE ハードウェアがネイティブに期待するフォーマットである [1, channels, 1, spatial] fp16 形式の IOSurface 共有メモリを介して渡されます。
重み埋め込み — 重みは BLOBFILE 定数として ANE プログラムに埋め込まれ、重みが更新されるたびにバッチごとに再コンパイルされます。
勾配フロー — フォワードパスの「タップ」は、バックワードパスに必要な中間アクティベーションを公開します。バックワードカーネルは ANE 上で入力勾配（dx）を計算します。重み勾配（dW）は CPU 上で cblas_sgemm を介して計算されます。

🧱 アーキテクチャ：トレーニングステップあたり 6 つの ANE カーネル

カーネル	機能
`kFwdAttn`	RMSNorm + QKV プロジェクション + SDPA + 出力プロジェクション
`kFwdFFN`	RMSNorm + SwiGLU FFN (W1, W3, SiLU, W2)
`kFFNBwd`	FFN バックワード (W2ᵀ + SiLU_bwd + W1ᵀ + W3ᵀ)
`kSdpaBwd1`	Woᵀ + SDPA バックワードパート 1 (dV, probs, dp)
`kSdpaBwd2`	SDPA バックワードパート 2 (softmax grad, dQ, dK)
`kQKVb`	QKV バックワード (Wqᵀ + Wkᵀ + Wvᵀ → dx)

CPU が担当: RMSNorm バックワード、残差接続、損失計算、dW 勾配累積（cblas_sgemm）、Adam オプティマイザの更新。

📊 ベンチマーク結果 (M4 Mac、単一 Transformer レイヤー、dim=768, seq=512)

最適化	ms/step	ANE 活用率
ベースライン (vDSP トランスポーズ)	33.5	3.1%
チャンネルファーストレイアウト	20.3	5.2%
vDSP ベクトル化 RMSNorm	14.2	7.4%
GCD 非同期 cblas オーバーラップ	11.4	9.2%
ANE RMSNorm フュージョン	11.4	9.2%
Wo^T フュージョン (7→6 カーネル)	11.4	9.2%
遅延 cblas 待機	9.3	11.2%

ベストリザルト: 9.3 ms/step、1.78 TFLOPS を維持（M4 ANE の 15.8 TFLOPS ピークの 11.2%）。

🔑 主要な最適化

チャンネルファースト CPU レイアウト — ANE のネイティブ [1,C,1,S] IOSurface フォーマットと一致し、すべてのトランスポーズオーバーヘッドを排除します。
vDSP ベクトル化 RMSNorm —Naive な実装よりも 10 倍高速（6.7ms → 0.7ms）。
GCD 非同期 cblas オーバーラップ — 重み勾配 sgemm 操作は、シリアルディスパッチキュー上で ANE 評価と並列に実行されます。
遅延 cblas 待機 — dW sgemm の待機は、最大の計算オーバーラップのために次のステップのフォワードパスに押し込まれます。
ANE RMSNorm フュージョン — RMSNorm は MIL オペレーションとしてフォワード ANE カーネルに直接折り込まれます。
フォワードタップ — Q、K、V、アテンションスコア、隠れ状態は、バックワードパス中の CPU による再計算を回避するために、連結出力経由で公開されます。
exec() 再起動 — チェックポイント/再開を使用して、プロセスあたり約 119 回の ANE コンパイル制限を回避します。

📁 リポジトリ構造

├── api_exploration.m       # 初期 ANE API 発見実験
├── inmem_basic.m           # インメモリ MIL コンパイル概念実証
├── inmem_bench.m           # ANE ディスパッチレイテンシベンチマーク
├── inmem_peak.m            # ピーク TFLOPS 測定 (2048×2048 行列乗算)
├── sram_bench.m            # ANE SRAM 帯域幅プローブ
├── sram_probe.m            # SRAM サイズ/レイアウト探索
└── training/
    ├── ane_runtime.h       # ANE プライベート API ラッパー (コンパイル、評価、IOSurface)
    ├── ane_mil_gen.h       # MIL プログラム生成ヘルパー
    ├── model.h             # モデル重み初期化とブロブビルダー
    ├── forward.h           # フォワードパス MIL ジェネレータ
    ├── backward.h          # バックワードパス MIL ジェネレータ
    ├── train.m             # ミニマルトレーニングループ (初期プロトタイプ)
    ├── tiny_train.m        # 2層の小さなモデルトレーニング
    ├── train_large.m       # メイン: 単一レイヤー dim=768 トレーニング (最適化済み)
    ├── test_*.m            # 個々のカーネルの単体テスト
    └── Makefile

🛠️ ビルドと実行

要件: Apple Silicon 上の macOS 15+ (M4 でテスト済み)。

# ビルド
xcrun clang -O2 -framework Foundation -framework IOSurface \
  -framework CoreML -framework Accelerate -ldl -lobjc \
  -o train_large training/train_large.m

# 実行
./train_large

外部依存関係なし — システムフレームワークと実行時に objc_msgSend を介して解決されるプライベート ANE API のみを使用します。

⚠️ 既知の制限事項

制限事項	詳細
低い ANE 活用率	ピークの約 11.2%; 多くの要素ごとの演算は依然として CPU にフォールバックします
約 119 回のコンパイル制限	ANE コンパイラはリソースをリークします。`exec()` 再起動で回避しました
単一 Transformer レイヤー	マルチレイヤーパイプラインスケジューリングはまだ実装されていません
合成データのみ	実トークン化データサポートは作業中です
SDPA 因果マスキング	ANE は SDPA の `attn_mask` を無視します。手動分解で回避
プライベート API	macOS のアップデートでいつでも破損する可能性のある、文書化されていない API を使用します

📚 関連研究論文

🏷️ テックスタック

言語: Objective-C, C
プラットフォーム: macOS 15+ / Apple Silicon (M1/M2/M4)
フレームワーク: Foundation, IOSurface, CoreML, Accelerate
プライベート API: _ANEClient, _ANECompiler, _ANEInMemoryModelDescriptor
フォーマット: MIL (Model Intermediate Language), fp16

⚖️ 法的免責事項

このプロジェクトは、Apple のプライベートで文書化されていない API を使用しています。これらの API は安定性の保証がなく、macOS のアップデートでいつでも破損する可能性があります。このプロジェクトは、公正使用および相互運用性の規定（Sega v. Accolade, 1992; DMCA §1201(f)）に基づく独立した研究です。Apple の専有的なコードやバイナリは含まれていません。Apple Inc. とは提携しておらず、承認されていません。