見出し画像

日本語LLMの推論速度検証

株式会社レトリバ

はじめに

こんにちは。横浜国立大学大学院 理工学府 修士2年の藤井巧朗と申します。8月24日から9月29日の５週間、株式会社レトリバにインターンとして参加させていただきました。インターンでは日本語LLMの推論速度に関する検証を行いました。本記事では、インターンでの成果の一部を紹介します。

想定する読者：

  • 自然言語処理に興味がある方

  • 自然言語処理を扱う会社のインターンに興味のある方

  • LLMの推論速度に関心のある方

時間がない方向けまとめ：

  • 推論速度に関して、特に大規模なモデルではCPUよりもGPUの方が圧倒的に速い

  • GPUで高速化したい場合vLLMかFasterTransformerを使うべき（モデルの大きさによる）

  • 特にバッチ処理をする場合はvLLMを使うべき

  • 本検証では量子化によるスループット向上は見られなかった


0. インターン参加の経緯

私はもともと研究スキルやエンジニアリング力を高めたい、学生のうちに企業で研究をしたいという理由でインターン先を探していました。その中で、2023年3月に言語処理学会年次大会（NLP2023）に参加した際に、レトリバのブースに「インターン募集」の文字を見つけ、声をかけたのがきっかけです（後から気づいたのですが、過去のRetrieva TECH BLOGでその時の私の発表について紹介していただいてましたmm）。

1. 背景：大規模言語モデルの高速化

ChatGPT（GPT-3.5）[1]を皮切りに世界中で大規模言語モデル（以下、LLM）の開発が急速に進んでいます[2,3,4]。LLMのモデルサイズは性能を向上させるためにスケーリング則に従って肥大化し、最近のLLMは数十億・数百億以上のパラメータを有します。そのため、LLMを動かすには多くのメモリを有する高性能なGPUが必要となり、非常にコストがかかります。
LLMを運用するにあたって、推論速度はサービスの質や運用コストに関わる重要な指標です。その中でも、数十万円のGPUあるいはCPUでLLMを動かすことができるのか、そして、どれくらいの速度なのかは気になるところです。
最近では、LLMを高速化するための様々な工夫が考案されています。次のセクションで高速化ツールとその実装について簡単に説明します。

2. LLM高速化の技術と実装紹介

2.1. 量子化

深層学習モデルは一般的に32bitあるいは16bitの浮動小数点で動作します。これを8bitや4bitで表現することを量子化といい、省メモリ化および高速化を実現できます。
量子化のより詳細な説明は、本インターンでメンターをしていただいた勝又さんが執筆した過去ブログ「深層学習の量子化に入門してみた～理論編～」をご覧ください。

実装
量子化はHuggingFaceにサポートされているbitsandbytesを用いて簡単に実装できます。（詳細はこちら）

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16)

model = AutoModelForCausalLM.from_pretrained("huggingface/model/dir", quantization_config=bnb_config)
tokenizer = AutoTokenizer.from_pretrained("huggingface/model/dir", padding_side='left')

in_tokens = tokenizer(["吾輩は猫である。", "名前はまだない。"], padding='max_length', max_length=64, truncation=True, return_tensors="pt")
out_tokens = model.generate(**in_tokens)

2.2. BetterTransformer

BetterTransformerは以下の2つの工夫によりTransformerモデルの計算ボトルネックを緩和することで高速化を実現するライブラリです。[huggingface]

  • カーネル融合: Transformer内の行列積やConcatなどの複数カーネルを同時に処理。

  • スパース性: パディングトークンに対する不要な計算を除去。

※ 本記事執筆時的では主にTransformerEncoder、TransformerEncoderLayer、MultiheadAttentionに対して実装されており、それ以外はPytorchの標準実装が用いられるためほとんど高速化はされないようです。

Currently, the BetterTransformer speedup only applies to transformer encoder models used in inference. … If the criteria are not met, control flows to the legacy PyTorch 1.11 Transformer implementation which has the same API, but lacks the fastpath performance boost.

実装
BetterTransformerはHuggingFaceのoptimumライブラリを通じて簡単に利用できます。（詳細はこちら）

