論文名

Transformer-Lite: High-efficiency Deployment of Large Language Models on Mobile Phone GPUs

arXivリンク

https://arxiv.org/pdf/2403.20041.pdf

ひとこと要約

モバイル端末のGPUで大規模言語モデル(LLM)を効率的に実行するTransformer-Liteを提案。

メモ

手法

• シンボリック式に基づくアプローチにより、動的な入力形状のモデル推論をサポート

  • LLMは入力トークン数が動的に変化するため(テンソルの形状も動的に変化するため)テンソルの不明な次元をシンボリック式で表現し、メモリ再利用やスケジューリングを行う。（入力トークン数をNとすると、テンソルの形状を以下のようにシンボリック式で表現できる。

    • [N, 128]

    • [N * 2, 256]

    • [batch_size, N, 1024])

• オペレーター最適化と実行優先度設定により、推論速度を高速化し、端末のラグを低減

  • 半精度活性化関数の計算と4bit量子化重みを直接扱う行列積を実装し、メモリ帯域幅を最大限活用（活性化関数の計算は精度を維持するために半精度が必要）

  • レイヤーノーム等の小さなオペレーターを融合

  • GPUの実行優先度を最低に設定してラグを低減

• FP4量子化手法M0E4

  • デクオンタイゼーションの整数から浮動小数点数への変換のオーバーヘッドを指数部0bit、仮数部4bitとして回避。(ビット演算のみで浮動小数点数に変換できるため)

• KVキャッシュのコピー最適化

  • KVキャッシュを1つの大きなテンソルに格納し、モデルの入力に用いることで、KVキャッシュの余分なコピーを回避（LLMの推論では、各トークンを順番に処理していく際に、直前のトークンから生成されたKey-Value（KV）キャッシュが次のトークンの推論時の入力として使用される。）

結果

モバイル推論エンジンTransformer-Liteの性能を、5つのLLMモデル（Gemma 2B、Qwen1.5 4B、ChatGLM2 6B、Llama2 7B、Qwen1.5 14B）で評価。比較対象としてGPUベースのMLC-LLMとCPUベースのFastLLMを用いて、2つのモバイルプロセッサ（Qualcomm Snapdragon 8 Gen 3、MediaTek Dimensity 9300）上で評価。

• Prefill（プロンプト入力）速度

  • プロンプト長によらず、最適な性能を発揮（1024トークン程度が最適）

  • パラメータ数が小さいモデルほど高速（Gemma 2Bで330token/s、ChatGLM2 6Bで121token/s）

  • Snapdragon 8 Gen 3の方がDimensity 9300よりも高速

• デコード（生成）速度

  • プロンプト長が長くなるほど低下

  • KVキャッシュの形式変更が必要なLlama2 7B、Qwen1.5 4Bは、ChatGLM2 6B、Gemma 2Bよりも速度低下が大きい

  • Gemma 2Bで30token/s、ChatGLM2 6Bで14token/s

• 他手法との比較

  • MLC-LLM（GPU）よりもPrefill速度で5倍以上、デコード速度で約2倍高速

  • FastLLM（CPU）よりもPrefill速度で10倍以上、デコード速度で2～3倍高速

• 最大デプロイ可能モデルサイズ

  • 12GBのメモリで最大7Bまで（900トークンまで）

  • 24GBのメモリでQwen1.5 14Bを2048トークンまで実行可能

• 理論限界との差

  • Prefill速度は理論ピーク（5.6TFLOPS）の約1/3（1.6TFLOPS）

  • デコード速度は短いプロンプト長ではメモリ帯域の限界