SiLUとSoftMaxを2倍高速化する新しい指数関数、精度は完全維持

 (github.com/ggerganov)
1 ポイント 投稿者 GN⁺ 2024-05-16 | 1件のコメント | WhatsAppで共有
  • llama.cpp PR #7154 は、GGML の CPU 向け SiLU と SoftMax 計算を、llamafile のベクトル化された expf() ベース実装で書き直すもので、2024年5月17日に master へマージされた
  • 従来の GGML は速度のために short[65536] の ルックアップテーブル を使用していたが、新実装は aarch64 と SSE2+ で最悪の丸め誤差を 2 ULP に保ちながら、より高精度な計算を目指している
  • SOFT_MAX の CPU 性能テストでは、SSE2+FMA は 1.5倍、AVX2+FMA は 1.9倍、AVX512 は 2.1倍高速化し、AMD Ryzen 9 5950X と M2 Ultra でも master 比で約 1.5倍高速という結果が確認された
  • 変更には ggml_v_expf() と ggml_v_silu() の追加、重複コードの ggml_vec_soft_max_f32() への抽出、GGML_SILU_FP16 関連関数の削除、SSE2 または ARM NEON 条件付き SiLU パスの調整が含まれる
  • マージ後、>1 slots でのサーバー実行時に 非決定的な結果 が再現され、その後 -ffinite-math-only が原因と特定され、-fno-finite-math-only が必要になるビルド上の制約へとつながった

PR の変更目的とマージ状況

  • PR #7154 は ggml : rewrite silu and softmax for cpu というタイトルで、llama.cpp の GGML CPU パスにおける SiLU と SoftMax の計算を書き直すもの
  • 変更は、llamafile の ベクトル化された expf() 関数 を upstream する形で始まった
  • PR は 2024年5月17日に ggml-org:master へマージされ、マージコミットは 934266c と表示されている
  • 作者は、従来の GGML が速度のために使っていた short[65536] ルックアップテーブルよりも、新方式の方が SoftMax と SiLU をより正確に計算できると述べている

精度と対応範囲

  • 新しい expf() ベースのパスは aarch64 と SSE2+ をサポートし、最悪の丸め誤差は 2 ULP とされている
  • 初期説明では AVX2 と AVX512 の実装も作成されたが、SSE2+FMA に比べてコード複雑性を受け入れるほどの利点が大きくないとして含めなかったという
  • その後、ベンチマーク結果に基づいて AVX2 と AVX512 のコード も含まれた
  • 別個のテスト出力では 4294967296 numbers tested successfully が示され、複数の入力値に対する exp と llamafile 実装の結果比較も含まれている

コード変更範囲

  • レビュアーが整理した主な変更点は以下の通り
    • コメントアウトされた #define の削除
    • 重複した 5 行を ggml_vec_soft_max_f32() に抽出
    • GGML_SILU_FP16 関連の複数関数を削除
    • ggml_v_expf() を追加
    • ggml_v_silu() を追加
    • ggml_vec_silu_f32() が SSE2 または __ARM_NEON フラグに応じて異なる関数を使うよう、プリプロセッサ条件を調整
  • 変更ファイル数は GitHub のメタデータ上では 1件 と表示されている
  • PR には refactoring と Review Complexity : High のラベルが付いており、後者には LLM または GPU に関する深い知識が必要な場合があるとの説明が含まれる

ベンチマークと性能結果

  • ggerganov は AMD Ryzen 9 5950X と M2 Ultra で、SOFT_MAX が master より約 1.5倍高速 であることを確認した
  • 使用したテストコマンドは次の通り
make -j tests && ./tests/test-backend-ops -o SOFT_MAX -b CPU perf
  • その後、作者は同じコマンドでの性能向上幅が次のように拡大すると述べている
    • SSE2+FMA: 1.5倍
    • AVX2+FMA: 1.9倍
    • AVX512: 2.1倍
  • 別の開発用スクリプトでは以下の数値が示されている
    • run_expf(): 2.98601 ns
    • run_llamafile_expf_sse2(): 1.35154 ns
    • run_llamafile_expf_avx2(): 1.16659 ns
    • run_llamafile_expf_avx512(): 1.18844 ns
  • GitHub Actions の llama.cpp server ベンチマークは、Standard_NC4as_T4_v3 上の phi-2 q4_0 構成で 543 iterations を記録した
    • 同時ユーザー数: 8
    • duration: 10分
    • HTTP リクエスト平均: 8626.19ms
    • p95: 21696.44ms
    • Prompt processing 平均: 94.59 tk/s
    • Token generation 平均: 33.43 tk/s