from transformers import AutoModelForCausalLM, AutoTokenizer
from optimum.bettertransformer import BetterTransformer

model = AutoModelForCausalLM.from_pretrained("huggingface/model/dir")
bt_model = BetterTransformer.transform(model, keep_original_model=False)
tokenizer = AutoTokenizer.from_pretrained("huggingface/model/dir")

in_tokens = tokenizer(['吾輩は猫である。','名前はまだない。'], padding='max_length', max_length=512, truncation=True, return_tensors="pt")
out_tokens = model.generate(**in_tokens)

2.3. ONNX Runtime

様々な機械学習モデルの推論を高速化・効率化するためのランタイム環境です。グラフ最適化により、複数カーネルの融合、定数値の事前計算および入出力に影響を与えない演算の削除などを行うことで高速化を実現します。様々な実行環境（C++/Python/javaなど）で実行できるという特徴があります。
※モデルファイルのPytorchからONNXへの変換処理が非常に重いというデメリットがあります。

実装
ONNXもBetterTransformerと同様、HuggingFaceのoptimumライブラリを通じて簡単に利用できます。（詳細はこちら）

from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForCausalLM

model = ORTModelForCausalLM.from_pretrained("huggingface/model/dir", export=True)
tokenizer = AutoTokenizer.from_pretrained("huggingface/model/dir")

in_tokens = tokenizer(["吾輩は猫である。", "名前はまだない。"], padding='max_length', max_length=64, truncation=True, return_tensors="pt")
out_tokens = model.generate(**in_tokens)

2.4. vLLM

vLLMはPagedAttentionの導入によりKVキャッシュのボトルネックを改善することで高速化を実現するライブラリです。本来のAttentionではキーとバリューをKVキャッシュとしてGPUメモリに保持しており、そのサイズは出力文章の長さに依存します。出力文章長は事前にはわからないため、事前に過剰にGPUメモリを先取りしてメモリ効率が悪くなります。vLLMでは、入力トークンを一定の長さごとに分割してGPUメモリをジャストインタイムに保持することでGPUメモリ効率を向上させるPagedAttention[5]を導入しています。他にもContinuous Batching[6]によりバッチ処理の効率を上げる工夫がなされています。
※量子化がサポートされていないというデメリットがあります。

実装
vLLMを用いた簡単な実装を紹介します。（詳細はこちら）

from vllm import LLM, SamplingParams

prompts = ["吾輩は猫である。", "名前はまだない。"]
sampling_params = SamplingParams(temperature=0.5, top_p=1.0, top_k=50)
llm = LLM(model="huggingface/model/dir", dtype='float16')
outputs = llm.generate(prompts, sampling_params)

2.5. FasterTransformer

FasterTransformerはTransformerベースのモデルのTransformerブロックをC++/CUDAで記述され最適化されたTransformerブロックに変換することで高速化を実現するバックエンドです。Transformerブロックの最適化には以下の工夫が含まれます。
レイヤーフュージョン: 複数レイヤーを1つのレイヤーに結合。
マルチヘッドアテンションの高速化: KVキャッシュの維持。
GEMMカーネルの自動チューニング: 行列計算の最適化。

利用方法
実装はFasterTransformer用に環境を準備する必要があります。少し複雑ですがFasterTransformerを動かす部分まで紹介します。（詳細はこちら）

# docker run
nvidia-docker run -ti --shm-size 5g --rm nvcr.io/nvidia/pytorch:22.09-py3 bash

# Install FT
git clone <https://github.com/NVIDIA/FasterTransformer.git>
mkdir -p FasterTransformer/build
cd FasterTransformer/build
git submodule init && git submodule update
pip3 install fire jax jaxlib
cmake -DSM=xx -DCMAKE_BUILD_TYPE=Release -DBUILD_PYT=ON -DBUILD_MULTI_GPU=ON BUILD_MIXED_GEMM=ON ..
make -j12

# Install git-lfs
apt update && apt upgrade
apt-get install git-lfs

# Install transformer module
pip install transformers && pip install transformers[sentencepiece]

