Tatsuya Urabe 2024.09.04 LLMの効率化・⾼速化を⽀えるアルゴリズム

基本事項

3 https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu- background/index.html https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron- hardware/trainium.html#trainium-arch Acceleratorのアーキテクチャ compute High bandwidth memory (HBM) • 処理時にcomputeとmemoryのどちらかがボトルネックになる。 (ex. memory-bound) • ⾼速化のアルゴリズムがどちらのボトルネックに寄与しているかを意識する • 現代のチップの性能においてはmemoryが⾜を引っ張ることが多い（HBM、全然Highじゃない説）

Decoderモデルのコンセプト initial (prefill) phase: 与えられた⼊⼒シーケンスを もとに初回のトークン⽣成を⾏うフェイズ decoding (generate) phase: 初回以降、逐次的に トークン⽣成するフェイズ ⼊⼒シーケンスのトークナイズ (The) (sky) (is) (blue) いわゆる「LMは次の単語を予測する」の話を もう少しブレイクダウン どちらのフェイズも次の単語を予測する意味 では変わらないが、トークン⼊⼒の粒度、実 際の (効率化) 計算において⼤きく異なる https://medium.com/@plienhar/llm-inference-series-2-the- two-phase-process-behind-llms-responses-1ff1ff021cd5

5 100トークンの⼊⼒で50トークンの出⼒を返す時 100トークンを⼊⼒して、⼀気に50の出⼒が返ってくるわけではない（UX的にはそうだが…） 1-100 → 101 1-101 → 102 1-102 → 103 ・ ・ ・ 1-149 → 150 というのが内部動作。出⼒は次の⼊⼒に回される。となると⼊⼒も出⼒も違いはない（結局全て⼊⼒になる）

6 1-100 → 101 1-101 → 102 1-102 → 103 ・ ・ ・ 1-149 → 150 初回 (prefill)とそれ以降 (generate) で明確な計算負荷の差がある (is 何︖) 初回は重い処理。⼊⼒シーケンスの⻑さで負荷が決まる それ以降は⽐較的軽い処理。出⼒シーケンスの回数分繰り返す 💡

Attention layer おさらい

Transformer (attention layer) Q N d ・ KT d N QKT = ＝ QKT N N Attention : QKT N N V N d ・ ＝ N d N: sequence length, d: head dimension 実装上はQ,K,Vのサイズはともに [batch_size, num_heads, N, d] Attention is all you need (arxiv)

QKT Q, K, V softmax(QKT/sqrt(d))V https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L313 実装例

10 attention layerの困りごと ⾏列積の計算負荷が⾼い（正確には無駄が多い）。上式だと 2dN2 FLOPS の演算が必要 ⼊⼒・出⼒トークン数 (N) が⻑いと、N2のオーダーで負荷がかかる （無邪気にcontext lengthを⻑くすると詰む） N d ・ d N ＝ N N 実際はこの演算がhead数、レイヤー数、バッチサイズ数で掛け算される

Attention計算効率化と KV cache

12 Attention layerの計算 1st sequence 3rd sequence グレー部分の値は未来の値になるので０でマスクする。(masked attention) トークンが追加されるたびに既存のトークンの⾏列を再計算している（無駄が多い） https://medium.com/@joaolages/kv-caching-explained-276520203249

13 最新のシーケンスに関する計算のみを⾏う。過去の値はHBMに格納してキャッシュしておく（KV cache） https://medium.com/@joaolages/kv-caching-explained-276520203249

KVキャッ シュ モデル （16.8GB） 14 KVキャッシュとデバイスメモリ KVキャッシュの最⼤サイズは以下で概算できる Where B: batch size, N: # of token Llama2 7B の場合、nlayers =32, nkv_attention_heads =32, dattention_head =128 なので、 B=1, N=4096 , bytes/param=2 (FP16) とすると （バッチあたりの）KVサイズは 約 2.15 GBとなる。 KVtotal_size (Bytes) =2 × B × N × bytes/param × nlayers × nkv_attention_heads × dattention_head 24GBのメモリ (ex. g5 instance ) の空き容量を考えると、モデルに14 * 1.2 = 16.8 GB使うと想定するとKVキャッシュ⽤に残されたメモリは7.2GB で、計算上3バッチしか同時実⾏できない (3バッチは同時実⾏できる)。 それ以上の同時リクエストはキューに積まれる（後述） そもそも固定 (Max) トークン⻑のキャッシュスペースを⽤意するのが無駄な のでは︖（伏線） KVキャッ シュ KVキャッ シュ GPUメモリ (HBM) の 使⽤内訳 レイヤー数 ヘッド数 cacheのdimension 最⼤トークン数 https://aws.amazon.com/jp/blogs/news/translate-jp-benchmark-and-optimize-endpoint-deployment-in-amazon-sagemaker-jumpstart/

Attention効率化の試み

16 解決⽅法1: Attention⾃体を改造する (ex. Sparse Attention) full attention (下三⾓形⾏列) に⽐べメモリフットプリントを節約できる https://www.hello-statisticians.com/ml/deeplearning/transformer3.html

