Metalによる機械学習の加速

こんにちは Kamal Ramamoorthyです GPU, Graphics and Displayチームの ソフトウェアエンジニアです このセッションでは同僚のSamと共に Metalを使用して機械学習モデルを 高速化する方法を紹介します Appleのプラットフォームにモデルを 導入する最初のステップはトレーニングです 次のステップは デバイスへの 導入に向けてモデルを準備することです 最後に 準備したモデルを アプリに組み込みます 今回のセッションでは このステップに 重点を置いて説明します Core MLを使用してモデルを導入する場合 MPSGraphでMetalを使用して GPUを高速化できます Core MLについて詳しくは 「Explore machine learning on Apple platforms」セッションをご覧ください 併せてお勧めするセッションが 「Deploy machine learning models on-device with Core ML」です モデルのトレーニングには PyTorch TensorFlow JAXなどの フレームワークも使用できます 詳しくは「Train your ML models with Apple silicon GPUs」セッションをご覧ください これらのフレームワークはすべて Metalを使用して汎用の計算グラフを 構築および実行するためのフレームワーク Metal Performance Shaders Graph上に 構築されています MPSGraphではGPU同期およびメモリの 低レベル制御を行えるため 場合によってはMPSGraphを直接 使用した方がよいこともあります 例えば アプリで Metalを使用している場合 MPSGraphを使用して 機械学習タスクと 他のGPUタスクを順番に処理できます またバッファなどの低レベルの Metalリソースを共有することもできます Metalを使用して機械学習を 高速化するのが初めてという場合は 昨年までのWWDCのセッションをご覧ください このセッションではSamと私が 3つのことについて説明します 1つ目はMPSとMPSGraphの改善です その多くはTransformerモデルの効率の 向上に重点を置いているため それらのモデルを例として使用します 2つ目は FFTベースの機械学習モデルを 高速化する新機能です 最後はSamが 機械学習モデルを可視化できる MPSGraphビューアを紹介します では Transformerモデルの性能の向上に 重点を置いた新機能から見ていきましょう 一般にTransformerは テキストを翻訳 予測 生成するために 言語モデル内で使用されます 入力値は一連のトークンです 例えば「The quick brown」という シンプルな文は3つのトークンの入力です 言語モデルは次のトークンを 「fox」と予測することで応答します この新しい文が繰り返しフィードバックされ モデルから新しいトークンが生成されます

MPSとMPSGraphにはTransformerモデルを 改善するための新機能が追加されました 改善点は次の3つのカテゴリに分けられます 1つ目は 計算性能の向上です 次は メモリ帯域幅の節約です そして最後は Transformerモデルの品質改善です まずは計算性能から見ていきます

Transformerベースのモデルは複数の層の Transformerブロックで構成されます Transformerブロックの内部構造に注目して さらに詳しく見てみると マルチヘッドアテンション 正規化 フィードフォワードの各ブロックがあります マルチヘッドアテンションブロックは 最も計算負荷の高いブロックの1つです このブロックで行われる計算は 大規模な多次元行列乗算であり 計算負荷の高い演算です 入力行列が行列乗算層を介して射影され Qと呼ばれるクエリ行列の他 キー行列Kとバリュー行列Vが生成されます そしてスケール化ドット積 アテンションブロックに供給されます これがTransformerモデルの心臓部です

このアテンションブロックの性能を 最適化する方法は2つあります アテンションブロックの内部を見てみると いくつかの演算で構成されています

ここからMPSGraphに処理が移され この一連の演算が 1つのカーネルにまとめられます これで より効率的に実行できます

この演算を使用するには MPSGraphオブジェクト でscaledDotProductAttention メソッドを呼び出します このメソッドは引数として query key valueの各テンソルを取ります

結合したSDPA演算を使用することで Transformerモデルの性能を 高めることができます マルチヘッドアテンションブロックを 再確認します これらのテンソルを使用して 計算性能を向上させる方法を もう1つ紹介します ここではquery key valueの各射影が 最初のトークンに対して どう機能するかを確認できます 行列乗算演算によって query key valueの 埋め込みベクトルが射影されます

次の出力トークンを生成するには 前に生成したすべてのトークンを 行列乗算に 供給する必要があります そのため前の反復処理で 既に計算されている KとVの射影が再計算されます

