Xeon上で 🤗 Optimum Intel を使用した SetFit 推論の高速化

TL;DR

Hugging Face は、🤗 Optimum Intel ライブラリを介した事後トレーニング静的量子化を利用することで、Intel Xeon CPU 上で SetFit モデルが推論スループットを最大 7.8 倍向上させることができることを実証しました。この最適化により、精度への影響を最小限に抑えた状態で、プロダクションレベルの少数ショット学習モデルのデプロイが可能になります。

SetFit: 効率的な Few-Shot 学習

SetFit は、ラベル付き学習データが限られた状況に対応するために設計された Sentence Transformers の Few-Shot ファインチューニング用フレームワークです。Large Language Model (LLM) ベースの手法に比べて、主に次の 2 つの利点があります。

  • プロンプトの不要化: LLM のインコンテキスト学習では脆弱で手作業のプロンプトが必要ですが、SetFit は少数のラベル付き例から直接埋め込みを生成します。
  • 学習効率: SetFit は、GPT-3.5 や Llama 2 などの LLM に比べて、学習および推論が通常 1 桁速く行えます。

性能面では、SetFit は 3-shot プロンプトにおいて GPT-3.5 を上回り、5-shot を使用した場合には Banking 77 金融インテントデータセットで 3-shot GPT-4 を上回ることが示されています。

🤗 Optimum Intel を使用した推論の高速化

Intel CPU 上で SetFit の推論を高速化するために、Hugging Face は Optimum Intel を利用しています。これは、Intel Advanced Vector Extensions 512 (AVX-512)、Vector Neural Network Instructions (VNNI)、Advanced Matrix Extensions (AMX) を活用するオープンソースライブラリです。これらのハードウェアアクセラレーションは、各コアに組み込みの BFloat16 (bf16) および int8 GEMM アクセラレータを備えています。

事後トレーニング静的量子化 (PTQ)

使用される主要な最適化手法は、Intel Neural Compressor (INC) を介した Post-Training Static Quantization (PTQ) です。量子化は、重みや活性化を表すビット数を削減（例: 高精度数値を INT8 に変換）し、メモリ使用量を減らし、計算を高速化します。

PTQ は、追加の学習を必要とせずにレイテンシとメモリ使用量を削減できる点で特に有効です。未ラベルのキャリブレーションセット（通常約 100 サンプル）だけで、未見データの分布を表現できます。

ベンチマークとパフォーマンス結果

パフォーマンスは、sst2 データセット上の bge-small モデルを使用して評価し、元の PyTorch/Transformers fp32 実装と IPEX-bfloat16 バージョン、量子化された Optimum-int8 バージョンを比較しました。

レイテンシとモデルサイズ

バッチサイズ 1 の場合、最適化された Optimum-int8 モデルは大幅な改善を示しました：

  • レイテンシ削減: レイテンシが 3.45 倍削減されました（fp32 モデルの 15.69 ms に対し 4.55 ms）。
  • モデルサイズ削減: モデルサイズが 2.85 倍縮小しました（127.32 MB に対し 44.65 MB）。
  • 精度の維持: 精度はほぼ変わらず、88.4%（fp32）から 88.1%（int8）へとわずかに変化しました。

スループットの向上

レイテンシの改善は大きいものの、最も顕著な向上はスループットに見られます。さまざまなバッチサイズで達成可能な最高スループットを比較すると、Optimum-int8 モデルは元の transformers fp32 モデルの 7.8 倍の速度です。

指標 bge-small (transformers) bge-small (ipex-bfloat16) bge-small (optimum-int8)
Model Size 127.32 MB 63.74 MB 44.65 MB
Accuracy (test set) 88.4% 88.4% 88.1%
Latency (bs=1) 15.69 +/- 0.57 ms 5.67 +/- 0.66 ms 4.55 +/- 0.25 ms

実装ワークフロー

これらの最適化を実装するには、次の 2 つの主要ステップに従います：

  1. 量子化: optimum.intel の INCQuantizer を使用し、PostTrainingQuantConfig を「static」アプローチかつ「ipex」バックエンドに設定します。
  2. デプロイ: 量子化されたモデル本体を IPEXModel でラップし、SetFit モデル構造に組み込んで推論に使用します。

Sources