17 解決⽅法2: K,V をヘッド間で共通化する 各Attention headのK,Vを全て共通化する (Multi-query attention)。 メモリ削減効率が⾼く、推論速度は上がるが性能が劣化する。 バランスを取るために、Headをグループに分けた上でK,Vを共通化する (Grouped-query attention) https://qiita.com/kernelian/items/9a61f3957ccd9155ee4e

18 いくらメモリの使⽤量を減らしてもそもそもHBMのI/Oが遅いのが根本原因 → Attentionの演算を⼯夫してHBMへのI/Oを極⼒減らそう、という試み 80 GB w/ 1.5-2.0 TB/s 192 KB w/ 19 TB/s 解決⽅法3: メモリのI/Oを減らす

19 Flash Attention Attention計算時にHBMへの読み書きを極⼒減らす⼿法 HBMのR/Wが⼤幅に削減し、実⾏時間が短縮。代わりに演算量は増えるが computeの能⼒が⾼いので無問題。 https://arxiv.org/abs/2205.14135

20 Flash Attentionのお気持ちと特徴 Flash Attentionのすごくざっくりとした説明: タイルの逐次読み込み中にmaxを更新しつつ、sumの値も巧妙に修正できる。そんなのsoftmaxの原理的に 無理じゃない︖と思うがsoftmax → 後段のvalueとのmatmulの⼯程まで考えるとできてしまう（興味があ ればここ⾒て⼿計算してみてください）。その結果、HBMの読み込みが極限まで最⼩化 (1read/1write)され る。dropoutとmaskもSRAM上で同時に実⾏できる。 pros: 速い、メモリ効率が良い (O(N))、計算結果は完全に等価、学習時には途中結果を再計算して対処する cons: 低レイヤーを制御するため実装が⼤変 (CUDA版はinternの学⽣に3週間で実装させたらしい)

Attention効率化の試み (advanced topic: PagedAttention and Speculative decoding)

22 ここまでの話 LLMのdecoder、特にAttention layerの計算効率を上げるための技術を紹介した 計算負荷を下げるために、過去の計算結果をキャッシュする⼿法が開発された（KV cache） モデルの巨⼤化も相まって、KVキャッシュの容量、HBMのI/Oの遅さが問題になってきた Attentionを改善する⼿法がいくつか開発された (Sparse Attention, Grouped query, Flash Attention) computeとmemoryのイタチごっこ感はあるが、切磋琢磨し合って性能が向上した 回収し忘れた伏線がもう⼀つあった。LLMの実際のcontext lengthは都度変わるのに、maxのlength (Claude だと200K!!!) 分のメモリ領域を常に静的に確保しておくのは⾮常に無駄

23 そもそもKV cacheのメモリスペースを動的に割り当てられたら良いのに… The sky is gray and ・・・ [EOS] 無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 固定⻑のKV cache領域 1 2 3 200K 😣

24 PagedAttention KV cacheを物理的に連続したメモリアドレスに配置するのではなく分割されたブロックに格納する。PagedAttention kernelはBlock tableを参照して論理的（仮想的）に連続した値として読み込む。これによりcache間でのトークン⻑を揃える必要がなくなり動的なメ モリ管理が可能になる。結果として2-4倍のスループットが向上が⾒込める コンピュータで古くより使われてきた仮想メモリとページングのテクニックをKV cacheに応⽤ Efficient Memory Management for Large Language Model Serving with PagedAttention (arxiv)

25 Speculative decoding の背景 なんだかんだでnext tokenの⽣成 (→) は早くなったけど、結局1 tokenずつ逐次的に⽣成する (↓) ことには 変わりない 逐次的なシーケンスを並列に扱えたら速くなるのになぁ でもそれって未来を先に知ることだから本質的に無理かぁ やりたいなぁ…

26 Speculative decodingのコンセプト Nトークン先 (↑ではN=3)までのトークン⽣成結果 (未来) を受け取る。 未来が正しいかの検証（3シーケンスの予測）は1 stepでバッチ処理できるので、⾃分で3step 予測するよりも 効率が良い。 https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/

27 Speculative decodingではtarget model (推論対象のLLM) よりも軽量なモデル (draft model)を使って、数 トークン先の予測結果を⽣成する target modelは予測結果を元に、各予測トークンの妥当性を並列（バッチ処理）で検証する draft modelのN歩先予測時間 + target modelのバッチ評価時間N が、target modelのN歩先予測よりも短いと きに、speculative decodingが有効になる。 ただし、投機的 (speculative) な予測のため、予測が外れた時のコストが発⽣する

リクエスト処理最適化と バッチング戦略

29 ここまでの話 単⼀のリクエストに関する最適化の技術、つまりモデル⾃体に対するパフォーマンス向上の話 ここからの話 エンドポイントにおける、時間差で降り注ぐリクエストをどう効率的にさばくか、の話。モデルのアーキテク チャやアルゴリズムだけでなく、推論⽅式、推論インスタンスの種類や数に依存する話。 req.1 req.2 req.3 LLM 例えば、リクエストの間⼝が⼀つのみで、 sequentialに受ける（そのリクエストが完了し たら次のリクエストを処理する）と効率が悪い

