オリオン：CoreMLをバイパスしてLLMをApple Neural Engine上で直接実行およびトレーニングする方法

LLMワークロードのための直接ANEアクセス

Orionは、CoreMLを完全にバイパスし、LLMをApple Neural Engine（ANE）上で直接実行およびトレーニングするエンドツーエンドのObjective-Cシステムを提供します。このアプローチにより、開発者はANEを直接制御できるようになります。これまでCoreMLによってブラックボックススケジューラとして扱われ、直接制御やトレーニングの能力が剥奪されていたANEに対して、直接的なコントロールが可能になります。

技術的実装と制約

このプロジェクトは、非公開のANEClientおよびANECompiler APIをマッピングしたリバースエンジニアリングの成果に基づいています。ANEは、開発者が「ハードウェアインピーダンスミスマッチ」と呼ぶ17のプログラミング制約を提示しており、そのうち11は完全に文書化されていませんでした。主な制約には以下が含まれます：

• concat操作は即座に、サイレントなコンパイラエラーを引き起こす
• BLOBFILEウェイトはチャンクヘッダーから64バイトのオフセットを必要とし、そうでないとサイレントな数値破損が発生する
• ANEは内部状態を維持しており、プロセスあたり約119回のコンパイルでハードキャップされ、その後サイレントに失敗する

トレーニング課題への解決策

以前のANEトレーニングの試みでは、1ステップ後にNaN発散が発生していました。Orionは以下によってこの問題を解決します：

• 遅延コンパイルパイプラインを構築する
• fp16オーバーフロー連鎖を防ぐための厳密なアクティベーションクランプを実装（アクティベーションを-65504から+65504にクランプ）
• 119回コンパイル制限を回避するために、各トレーニングステップ後にexec()プロセス再起動ループを使用する

パフォーマンス結果

コンパイラは、27オペレーションのグラフIRを5つの最適化パスを通じてANEネイティブのMILに変換します。現在のパフォーマンスには以下が含まれます：

• GPT-2 124Mデコードで170+トークン/秒
• 110Mパラメータのトランスフォーマーでの機械的に安定した多段階トレーニング（ハードウェアの「コヒーレンス上限」）
• 1,000ステップ以上にわたり、損失が12.3から6.2に減少し、NaNはゼロ

現在の制限

ANEはコンパイル時にウェイトを固定するため、トレーニングの更新ごとに約4.2秒の再コンパイルペナルティが発生します。ANEはfp16で約19 TFLOPSを発揮しますが、その使用における根本的な制約は計算能力ではなく、ネイティブなオーケストレーションレイヤーの完全な欠如でした。