Lock in $30 Savings on PRO—Offer Ends Soon! ⏳

LLMの効率化を支えるアルゴリズム

triwave33
September 24, 2024

 LLMの効率化を支えるアルゴリズム

2024.09.04

triwave33

September 24, 2024
Tweet

More Decks by triwave33

Other Decks in Technology

Transcript

  1. 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じゃない説)
  2. 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
  3. 5 100トークンの⼊⼒で50トークンの出⼒を返す時 100トークンを⼊⼒して、⼀気に50の出⼒が返ってくるわけではない(UX的にはそうだが…) 1-100 → 101 1-101 → 102 1-102

    → 103 ・ ・ ・ 1-149 → 150 というのが内部動作。出⼒は次の⼊⼒に回される。となると⼊⼒も出⼒も違いはない(結局全て⼊⼒になる)
  4. 6 1-100 → 101 1-101 → 102 1-102 → 103

    ・ ・ ・ 1-149 → 150 初回 (prefill)とそれ以降 (generate) で明確な計算負荷の差がある (is 何︖) 初回は重い処理。⼊⼒シーケンスの⻑さで負荷が決まる それ以降は⽐較的軽い処理。出⼒シーケンスの回数分繰り返す 💡
  5. 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)
  6. 10 attention layerの困りごと ⾏列積の計算負荷が⾼い(正確には無駄が多い)。上式だと 2dN2 FLOPS の演算が必要 ⼊⼒・出⼒トークン数 (N) が⻑いと、N2のオーダーで負荷がかかる

    (無邪気にcontext lengthを⻑くすると詰む) N d ・ d N = N N 実際はこの演算がhead数、レイヤー数、バッチサイズ数で掛け算される
  7. 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/
  8. 20 Flash Attentionのお気持ちと特徴 Flash Attentionのすごくざっくりとした説明: タイルの逐次読み込み中にmaxを更新しつつ、sumの値も巧妙に修正できる。そんなのsoftmaxの原理的に 無理じゃない︖と思うがsoftmax → 後段のvalueとのmatmulの⼯程まで考えるとできてしまう(興味があ ればここ⾒て⼿計算してみてください)。その結果、HBMの読み込みが極限まで最⼩化

    (1read/1write)され る。dropoutとmaskもSRAM上で同時に実⾏できる。 pros: 速い、メモリ効率が良い (O(N))、計算結果は完全に等価、学習時には途中結果を再計算して対処する cons: 低レイヤーを制御するため実装が⼤変 (CUDA版はinternの学⽣に3週間で実装させたらしい)
  9. 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!!!) 分のメモリ領域を常に静的に確保しておくのは⾮常に無駄
  10. 23 そもそもKV cacheのメモリスペースを動的に割り当てられたら良いのに… The sky is gray and ・・・ [EOS]

    無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 無駄 固定⻑のKV cache領域 1 2 3 200K 😣
  11. 25 Speculative decoding の背景 なんだかんだでnext tokenの⽣成 (→) は早くなったけど、結局1 tokenずつ逐次的に⽣成する (↓)

    ことには 変わりない 逐次的なシーケンスを並列に扱えたら速くなるのになぁ でもそれって未来を先に知ることだから本質的に無理かぁ やりたいなぁ…
  12. 27 Speculative decodingではtarget model (推論対象のLLM) よりも軽量なモデル (draft model)を使って、数 トークン先の予測結果を⽣成する target

    modelは予測結果を元に、各予測トークンの妥当性を並列(バッチ処理)で検証する draft modelのN歩先予測時間 + target modelのバッチ評価時間N が、target modelのN歩先予測よりも短いと きに、speculative decodingが有効になる。 ただし、投機的 (speculative) な予測のため、予測が外れた時のコストが発⽣する
  13. 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
  14. 32 KVキャッ シュ モデル (16.8GB) KVキャッ シュ KVキャッ シュ Llama2

    on g5.2xlarge メモリ容量とスループット・レイテンシの関係 request 1 request 2 request 3 3リクエストまでなら 同時実⾏余裕🦙🦙🦙🚀︕︕ KVキャッシュが⾜りるリクエスト数 (この図の場合N = 1-3) まではレイテンシーを損なわずにスループット が上がる ただし、バッチで同時に⾏列演算を実⾏する必要がある。(揃うまで待つ必要がある)
  15. 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巡⽬以降になるので)、 スループットは頭打ちになる (⾏列がいくら増えようが厨房の回転率は頭打ちなので)