30 理想論 req.1 req.2 req.3 リクエストに対応できるだけのモデルレプリカを⽤意する（ノード内もしくは複数ノード）。各リクエストに すぐに取り掛かれるため、レイテンシも最⼩限に抑えられる。ただし、モデルパラメータ⾃体を複製する必要 があるため⾼性能なメモリ領域を多く使い、コストがかかる💰 KVキャッ シュ モデル （16.8GB） KVキャッ シュ モデル （16.8GB） KVキャッ シュ モデル （16.8GB） 以下、↑の形式ではなく⼀つのモデルでなるべくスループットを上げる条件・⽅法を掘り下げる

KVキャッ シュ モデル （16.8GB） 31 KVキャッシュとデバイスメモリ （再掲） KVキャッシュの最⼤サイズは以下で概算できる Where B: batch size, N: # of token Llama2 7B の場合、nlayers =32, nkv_attention_heads =32, dattention_head =128 なので、 B=1, N=4096 , bytes/param=2 (FP16) とすると （バッチあたりの）KVサイズは 約 2.15 GBとなる。 KVtotal_size (Bytes) =2 × B × N × bytes/param × nlayers × nkv_attention_heads × dattention_head 24GBのメモリ (ex. g5 instance ) の空き容量を考えると、モデルに14 * 1.2 = 16.8 GB使うと想定するとKVキャッシュ⽤に残されたメモリは7.2GB で、計算上3バッチしか同時実⾏できない (3バッチは同時実⾏できる)。 それ以上の同時リクエストはキューに積まれる（後述） KVキャッ シュ KVキャッ シュ GPUメモリ (HBM) の 使⽤内訳 レイヤー数 ヘッド数 cacheのdimension

32 KVキャッ シュ モデル （16.8GB） KVキャッ シュ KVキャッ シュ Llama2 on g5.2xlarge メモリ容量とスループット・レイテンシの関係 request 1 request 2 request 3 3リクエストまでなら 同時実⾏余裕🦙🦙🦙🚀︕︕ KVキャッシュが⾜りるリクエスト数 (この図の場合N = 1-3) まではレイテンシーを損なわずにスループット が上がる ただし、バッチで同時に⾏列演算を実⾏する必要がある。（揃うまで待つ必要がある）

33 KVキャッ シュ モデル （16.8GB） KVキャッ シュ KVキャッ シュ Llama2 on g5.2xlarge メモリ容量とスループット・レイテンシの関係 request token 1 request token 2 request token 3 これ以上は無理っす… 2巡 (ロット) ⽬でよろしく🦙 request token 4 queue 同時リクエスト数がKVキャッシュの容量を超えると、キューに回される。 その結果レイテンシーが急激に悪化し (2巡⽬以降になるので)、 スループットは頭打ちになる (⾏列がいくら増えようが厨房の回転率は頭打ちなので)

34 レイテンシーを損なわずに同時実⾏を捌ける領域 座席に空きがあるので突っ込める スループットが飽和してくる領域 座席は満席で⾏列ができて厨房は フル稼働 スループット・レイテンシの実例 (ベンチマーク結果) https://aws.amazon.com/jp/blogs/news/translate-jp-benchmark-and-optimize-endpoint-deployment-in-amazon-sagemaker-jumpstart/ 同時リクエスト数を変えた時に、レイテンシ (縦軸)・スループット (横軸)

35 3x 2x スループットが頭打ち 同時バッチ数で捌けないリクエストは2ロット ⽬、3ロット⽬に回されている https://aws.amazon.com/jp/blogs/news/translate-jp-benchmark-and-optimize-endpoint-deployment-in-amazon-sagemaker-jumpstart/

36 KV cacheのメモリスペースを担保している範囲内でなるべくリクエストをバッチとして扱うとよい。 ただし、バッチにまとめる分の待ち時間はレイテンシに乗ることになる。 req.1 req.2 req.3 req.1 req.2 req.3 同⼀バッチ

37 Dynamic batchingの場合、すべてのリクエストが完了してから次のバッチ処理に取り掛かるため その間、アクセラレータのIdle時間が発⽣する。 https://aws.amazon.com/jp/blogs/machine-learning/improve-throughput-performance-of-llama-2-models-using-amazon-sagemaker/ Dynamic batching リクエストを⼀定時間でまとめ、バッチ処理する⽅式。LLMに限らず、CVやNLPなどでも汎⽤的に使える。 シングル処理する⽅式に⽐べ、スループットが向上する

38 https://aws.amazon.com/jp/blogs/machine-learning/improve-throughput-performance-of-llama-2-models-using-amazon-sagemaker/ Continuous batching (a.k.a. rolling batching) Dynamic batchingの⽋点を補った⼿法 バッチ処理に空きが出たら、次のリクエストに回す (rolling)。 コンセプト的にはシンプルで分かりやすいが実際の実装はもっと複雑

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