アップルがメモリ容量が限られたデバイス上で効率的な大規模言語モデルの推論を実現
3つの要点
✔️ 搭載するメモリ(DRAM)をオーバーするサイズの大規模言語モデルの推論を実行する方法を提案
✔️ フラッシュメモリに保存されたモデルパラメータのうち、目下の推論に必要な最小限のモデルパラメータだけをDRAMに高速に転送するため、windowingとrow-column bundlingを提案
✔️ 大規模言語モデルのモデルパラメータの半分しかDRAMに載らない場合において、提案手法は、素朴な方法に比べ、CPUでは4-5倍、GPUでは20-25倍の高速化を達成
LLM in a flash: Efficient Large Language Model Inference with Limited Memory
written by Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, Mehrdad Farajtabar
(Submitted on 12 Dec 2023 (v1), last revised 4 Jan 2024 (this version, v2))
Comments: Published on arxiv.
Subjects: Computation and Language (cs.CL); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
近年、ChatGPTを代表に大規模言語モデル（LLM）が社会の注目を集めています。
LLMはモデルパラメータ数が増えるほど、性能向上することが確認されているため、大規模言語モデルのさらなる性能向上に向け、より大規模なLLMが開発されています。
LLMの大規模化に伴い、大量のモデルパラメータ値を記憶し、それらのモデルパラメータ値に基づき、大規模な演算をする必要があるので、そのようなモデルパラメータ値を記憶し演算する”作業場”として大規模なメモリを積んだパソコンがなくては、LLMが動作しません。
多くの人にとっては、最先端の大規模なLLMを自分のパソコン上で動作させることはできず、ChatGPTのようなクラウド上のWebサービスとして利用することしかできません。ましてや、スマホなどの小型の携帯端末上で直接LLMを動作させることは難しくなっています。
クラウドサービスとは違い、スマホなどの携帯端末、いわゆるエッジデバイスに搭載され動作するAIはエッジAIと呼ばれ、ネットワークを介さず処理できることによる、処理速度の向上、個人情報や機密情報の漏洩リスクの低減が期待されています。
しかし、現状はすでに述べたように、LLMの大規模化のトレンドは止まらず、高性能なLLMをエッジデバイスに載せ動作させることは難しくなっています。
今回解説する論文では、iPhoneで有名なアップルの研究者が大規模なメモリを積んでいない、メモリが限られたデバイス上でLLMを効率よく動作させることが可能な手法を提案しています。※あくまでデバイス上で学習ができるものではなく、学習済みのモデルパラメータ値に基づき、効率よく推論ができる技術です。
Googleは2023年12月7日にLLMであるGeminiをPixel 8 Proに実装すると発表しています。スマートフォンの付加価値として、各社注目している可能性が高いです。Googleの場合は、デバイスのメモリに載りきるように、モデルパラメータ数をそもそも抑えたLLMをスマートフォンに実装しているようですが、今回のアップルの論文の技術は、デバイスのメモリに載りきらないようなLLMを効率よく実行することを狙った手法です。この提案手法をベースに実装されたLLM搭載のiPhoneが将来販売されるかもしれません。
では、その効果と仕組みを見ていきましょう。
素朴な手法と提案手法の推論速度比較
まず、提案手法の効果を示します。
図1に素朴な手法と提案手法の推論速度の比較を示します。これは、フラッシュメモリから素朴にLLMのモデルを読み込んだ場合と提案手法の比較です。
ここで、メモリとフラッシュメモリという言い方が分かりにくいかもしれませんが、これらは別の概念を指しています。
メモリは、通常よく言われるPCのメモリと同じものです。比喩的には”作業場”に対応する概念で、本論文ではDRAMという記憶装置を指しています。
PC上で、何らかの処理をさせようと思ったら、一旦データがメモリ上に展開されなければならないため、作業場に対応するのですが、PCの電源が切れると、データも失われてしまいます。
そのため、保持し続けたいデータは、SSDやSDカードなどのストレージに保存しておく必要があります。比喩的には”倉庫”に対応する概念です。
スマホで撮影した画像などはこのストレージに保存されており、ストレージとしてよく利用されているSSDやSDカードの記憶装置のコアは、フラッシュメモリと呼ばれるものです。基本的に、デバイスの起動時は学習済みのモデルパラメータ値もストレージに保存されています。
このように、本論文でのメモリはDRAM、フラッシュメモリはフラッシュメモリに対応するものとなっています。
推論速度の比較に話を戻すと、本論文ではLLMのモデルサイズの半分のメモリしか利用できないという制約の下、LLMの推論処理速度を比較しています。つまり、フラッシュメモリに保存されたLLMのモデルを丸ごとメモリ上に展開することが不可能な場合を考えています。
ベースライン（図のNaive）は、LLMのモデルサイズの半分はすでにDRAMに入っており、残り半分をフラッシュメモリから読み込んで計算する場合を考えているようです。提案手法と前提がずれているように見えて分かりにくいのですが、ベースライン自体はDRAMにLLMのモデルを丸ごと入れることが可能な設定になっているものの、1トークン(LLMでテキストを処理するときの処理単位。テキストを単語に分解して単語単位で処理しているようなことと思えばよいです)を推論するためにLLMのモデルの半分をフラッシュメモリから読み込む必要があるという設定にしているようです。
今回、図にあるように、LLMのモデルとしては、約70億のモデルパラメータを持ち、後述するモデルパラメータのスパース性を持つとされるFalcon 7BとOpt 6.7Bを対象としています。
図のComputeは計算時間、Load From FlashはフラッシュメモリからDRAMへのモデルパラメータの読み込みにかかる時間、Memory Managementはメモリ管理にかかる時間です。どちらのモデルにおいても、ベースラインの推論時間の割合から分かるようにフラッシュメモリからモデルパラメータを読み込む時間が推論のボトルネックになることがわかります。
本論文の提案手法では、このフラッシュメモリからのモデルパラメータの読み込みを効率的に実行することで、どちらのモデルにおいても高速化を達成しています。特に、Opt 6.7BにおいてGPU上で推論する場合は、素朴な方法では1トークンの推論（生成）にかかる時間は2秒以上かかるのに対して、提案手法では0.1秒以下に高速化できます。つまり、20倍以上の速さです。
フラッシュメモリとDRAMの記憶容量と転送速度、LLMのモデルサイズの関係
先程、実際に作業するのはメモリ（DRAM）で、データを保持しておく倉庫はフラッシュメモリであり、この往復時間のロスになることを説明しました。では、DRAMとフラッシュメモリの特性は具体的にどのように違うのでしょうか？データの記憶容量と転送速度の違いがポイントです。図2に、DRAMとフラッシュメモリ（Flash Memory）の転送速度（バンド幅）の比較を示します。
フラッシュメモリは100GB程度、DRAMは10GB程度の記憶容量を持っています。このように、一般には、フラッシュメモリは、CPUやGPUの内部メモリ(DRAM)に比べ、大きな記憶容量を持っています。フラッシュメモリは10倍のモデルパラメータを保存することができます。
そうであれば、フラッシュメモリで作業すればよいではないかとなりそうですが、フラッシュメモリはDRAMよりも低速です。フラッシュメモリのバンド幅（データの転送速度）は、フラッシュメモリが1GB/s程度に対し、DRAMは100GB/s程度です。つまり、フラッシュメモリは、DRAMに比べ10倍の記憶容量があるけれども、転送速度は1/100になるのです。計算が行われるCPUとDRAMの往復に比べ、CPUとフラッシュメモリの往復は推論時間にかなり響いてしまうのです。
LLMのモデルサイズ以上の記憶容量を持つDRAMがある場合、フラッシュメモリに保存しておいた学習済みのモデルパラメータ全体をDRAMに一度だけ読み込んで推論すればよいのです。
フラッシュメモリからモデル全体を読み込む時間はかかりますが、その後の推論はDRAMに読み込み済みのモデルで計算できます。さらに、LLMは計算量が多いので、CPUよりも大規模に並列計算できるGPUで処理されることが多いですが、GPUのメモリに必要なモデルパラメータを転送し、計算することになります。
対して、LLMのモデルサイズ以上の記憶容量を持つDRAMがない場合、どうすべきか？というのが本論文の課題です。
提案手法の仕組み
概要
本論文の提案手法LLM in a flashには、大きく2つの工夫ポイントがあります。
ポイント1は、データ転送量の削減です。具体的には、フラッシュメモリからモデルパラメータ全体を一気にDRAMに読み込まないことです。目下の1トークンの推論に本当に必要なモデルパラメータだけを転送するため、少数の非零のモデルパラメータを予測してそれのみを転送するスパース性の活用とフィードフォワードネットワークへの入力シーケンスを分割して、windowのスライド前後で必要なモデルパラメータの差分のみを逐次転送してゆくwindowingを提案しています。
ポイント2は、フラッシュメモリから読み込むデータのチャンクサイズ（一度に転送するデータサイズ）の最適化です。フラッシュメモリは、チャンクサイズが大きいほど、リード（読み込み）のスループット（単位時間当たりの処理量）が上がる傾向を持っています。スループットを上げ、フラッシュメモリからモデルパラメータを読み込む際のボトルネックを解消するため、チャンクサイズを向上するrow-column bundlingを提案しています。
ポイント1: データ転送の削減（スパース性の活用とwindowing）
LLMはトランスフォーマがベースモデルとなっていますが、トランスフォーマの各層にはアテンションレイヤーとフィードフォワードネットワークがあります。本論文では、アテンションレイヤーのモデルパラメータは常にメモリに保持します。これはモデルパラメータ全体の1/3を占めます。残りの2/3に当たるフィードフォワードネットワークのモデルパラメータに着目して、データ転送の削減を図っています。
今回対象にしているLLMのモデルであるOPTやFalconのフィードフォワードネットワークのモデルパラメータは非常にスパース（大半のモデルパラメータ値が0)です。フィードフォワードネットワークのモデルパラメータ値のうち、OPTは97%、Falconは95%が0です。0のモデルパラメータ値は計算に寄与しないので、本当に必要なモデルパラメータはそのうちの非零なモデルパラメータだけです。そこで、この非零のモデルパラメータだけを必要に応じて動的に読み込む、windowingを提案しています。実際にスパースかどうか確かめるとすれば結局DRAMに読み込んでいることになると思いますので、スパースかどうかを予測する予測モデルも作っているようです。これは既存にもある考えのようですが、現在の層のアテンションレイヤーの出力が分かれば、その先のフィードフォワードネットワークの出力が0になるかを予測できるのが特徴と示されています。アテンションレイヤーのモデルパラメータはメモリに常に保持しておくという設定は、この予測を行うためにも採用している設定だと思われます。
予測で本当にモデルパラメータ値が0かどうかを事前に見抜けているのかという懸念もあると思いますが、本論文では、3種類のZero shot task(LLMのファインチューニングをしない場合の評価タスク)で、精度評価を行っていますが、予測の有無でOPT 6.7Bモデルの精度はほとんど変わらないことが示されています。
windowingの概念図を図3に示します。
LLMではユーザから入力された、図のような単語の列を処理していきます。このような単語の列（入力シーケンス）を処理する際に、本論文ではSliding Windowを設定し、window内の単語にかかわるモデルパラメータを読み込んで推論を行います。Opt 6.7Bの場合、windowサイズは5です。window内の非単語の処理に関係する、非零と予測されるモデルパラメータ値のみを読み込みます。
ここで、さらに活性ニューロンに注目します。本論文では活性ニューロンとは入力トークンに対してフィードフォワードネットワークの各層において出力が正になるものと定義されています。不活性ニューロンの出力は足しても0で事実上計算にかかわらないので活性ニューロンだけ読み込めばよいと言えます。
windowingのポイントは、次のwindow内の単語の入力に対する活性ニューロンの大半は、前のwindowの活性ニューロンと共通するという点です。図3の青は、読み込む必要のあるニューロン（モデルパラメータ）、赤は読み込み不要のニューロン（モデルパラメータ）となります。少し濃い青がNew Neuronsで、前のウィンドウと異なり、新しく読み込む必要のあるニューロンです。そこで、windowingでは過去5トークンに関わる活性ニューロンを記憶しておきます。これで次のwindowで新しく読み込む必要のあるニューロンを削減し、データ転送を削減できます。
Falcon 7Bでウィンドウサイズにより必要なニューロン（モデルパラメータ）の転送量を図4に示します。
ｘ軸にウィンドウサイズ、y軸にDRAMに読み込まれたモデルパラメータ数の割合を示しています。Aggregated Usage（前のウインドウを記憶して差分だけを読み込む工夫をしない）の場合、ウィンドウサイズが大きいほど、より多くのモデルパラメータが計算に関わります。Incremental Transfer（windowingでの新しいニューロン分の読み込む）の場合、ウィンドウサイズが大きいほど、前のウィンドウと共通するモデルパメータが多くなり、新たに読み込む必要のある差分は少量で済むことを示しています。
ポイント2: チャンクサイズの最適化 row-column bundling
チャンクサイズとフラッシュメモリのリードスループットの関係を図5に示します。
横軸はチャンクサイズ、縦軸はリードスループット、線の違いはスレッド数の違いです。フラッシュメモリからデータを読み込む際に、チャンクサイズ、スレッド数を増やすほど、フラッシュメモリからデータを高速に読み込めることが示されています。このフラッシュメモリの特性を生かすため、チャンクサイズを増やそうとするのがrow-column bundlingになります。
OPT、Falconモデルにおいては、i番目の中間ニューロンの計算に、upward projectionのi番目の行とdownward projectionのi番目の列が必要なようです。フィードフォワードネットワークに存在する写像行列の値（モデルパラメータ）と考えておけばよさそうです。この対応する行列データ（モデルパラメータ）をまとめてリードすること（row-column bundling）でチャンクサイズを向上させます。
図６にrow-column bundlingの概念図を示します。
図6に関し直接の説明が本論文にあまりないですが、Predictor's Outputというのが先ほど説明したwindowingにかかわるところだと思われます。目下の1トークンを推論するの必要なニューロンの予測で、図ではPredictor's Outputの紫、黒、青、赤が必要なニューロンと考えられます。つまり、本来は、図6の8行のニューロンを読み込む必要がありますが、必要な4つのニューロンのみ読み込んでいることが示されていると考えられます。その必要なニューロンに関するUpward projectionの行（Up Proj Columns）とDownward projectionの列（Down Proj Rows）をまとめている様子を示しているのが右側で、それらをまとめてフラッシュメモリから様子を表していると考えられます。Upward projectionの行、Downward projectionの列のサイズがそれぞれd_modelとすると、二つをまとめることで2d_model、つまり2倍のチャンクサイズに増やせることが示されています。このように、提案手法では必要なフラッシュメモリからの読み込みデータ量を削減しつつ、できるだけチャンクサイズを大きくすることで転送速度を上げ、効率的な推論を実現できる仕組みになっていることが分かります。
終わりに
限られたメモリを持つデバイス上で推論が実行可能なLLMを実現する技術のポイントを解説しました。本論文は、LLMの民主化を追求し、より幅広い個人やデバイスでLLMの推論を利用可能にするための最初の試みと位置付けられています。今後もLLMの民主化を目指し、LLMを省メモリで利用するための多くの発展的な研究開発が行われていくと予想されます。
この記事に関するカテゴリー