# EXAMPLE: rinna/japanese-gpt-neox-3.6b
git lfs install && git lfs clone <https://huggingface.co/rinna/japanese-gpt-neox-3.6b>
# convert
python ../examples/pytorch/gptneox/utils/huggingface_gptneox_convert.py -i japanese-gpt-neox-3.6b -o japanese-gpt-neox-3.6b/c-models/ -i_g 1 -m_n gptneox
# test
python ../examples/pytorch/gptneox/gptneox_example.py \\
    --ckpt_path japanese-gpt-neox-3.6b/c-models/1-gpu \\
    --tokenizer_path japanese-gpt-neox-3.6b \\
    --lib_path lib/libth_transformer.so \\
    --max_seq_len 64 \\
    --output_len 256 \\ # この場合、256トークンになるまで文末トークンでpaddingされます
    --max_batch_size 1 \\
    --inference_data_type fp32 \\
    --enable_random_seed \\
    --time \\

3. 速度検証

概要
本記事の実験では、様々な設定（特に安価なデバイス）で、上で紹介した技術を用いてLLMの速度検証を行います。ただし、本実験ではサービングではなくライブラリとしての速度検証を行いました。具体的には、以下の4つの実験で検証しました。

  1. CPUはGPUに推論速度でどこまで迫れるのか

  2. 入力系列長による推論速度変化

  3. バッチサイズによる推論速度変化

  4. 量子化で推論速度はどれだけ高速化されるのか

設定
モデルをcyberagent/open-calmからsmallおよび1bモデル、rinna/japanese-gpt-neoxからsmallおよび3.6bモデルとし、種類とモデルサイズの異なる4モデルで検証します。また、デフォルト設定として、入力系列長を64、最大出力系列長を256、浮動小数点をfp32としました。入力テキストは「吾輩は猫である」から入力系列長分用いました。

評価指標
評価指標にはスループットを用いました。本検証で用いたスループットは以下の式で表され、1秒間に何トークン生成できるかを指します。

補足
FasterTransformerに関して、出力文が途中で終了した場合、文末トークン（<|endoftext|>や</s>など）により最大出力系列長までパディングされます。FasterTransformerによって算出されるスループットはトークン数に文末トークンも含まれるため、文末トークンが多い場合はスループットが異常に高くなることがあります。そのため、FasterTransformerの実験結果のスループット値は、何度か実行して、出力が最大出力系列長（文末トークンによるパディングがない状態）になった時の値となっています。

3.1. CPUはGPUに推論速度でどこまで迫れるのか

本実験では、2種類のCPU（i9-13900KF、Intel-Xeon）と3種類のGPU（RTX4090、A10、T4）を用いて4つのモデルを比較しました。下図は各デバイスで各高速化技術を用いたスループットを示しています。（vLLMおよびFasterTransformerではCPUがサポートされておらず、ONNXではモデル変換時の処理が重すぎて動かせないものがいくつかあったため、一部測定できていません。）

（PTはPyTorchを用いたHugginFaceのデフォルト実装、BTはBetterTransformer、FTはFasterTransformerを指しています。以下、同様。）
結果として、RTX4090でFasterTransformerあるいはvLLMを用いた場合が最も速いです。
CPUとGPUの比較では、smallモデルのPyTorchおよびBetterTransformerでi9がA10やT4に迫る性能を獲得しています。また、smallモデルではONNXのスループットが低いため、CPUでsmallモデルを動かしたい場合はHuggingFaceのデフォルト実装が良いという結果となりました。 （PyTorchとBetterTransformerの性能がほとんど変わらないのは、BetterTransformerは現状Encoderモデルのみ対応しておりDecoderモデルはまだサポートされていないためだと考えています。）
一方、GPUを用いる場合、smallモデルではFasterTransformerが最もスループットが高く、1b以上のモデルではvLLMがFasterTransformerを上回ることも確認できました。そのため、GPUで動かす場合はFasterTransformerかvLLMが良いという結果となりました。その中でも、実装のしやすさを考慮するとvLLMを使うのが良いのではないかと思います。