AVX512 最適化の議論

  • chriselrod は AVX512 で vscalefps を使うことを提案した
    • vscalefps は zmm0 = zmm1 * 2^{zmm2} を計算する
    • overflow と underflow を適切に処理でき、checks や blends を除去できるという
  • Julia の実装例とアセンブリループが共有され、テストが正しければ x=47.483456f で最大誤差は 1 ULP 未満 だったとされる
  • vscalefps アプローチは lookup table を使わず、Float64 / double 実装では vpermi2pd を通じた 16 要素の lookup table を使うと説明されている
  • その後、C++ 実装へのリンクも共有された
    • ExpAVX512
    • ソースは include/ExpAVX512.hpp にある
    • README にはベンチマークが含まれるが、他実装との比較ベンチマークは行っていないと述べられている

マージ後の非決定性問題

  • マージ後、サーバーで >1 slots を使うと 非決定的な結果 が出る再現例が報告された
  • 最小再現手順は次の通り
make clean && make server
./server -m models/opt/llama_2-7b-q4_0.gguf --parallel 2 --threads 1
  • 別のシェルで実行したリクエストは以下の通り
curl --request POST --url http://localhost:8080/completion --header "Content-Type: application/json" --data '{"prompt": "", "n_predict":10, "n_probs": 2, "temperature": -1}' | python3 -m json.tool
  • 最後のトークンの token probabilities が curl 呼び出しごとに 2つの値の間を循環し、4 slots を使うと 4つの可能な値の間を循環するという

-ffinite-math-only とビルド制約

  • その後の関連コミットでは、-ffinite-math-only が問題の原因だと絞り込んだ内容が参照されている
  • この問題は、SiLU が小さな値を 0 に flush する代わりに NaN やその他の garbage 値を返すことが原因と推定されたと記録されている
  • 修正では、-fno-finite-math-only が設定されているかを確認し、コンパイルモードが finite math モードではないことを強制的に検査するようにした
  • エラーメッセージでは、GGML の一部ルーチンが non-finite math arithmetic を必要とし、コンパイラに -fno-finite-math-only を渡すよう案内している
  • その後、ユーザーからは -Ofast や -ffast-math が -ffinite-math-only を含むことでビルドを壊しうるという経験が共有された
    • GCC 13.2 までは -Ofast を使えていたが、GCC 14 から結果が garbage になったとの報告がある
    • 一部のテストでは -fno-finite-math-only に加えて -fmath-errno も必要だったという
    • -ffast-math を削除するか -fno-finite-math-only を明示して ggml のコンパイルエラーを解消した後続コミットが複数のリポジトリで参照されている

関連記事

