限られたメモリで効率的なLLM推論を可能にする「LLM in a Flash」
(huggingface.co)- LLM in a Flash は、DRAMより大きいLLMのパラメータをフラッシュメモリに置き、必要な部分だけをDRAMに取り込むことで、メモリが限られたデバイスでも推論を可能にする
- フラッシュメモリのアクセス特性を反映した 推論コストモデル に基づき、転送データ量を減らし、より大きく連続したチャンクの読み出しを促す
- windowing は以前に活性化されたニューロンを再利用し、row-column bundling はフラッシュの逐次アクセスの強みを活用してボトルネックを減らす
- 2つの手法を組み合わせると、利用可能な DRAMの最大2倍のサイズのモデル を実行でき、単純なロード方式と比べて CPU で4〜5倍、GPUで20〜25倍高速な推論を示す
- スパース性の認識、文脈適応型ローディング、ハードウェア指向設計が組み合わさることで、メモリ制約のあるデバイス でのLLM実行可能性が高まる
フラッシュメモリベースのLLM推論
- 大規模言語モデルはさまざまな自然言語処理タスクで高い性能を発揮するが、計算量とメモリ要求量 が大きく、DRAM容量が限られたデバイスでは実行が難しい
- LLM in a Flash は、モデルパラメータを フラッシュメモリ に保存し、推論中に必要なデータだけをオンデマンドでDRAMに取り込む方式でこの制約に対応する
- アプローチの中心には、フラッシュメモリの動作に合わせた 推論コストモデル がある
- フラッシュから転送されるデータ量を減らすよう最適化する
- データをより大きく連続したチャンクで読み出すよう設計する
転送量と読み出しパターンを減らす2つの手法
- windowing は、以前に活性化されたニューロンを再利用して、フラッシュからDRAMへ移さなければならないデータ転送量を減らす
- row-column bundling は、フラッシュメモリが逐次データアクセスに強いという特性に合わせて、フラッシュから読み出すデータチャンクのサイズを大きくする
- 2つの手法を併用すると、利用可能なDRAMより最大2倍大きいモデルを実行できる
- 単純なロード方式と比較した推論速度の向上は、ハードウェアごとに異なる
- CPUでは 4〜5倍 高速化
- GPUでは 20〜25倍 高速化
- この設計は、スパース性の認識、文脈適応型ローディング、ハードウェア指向設計を統合し、メモリが限られた環境でLLM推論を可能にする
1件のコメント
Hacker Newsのコメント
その論文は、すでにかなり複雑な Deja Vu論文のスパース性活用手法の上に構築されているので、理解するのに時間がかかった
Deja Vuは、重みのスパース性が低いモデルでも「文脈的スパース性（contextual sparsity）」は高いと見なしている。つまり、行列積の結果ベクトルには0が多く生じるが、どの位置が0になるかは入力によって変わる
このスパース性を利用すると、行列の一部の行の読み込みをスキップでき、性能上の利得を得るにはどの行をスキップするかを事前に予測する必要があり、その予測は低ランク行列で可能だ
Appleの論文は、この発見がRAMの読み込み性能だけでなく、帯域幅を大きく犠牲にせずにフラッシュメモリから読み込むうえでも役立つと示唆している
アテンション行列は比較的軽く、スパースに読み込むべき対象はFFNだ。また、FFNの入力ではなくReLU層の出力を予測すると、はるかに良いスパース性が得られる。「行列積の後、ReLUの前でこのベクトルスロットは負になる」と予測できれば、その行列の列を読み込まずに0を出力すればよい
FFNの大半の行はそもそも読み込む必要がなく、各FFNごとに最近使われた行のキャッシュを維持し、必要なときにフラッシュから更新すればよい、という提案だ。チャンク読み込みと射影層間の相関関係などもあるが、核心となる洞察はこの部分に見える
FFNはFeed Forward Networkで、トランスフォーマーでは最大級のブロックにあたる
実際に性能差が出始める前に、モデルのどれくらいを読み込まずに済むのかが気になる
たとえば全体をRAMに載せたときの性能の90%を維持するには、メモリは半分で足りるのか、90%が必要なのか、95%が必要なのかを知りたい
結局のところ、RAMを減らしたときに最大性能に対してどれだけ早く性能を失うかが問題だ。グラフは少ないRAMの状況でこのアルゴリズムと基本アルゴリズムを比較しているが、それとは別に良い疑問だ
8GBのモデル全体をスマートフォンのメモリに載せなくても良い性能を出せるなら、明らかに有用だ
論文全体は読む価値があり、かなり興味深い仕事だ: https://arxiv.org/pdf/2312.11514
論文で強調されていたのは、データ転送を最小化し、フラッシュメモリのスループットを最大化するための2つの補完的な手法だ
Windowingは、直近に計算したトークンの活性値を再利用しつつ、過去数トークンに必要なパラメータだけを読み込む方式で、重み読み込みのI/O要求数を減らす
Row-column bundlingは、up-projectionとdown-projection層の行と列を連結して保存し、フラッシュからより大きな連続チャンクを読めるようにすることで、大きなチャンクを読んでスループットを高める
理論的には、モデルの一部がフラッシュ上にあるときの単一トークン推論時間は、モデル全体がRAMにあるときの推論時間に、フラッシュ上にあるモデル部分を読み込む時間を加えたものに等しくなるはずだ
フラッシュへの書き戻しは不要だと仮定しているが、LLMの専門家ではないので間違っているかもしれない
層が10個よりはるかに多いなら、RAMには小さな一部だけを残して1層ずつ読み込めそうだ。たいていのLLMは層が数十個あるので、あり得そうに思える
推論中にRAMがボトルネックでなければ、現在の層を推論している間に次の層をフラッシュからRAMへDMA転送することもできそうだ。単一プロセッサのシステムでは、RAMのボトルネックのためうまくいかなさそうだ
デュアルプロセッサのシステムなら、一方のプロセッサが次の層をRAMに読み込んでいる間に、もう一方のプロセッサが前の層を推論して、少ないRAMでも非常に大きなLLMを動かせるのではないかと思う
今、新しいLLM AIマシンを組むための部品の山の横に座っている。z840のデュアルプロセッサなので、こういうものを自分で触ってみるのが楽しみだ
Appleデバイスは、競合他社の同等デバイスと比べるとRAMが非常に少ない傾向が目立つ。
その一因は、AppleのソフトウェアチームがObjective-Cのような、より効率的な言語を使っていることにあり、また一因は、iOSアプリが非常に多様な画面解像度をターゲットにしなくてもよいことにある。そのため、高解像度テクスチャを読み込んでから縮小することが比較的少ない。
さらに、Appleほどの規模で購入してもRAMは大きく安くならないため、RAM増設は他の機能を追加するよりもマージンに大きな打撃を与える。
しかしLLMは本質的にRAMを大量に消費するため、こうした選択がすべて跳ね返ってくる。どんなメモリ節約手法を使っても、RAMをより多く搭載した競合他社は、より大きく、より良く、より賢いモデルを載せられるようになる。
近いうちにMacを買いたいのだが、RAMをどれだけ注文すべきか本当に悩んでいる。予算が限られているのでさらに難しい。制約がなければ最低でも32GBにはしていたと思う。
AppleがRAMの価格設定を変えてくれることをまだ期待しているが、おそらくむなしい期待である可能性が高い。
関連する実装も見られる: https://opensource.apple.com/source/xnu/xnu-2050.18.24/libke...
またAppleデバイスは「jetsam」という機能もサポートしており、優先度の高いアプリが滑らかに動き続けられるよう、使っていないアプリやバックグラウンドにあるアプリを終了してメモリを確保する: https://developer.apple.com/documentation/xcode/identifying-...
静か、またはほぼ無音のLinuxノートPCのおすすめがあれば受け入れる。ほとんどはファンが強く回るし、静かさのためならCPU性能の一部は喜んで諦められる。静音モードを簡単にオン・オフできるならそれでもよい。
これまで見た中でMacBook Airの静粛性に匹敵するものはなく、自分が間違っていることを示す製品があるなら喜んで聞きたい。
当然、より安い、あるいはRAM交換が可能という利点もあればよい。主にMacBook Airは、Webベースのサービスと、大規模プロジェクトのコンパイルやホーム／セルフホスティングに使うLinuxサーバーへ接続するリモート端末として使っている。
スケジュールが前倒しになったのは確かだが、彼らの仮説が間違っているとは思わない。
このテーマについては限定的にしか理解していないのだが、この方式ならスマートフォンでオフラインLLMを動かせるようになるのか気になる。
可能なら、機密データをサーバーに送らずにAI支援のコンテンツモデレーションを行うなど、興味深いアプリケーションが多数開けそうだ。
例えば昨日、Apple A-seriesチップ（iPhone）でLLMを動かすGitHub Discussionがここに上がっていた: https://news.ycombinator.com/item?id=38703161
最近の記事が「AI」ではなくLLMと呼んでいる点が気に入っている。そうすればマーケティング上の誇張ではなく、特定の技術についての内容だと分かる。
これがFlashAttentionとどう違うのか気になる。似た用語を使っていながら、アブストラクトで違いを説明していないと混乱する。
追記: フラッシュフレームワーク内で2つの異なるメカニズムを拡張したもののようだ。論文タイトルはもっと良くできたかもしれないが、最初の数ページ以内に説明されている。
結論で「この機能がユーザーにどのように露出するのか」といった節を期待していたが、おそらく論文の範囲外の話なのだろう。
こうした機能がCoreMLのAPI呼び出しや設定として出てきて、例えば use_flash フラグを指定する必要があるのか、それともユーザーには見えないランタイム最適化になるのか気になる。
AppleがCoreMLやMetalなどの開発ロードマップを説明している良い発表やトークを知っている人がいるかも気になる。
Appleはイランの会社を買収したのか？
[0]: https://www.geekwire.com/2020/exclusive-apple-acquires-xnor-...
アプローチが違うのは分かるが、どちらもフラッシュメモリを活用するので、この論文が FlashAttention [1] に少なくとも言及するだろうと期待していた
[1] https://arxiv.org/abs/2205.14135
私の理解では、FlashAttention はローカルメモリ、特に SRAM をよりうまく活用するためのアクセスパターンに関するものだ。たとえば CPU の L1 キャッシュや GPU の同等の階層にデータを保持するようなもの
言い換えると、FlashAttention は DRAM より速い部分に関するもので、この論文は DRAM より遅い部分へよりうまくオフロードする問題を扱っている
「OPT 6.7B モデルは、たとえば FFN 層の内部で顕著な 97% の疎性を示す」という箇所で、この指標が正確に何を意味するのか分かる人はいる？
層の値の 97% が 0 という意味なのか、サイズの 3% まで圧縮できるという意味なのか気になっている
ただし、活性化される 3% は固定されていないため、残りの 97% を完全になくすことはできない。論文は、活性化する 3% をかなり正確に予測し、精度を大きく損なわずに高速化できると言っているようだ