このコストは配列の 長さに応じて加算されます この問題を軽減するには これらの射影を 生成された時点でキャッシュし 後の反復処理で再利用できるようにします 最初にすべきことは キャッシュしたK値とV値を格納する KテンソルとVテンソルの作成です これらはKVキャッシュと呼ばれます 最初の反復処理では 最初の トークンに対してK値とV値が計算され KVキャッシュに挿入されます

これでキャッシュした値を再利用できるので 2回目の反復処理では 2つ目のトークンのK値とV値を 計算するだけで済みます これで行列と行列の乗算を 行列とベクトルの乗算にまで 簡略化できます

反復処理ごとに新しいテンソルを作成して KVキャッシュの末尾にK/V射影を 加えられますが その場合は大量のメモリを消費します 代わりにスライス更新操作を使用すれば 既存のテンソルを インプレースで更新できます

その後は スライス操作を使用して KVキャッシュから計算された 部分だけを抽出できます

これをコードで行う方法を確認しましょう

まずキャッシュを表す プレースホルダを作成します このテンソルの次元は モデルの詳細によって異なります この例ではキャッシュの key部分だけに絞りますが value部分も同じように機能します KVキャッシュをインプレースで 更新できるようにするには 現在のキャッシュの状態を表す 変数を作成する必要があります 通常のグラフ演算の結果とは異なり この変数は後で別の値を 参照するように更新できます

すべてのトークンでキー射影を キャッシュに挿入する必要があります これにはMPSGraphオブジェクトの sliceUpdateDataTensorメソッドを 使用します starts配列とends配列は 新しい値を挿入する場所を示しています この例ではキャッシュの有効部分の 末尾に追加しています この場合 ストライドは一定です

これで更新したキャッシュを 元の変数に代入し直すことができます MPSGraphがこれを最適化して キャッシュの 割り当てをインプレースで更新します 最後に スライス操作を使用することで KVキャッシュの計算された部分だけを 抽出できます これはキャッシュの先頭から 最後に挿入されたkey射影までの部分です

その後は 更新したkeyキャッシュを SDPA演算に渡すことができます

このように計算性能を向上させた後は メモリ帯域幅が新たな ボトルネックになります

大規模言語モデルの重みを 格納するために必要なメモリは 数十ギガバイトほどになる場合もあります 通常 これらの重みは 16ビット浮動小数点を使用して表されます ただし MPSはこれらの重みの 8ビット整数への量子化に対応しており メモリ使用量を半分に減らすことができます

MPSは今年から新たに 4ビット整数形式にも対応しています これによって重みのサイズを さらに減らすことができます MPSはこれらの量子化された形式に 重みをマッピングする 一部の手法に対応済みです

こちらのテンソル例では 各要素が数直線上に分布しています 8ビットの量子化の場合 可能な量子化点は256個あり それらが数直線に沿って 直線的に分布します 4ビットの量子化の場合 量子化点は16個あります 量子化の際 これらの点は最も近い量子化値に 合わせて調整されます 量子化の倍率は 右側の計算式で求めることができます 若干の誤差は生じますが 最終的に2倍または4倍のメモリ空間と 帯域幅を節約できます MPSGraphオブジェクトの dequantizeメソッドを使用すると 値を逆量子化できます

他にもルックアップテーブルを 使用する量子化手法があります この手法は 重みが数直線上の 様々な領域に集まっている場合に便利です アフィン量子化では 量子化された値が均一に分布しますが 入力値は均一になりません ほとんどの入力値がいくつかの 量子化点にのみ集まるため 量子化ビットの大部分が 使用されなくなります ルックアップテーブルを使って 量子化ビットをより効果的に活用できます この手法では データの分布に基づき 独自の量子化点を自分で選択します 量子化したそれらの値を ルックアップテーブルに格納します その後は 各重みの4ビットまたは8ビットの インデックスをこのテーブルに割り当てます この方法であれば 柔軟性を大幅に高めながら テーブル内の値の検索性能は ほとんど損なわれません

dequantizeメソッドを使い これらの量子化した値を 32ビット浮動小数点値に戻します 単に32ビットのルックアップテーブルで 量子化した重みを渡すだけです その後は 逆量子化したテンソルを 通常通りに再利用できます 例えば 行列乗算への 入力として使用できます 実際このようなケースでは MPSGraphはさらに一歩先まで踏み込みます

重みに対する逆量子化演算が グラフに含まれていて その後に行列乗算を実行している場合