1件のコメント

 
GN⁺ 2024-05-16
Hacker News のコメント

  • 20年ほど前、Hughes のレーダー信号プロセッサ向けにプログラミングしていたとき、0 < x < 1 の範囲の e^x を計算する必要があった。
    そのプロセッサには乗算があったので、32ビットワード内の4つの8ビットブロックそれぞれについて、取り得る256個の値に対応する e^x テーブルを4つ作り、それらを掛け合わせて最終値を得ていた。
    それまで最速だった e^x ルーチンより約5倍速く、今では旧式だが、しばらくの間は名目上ずっと高速なプロセッサよりもレーダー信号を速く処理していた、面白いマシンだった。

    • 追いにくかったなら、発想としてはおおよそ e^x = e^(a+b+c+d) で、ここで a/b/c/d は x の各バイト、これを e^a * e^b * e^c * e^d に変換し、それぞれの e^a、e^b lookup table を作る、という方式に見える。
      厳密には a が high byte << 24 のような形なので、e^a テーブルは a => e^(a<<24) の対応になり、他のバイトも同様に扱うということになる。

  • こうした silu と softmax の改善が LLM 推論速度全体にどれくらい影響するのか気になる。
    間違っていたら直してほしいが、時間の大半は行列乗算に使われるので、この変更の効果は小さいのではないかと思う。

    • 浮動小数点演算の大半が行列乗算に使われるのはその通りだが、softmax はメモリ帯域を不釣り合いなほど多く使うため、演算量だけを見て予想するより通常はずっと時間がかかる。

  • 少し脱線するが、ざっと見て「これはかなり狂った最適化っぽいな。複雑だし、すでに多くの人が見ているコードなのに」と思い、コントリビューターを見て「やっぱり jart か。狂ったように良い解法はいつも jart だな」と思った。

    • 主に恐ろしく見える理由は、C/C++ の intrinsics 構文がもともとそういうものだから。
      あの分野の多くのことと同じく、この苦痛もかなり自業自得な面がある。
      C# 風の SIMD とハードウェア intrinsic 構文を可能にする C++ ライブラリもあると認識しているが、命令セットのドキュメントでニーモニックを直接探しにくいという欠点がある。
      ここで行われた作業の重要性を貶めたいわけではなく、より広い読者にとって近づきやすくできたかもしれない、という意味だ。ただし、まだ推論バックエンドを C# で書き直そうという、ここにいる全員が馬鹿げていると思うであろう提案をするつもりはない。
    • adapted from arm limited optimized routine ということは、結局は巨人の肩の上に立っているわけだ。
    • こういうものは漸近解析の講義で教える内容ではない気がする。
      「誰もが無視するあの定数が、工学的には頭を丸ごと食い尽くすことがある」と有名な言い方をした教授を思い出す。

  • short[65536] の lookup table を置き換えるというが、そもそも少し鈍い選択ではないかと思う。
    L1 キャッシュ全体サイズの lookup table を置くようなものだが、確率的にたまたまうまくはまって、意外とうまく動くのだろうか？

    • lookup table が意外とうまくいく理由は、ワークロード自体が極端にキャッシュ非友好的だから。
      L1 キャッシュを吹き飛ばしてもあまり問題なく、LUT を入れるために追い出されたデータはどうせほとんど再利用されなかった。
      機械学習の負荷は概して、各反復でデータセット全体を線形に読むストリーミングロードである。
    • lookup table をなぜおそらく使うべきでないかについての記事 https://specbranch.com/posts/lookup-tables/ は、一般にいつ適切かを扱っている。
      限られた経験では、lookup より速くなるまでに、リアルタイム計算をかなり多く行える。

  • llama.cpp では、CPU向けの話だ。

    • もともとは llamafile 向けに開発され、最近の2つのリリースに含まれている: https://github.com/Mozilla-Ocho/llamafile/releases/tag/0.8.2
      いまこれを llama.cpp プロジェクトへ upstream しているところ。
      現在 llamafile でしか得られない他の性能改善もあり、たとえば Kawrakow が K quants を大幅に高速化した作業がある。

  • 少し脱線かもしれないが、ggml のようなものが tensorflow lite、onnxruntime のようなランタイムと比べてどうなのか知っている人はいる？

    • ONNX と llama.cpp Flutter ライブラリを6つの True Platforms すべてで保守しているので、かなりよく知っている。
      短く言うと、LLM には llama.cpp が適しており、主要依存である GGML で whisper も可能。
      それ以外には ONNX を使えばよい。
      TF は機械学習界の Apple のようなもので、Google ML エコシステムに完全に縛られているなら素晴らしいが、それ以外では事実上死んでいる。HF モデルの途方もない割合、およそ94%が PyTorch だ。
      直接の推論性能比較をする価値があるのは ONNX の Whisper と GGML くらいだが、誰かが私の llama.cpp ライブラリを Whisper と一緒に動かしており、有意な性能差は報告していない。
    • 正確にどのハードウェアを基準にしているかが重要だ。

  • 現時点で CUDA デバイスの非バッチ推論では、gguf/llama.cpp のほうが性能の良い解法なのか、それとも依然として exllamav2+flashattention が優位なのか？

    • 2x 4090 では差は無視できる程度。
      4ビット KV キャッシュのような、もっと重要な違いがある。

  • LUT もベクトル化できる。
    https://www.intel.com/content/www/us/en/docs/intrinsics-guid...
    以前、LUT で可能なことについて書いたこともある https://darkcephas.blogspot.com/2018/10/validating-utf8-stri...

    • その通りだが、exp を直接実装しても、求める精度によっては FMA 10〜20個程度で済む。
      gather や permutation が純粋な計算と競うのは難しい。

  • 似た文脈で、より高速な tanh もある https://github.com/microsoft/onnxruntime/pull/20612

    • 素晴らしい仕事だ。
      ところで目的は何だろう？ その GeLU 近似をさらに速くすることなのか？
      erff() に戻せば、おそらくずっと速くなると思う。

  • これは gguf の GPU 部分オフロードのユースケースにも役立つのか？
    CPU 側も速くなるのか？