3.2. 入力系列長による推論速度変化

本実験では、実験3.1で最もスループットの高かったRTX4090を用いて、入力系列長に対するスループットの変化を確かめます。LLMを実際に運用する場面によって想定される入力系列長が異なるため、想定される入力系列長に対してどれだけのパフォーマンスなのか知りたいというモチベーションです。下図は4つのモデルに対して、入力系列長を$${2^4 \sim 2^{10}}$$で変化させたときのスループットを示しています。

全体の結果として、入力系列長を大きくするにつれてスループットは低下する傾向にあることが分かります。これは入力系列長が大きいと、入力テキストを処理するのに時間がかかるためです。
また、smallモデルではFasterTransformer、1b以上のモデルではvLLMのスループットの低下が顕著になっています。これは実験3.1とも一致していて、モデルが小さい場合はFasterTransformer、大きい場合はvLLMを用いるべきであることが分かります。
結論として、入力系列長の増加に伴ってスループットは低下します。また、モデルサイズが小さい場合はFasterTransformer、モデルサイズが大きい場合はvLLMを用いるべきです。

3.3. バッチサイズによる推論速度変化

本実験では、RTX4090を用いて、バッチサイズに対するスループットの変化を確かめます。LLMの実運用においてバッチ処理を想定する場合、バッチサイズを増やすとスループットはどうなるか知りたいというモチベーションです。下図は4つのモデルに対して、バッチサイズを$${2^0 \sim 2^6}$$で変化させたときのスループットを示しています。（3.6bモデルのONNXはモデル変換の処理が重すぎて実行できませんでした。また、FasterTransformerは途中で出力が終了すると、残りの長さまで<|endoftext|>でパディングされ、スループットが異常に高くなることがあるため、本実験では除きました。）

結果として、モデルサイズによらずvLLMのみバッチサイズの増加に伴ってスループットも向上しており、Pytorch・ONNX・BetterTransformerはバッチサイズの増加に伴ってスループットはやや低下します。これは[7]の理論と一致していて、vLLMはPagedAttentionによる高速化の他にContinuous Batchingという技術でバッチサイズに対する工夫がなされており、バッチサイズが大きいほどその恩恵を受けられるためだと考えられます。
結論として、バッチ処理を想定とする場合はvLLMを用いるべきです。

3.4. 量子化で推論速度はどれだけ高速化されるのか

本実験では、RTX4090を用いて低精度化・量子化によるスループットの比較を行います。下図は4つのモデルをFP32、FP16、int8、int4で動作した時のスループットを示しています。（vllmは量子化がサポートされていないません。）

結果として、FasterTransformer以外は低精度化・量子化によってスループットはほとんど向上しませんでした。

おわりに

本記事では、日本語LLMの実運用を想定した速度検証の一部をご報告しました。今回は日本語LLMのGPT-NEOXモデルに対応した高速化技術を用いましたが、他にもOPT[8]のLLMに対応するものであればFlexgenやDeepSpeed-FastGennなどもあります。
インターンではレトリバの新技術開発室の方々と毎朝の定例やミーティング、勉強会に参加させていただいて、チームでのタスク・進捗管理などのソフトスキルからコードレビューや幅広い技術までたくさんのことを学ばせていただきました。また、悩みごとを共有するとすぐに対応してくださるスピードも速くとても良い環境だと思いました。
5週間という短い間ではありましたが、メンターの勝又さんをはじめ、チームの西鳥羽さん、飯田さん、本当にありがとうございました！！

参考

  1. https://chat.openai.com

  2. [2303.08774] GPT-4 Technical Report (arxiv.org)

  3. https://stability.ai/stable-lm

  4. https://www.cyberagent.co.jp/news/detail/id=29479

  5. vllm/vllm/model_executor/layers/attention.py at bdd6b4c8bc3e5ac93553436514171ffad5926f0c · vllm-project/vllm (github.com)

  6. https://www.anyscale.com/blog/continuous-batching-llm-inference

  7. https://zenn.dev/rinna/articles/7d10e61f694611

  8. https://huggingface.co/docs/transformers/model_doc/opt