MPSGraphはそれら2つの演算を 単一の量子化された 行列乗算演算に置き換えます この演算では 量子化した重みの一時コピーを 格納するのではなく 必要に応じて 重みをその場で 逆量子化します

量子化はメモリとメモリ帯域幅を 節約できますが 数値が不正確になる可能性もあります 次は Transformerモデルの品質を向上させる 2つの方法を見ていきましょう

重みを量子化すると それぞれの重みがより低い精度の値に マッピングされます またスケールも選択する他 必要に応じて 逆量子化時に量子化した値に 適用されるオフセットも選択します ただし単一のスケールとオフセット値を すべての重みに適用すると 再構築される値の精度が 制限されることになります 代わりに 複数の要素からなるブロックを 独自のスケールとオフセット値で 個別に量子化できます そうすることでスケールと オフセット値をブロックごとに より正確に一致させることができます

これを実行するコードは 先ほどの例と同様ですが 単一のスケールとゼロポイント値を 渡す代わりに ブロックごとのスケールとゼロポイントが 含まれたテンソルを渡します

以上で量子化の説明は終わりです 次に紹介する別の方法では アダプタを使用してTransformerモデルの 品質を向上させることができます

アダプタとはモデルに挿入できる 小さなレイヤであり ごくわずかな演算とパラメータで 構成されています モデルを微調整した時は アダプタ内のパラメータのみが更新されます トレーニング済みの基本モデルを 新しいタスクに適合させることができます 量子化によって生じた誤差を 補正するためにも使用できます MPSGraphの呼び出し可能オブジェクトを 使用して アダプタをモデルに追加できます

各アダプタは メイングラフからの呼び出しが可能な 個別のMPSGraphとして機能します まず アダプタごとに一意の名前を指定して 基本グラフからのアダプタの 呼び出しを挿入します これをコード化する際には アダプタの呼び出しによって 生成される出力の形状と型を 定義する必要があります 次に メインMPSGraphオブジェクトの callメソッドを使用して アダプタの呼び出しを追加します その際 呼び出し可能オブジェクトの 名前 入力 出力の型を指定します

次は 各アダプタのMPSGraphを作成します この例では ランク付けされていないテンソルとして 「input」のプレースホルダを作成します 次に「input」テンソルを2倍することで 「output」テンソルを作成します

最後に各アダプタのグラフをコンパイルして グラフの実行可能ファイルを生成します これらは他のMPSGraphと同じように コンパイルされます まずは正確な形状を指定して グラフへの入力の型を定義します 次に グラフオブジェクトの compileメソッドを呼び出します その際 Metalデバイス 入力の型 outputTensorを指定します

メイングラフからのアダプタの 呼び出しを追加し 各アダプタのグラフ実行可能ファイルを コンパイルしました 最後に必要なのは メイングラフの実際の グラフ実行可能ファイルに アダプタ名をマッピングすることです これは ネットワークの メインMPSGraphをコンパイルする時に GraphCompilationDescriptorを 使用して行います まず 各アダプタの名前を グラフ実行可能ファイルに マッピングする辞書を作成し このディスクリプタに設定します 後は メイングラフのコンパイル時に CompilationDescriptorを指定するだけです アダプタの設定に必要な作業は これですべてです まとめると アダプタと 呼び出し可能オブジェクトを使用することで 新しいタスクを実行するよう モデルをカスタマイズできる他 モデルの出力の微調整もできるため 品質が向上します 次は MPSとMPSGraphに今年追加された フーリエ変換の新機能を紹介します 高速フーリエ変換(FFT)は 信号や画像などのデータを 時間領域または空間領域から 周波数領域に変換します 音声からテキストへの変換モデルや 単一のトラックから 様々なオーディオソースを 分離するモデルなど オーディオを処理する機械学習モデルでは 一般的な前処理ステップです また 特定の畳み込みレイヤを 高速化するためにも使用でき 多くの画像処理アプリケーションや 科学計算アプリケーションで利用されています 例えば オーディオ信号から テキストを抽出するには 最初に入力波形で 短時間フーリエ変換(STFT)を実行します 次に 周波数スペクトルを Transformerモデルで分析して テキストを抽出します GPUで機械学習モデルを 効率的に実行するために MPSGraphをどのように使用すればよいかは 既に説明しました しかし MPSGraphの高速フーリエ変換の サポートを利用すれば このパイプライン 全体をGPUに移行することもできます 最初のステップは 短時間フーリエ変換の実装です この手法では 入力波形を短い複数の ビュー（または窓）に分割します これは互いに重なり合っている 場合もあります 各窓は実質的に独立した信号です スペクトル漏れを軽減するため 個々の窓を窓関数で乗算します 最後に 通常のバッチ処理 1次元FFT演算を使用することで 各窓のSTFTを計算できます 波形をより短い窓に分割するには ストライドビューを作成します まず 窓に分割するビューの 形状を設定します この場合の窓幅は512要素になります 次に 各次元のストライドを設定します この例では256を使用します つまり 基盤となる 1次元配列の256要素を 2次元目の各ステップでスキップします 最後のバッチ次元は1に設定されていますが より大きなバッチサイズも使用できます

最後にinputテンソルのarrayViewメソッドを 呼び出すことで ストライドビューを 作成できます 何より素晴らしいのは 入力配列のメモリのエイリアシングにより コピーを行わずにビューを演算して メモリとGPU時間を節約できることです

これですべての窓に対して FFTを計算できます まず ストライドビューのデータ用に プレースホルダを作成します ストライドビューのNDArrayから データを読み込んで 後でグラフの実行時に 指定する必要があります 次に窓関数で乗算を行います 通常はハン窓またはガウス窓で行います 例えば これにMPSGraphの 定数テンソルを使用できます 最後にFFTTensor演算を作成します 高速フーリエ変換については以上です ここからはSamに引き継ぎます MPSGraphの構造について 理解を深めたい方に向けて 朗報があります ありがとうKamal こんにちは Sam Colbranです GPU, Graphics and Displayチームの ソフトウェアエンジニアです ご存知ない方もいるかもしれませんが Apple GPUを最大限に活用できるように XcodeとInstrumentsにはMetalの 高度なツールが組み込まれています パワフルな機能を簡単に利用できるため デバイス上で実行する MetalパイプラインやAIモデルは より大規模で複雑なものになりがちです MetalパイプラインはXcodeの 依存関係ビューアで可視化できますが MPSGraphを可視化することは これまで不可能でした 本日はXcode 16でMetalツールに 追加された最新機能 MPSGraphビューアを紹介します これは機械学習とAI向けに特別に 設計されたまったく新しいツールです 現在はXcodeで MPSGraphパッケージを直接開き 各演算の接続状況を可視化できます デモを始める前に まずはMPSGraphパッケージを 実際に作成する方法を振り返りましょう MPSGraphを直接使用する場合と 別のフレームワークでMLモデルを 開発している場合の両方についてです MPSGraphで直接 モデルを作成している場合は まず グラフをMPSGraphの 実行可能ファイルにコンパイルします 次に 実行可能ファイルのシリアル化APIを 使用してパッケージを作成します

今年からは MPSGraph実行可能ファイルを CoreMLパッケージから 直接作成することもできます その後は 以前と同様 実行可能ファイルを MPSGraphパッケージにシリアル化できます

またONNXにエクスポートする フレームワークなど 別のフレームワークで開発する際は mpsgraphtoolを使用して モデルを変換できます では 例を見てみましょう

ここでは70億個のパラメータを含む Mistralのモデルを使用しています 今年の「Bring your machine learning and AI models to Apple silicon」セッションで CoreMLに変換されたモデルと同じものです mpsgraphtoolには コマンドラインからアクセスできます お好みのターミナルを開き 引数にconvertを使用して mpsgraphtoolを実行してください

これで完了です 新たに作成したMPSGraphパッケージは いつでも使える状態です また 現在は新しいMPSGraphビューアで 簡単にパッケージを確認できます Xcode 16で変換した Mistralパッケージを開いたところです 画面上の内容を説明していきます まず 左上には コンパイルオプションがあります デフォルトでは ビューアに グラフがそのまま表示されます つまり 特定のデバイスに 最適化されていません 使用しているデバイスにかかわらず 演算は同じように表示されます

その下には演算ナビゲータがあります ここには グラフで使用されている すべての演算がリスト表示されます 中央にあるのはグラフそのものです 最後に 右側には演算インスペクタがあります 後でまた説明します

このレベルでは少し見えにくいので 拡大しましょう

大まかな構造が見えるようになりました さらに拡大すると より細部まで見ることができます この状態でスクロールすれば 各演算の入力と出力に加え それらの接続状況をすべて確認できます これによってグラフの構造の可視化と 理解が容易になります

MistralはTransformerモデルです Kamalが先ほど説明したように Transformerモデルは複数層の Transformerブロックで構成されます それを探してみましょう まず 新しいScaled Dot Product Around 演算子を探します 各Transformerブロックの マルチヘッドアテンションの 一部になっているはずです 検索することもできますが 左側の演算ナビゲータを見ると この演算子が 32個あることがわかります このグループを展開して 最初の項目をクリックすると グラフの該当箇所にジャンプします

この演算には5つの入力があるようです 先ほどKamalが説明した Query Key Valueがあることに 気付いたかと思います

少し縮小して Transformerブロック全体が 見えるようにしましょう

接続をたどってみると Query Key Valueを 構成しているブロックを 確認できます このレベルでも KeyとValueの両方の変数を確認できます Keyのブロックを拡大します このモデルは状態とともに CoreMLからエクスポートされたため KVキャッシュを使用してMPSの 新しい「変数に代入」演算や 「変数からの読み取り」演算を活用しています Kamalが説明したように これによって計算性能が向上します グラフを簡略化するため ビューアの複数の場所で 一部の演算や変数がこのように 表示されることがあります この変数を選択してみると 右側のインスペクタで 最初に作成された場所と使用された すべての場所を確認できます

Transformerブロックを1つ確認しました 残りも確認します インスペクタと同じように 演算ナビゲータでも 複数の演算を同時に選択できます このように大まかな構造が わかりやすくなり 何度も繰り返されている レイヤをすべて確認できます とても便利ですね 次は定数について説明しましょう お気付きのように グラフ内に緑色のプレビューが 直接表示されています 左側のナビゲータタブでは これらをサイズ順に表示することもできます 一番大きいものを選択して ダブルクリックすると 定数ビューアが開きます ここではトレーニング済みの重みを調べ モデルが学習した内容についての インサイトを得ることができます デバイスへのスムーズな統合に向けて モデルを最適化する機会を 見つけるのに役立ちます ただし このビューアには グラフがそのまま表示されます 特定のデバイスに 最適化されるわけではありません

実際のグラフは 異なるものになる可能性があります 例えば MPSGraphは 複数の演算をスティッチングして 1つのMetalカーネルにまとめ 自動で演算を最適化する場合があります ビューアを使用すると これを可視化できます 具体的に説明します ResNet50が含まれている MPSGraphパッケージを開きます 先ほどと同様に拡大すると すべての演算と定数を確認できます 今回は 私のデバイスでグラフが どのような見た目になるか確認しましょう

左上のコンパイルオプションで 自分のデバイスを選択します

拡大してみます これらのMetal Shader (Stitched)領域内で 演算がグループ化されていることがわかります この領域を拡大して内部を確認できます これらの演算は 最適化された1つの Metalシェーダ内で結合されているので 内部的にメモリのオーバーヘッドが発生せず 性能が劇的に向上します 一般に グラフがハードウェア上で 最終的にどのように実行されるかわかれば 実行時のパフォーマンス特性を より深く理解できます 新しいMPSGraphビューアについては以上です では 本日の要点をまとめましょう Kamalが説明したように Metal Performance Shaders Graphを 使用することにより Metalで機械学習を高速化できます Appleシリコンで 最適なパフォーマンスを実現するため CoreMLなどの一般的なフレームワークでは 既に利用されています 今年は 計算性能の向上に役立つ Transformerの新機能に加え 新しいScaled Dot Product Attentionと KVキャッシュが追加されました 量子化によるメモリ帯域幅の節約や 呼び出し可能オブジェクトを通じた アダプタによる品質向上も可能です MPSGraphでは 新しいストライドNDArray APIによって フーリエ変換の計算をさらに高速化できます 最後に 新しいMPSGraphビューアでは Appleシリコンで モデルがどのように実行されるかを 簡単に把握してインサイトを獲得できます MPSGraphのドキュメントと サンプルコードを忘れずにご確認ください そして当然ながらパズルの 最後のピースはモデルの統合です これらの素晴らしいセッションについても 未視聴の場合はぜひご覧になり トレーニングと導入についての 詳細をご確認ください ご視聴ありがとうございました WWDCをお楽しみください