もうでかいVRAMに高い金を払う必要は無くなるかもしれないのか?超分散LLM推論環境が爆誕

2024年10月4日 00:28

世は大容量GPU時代。
吾輩も、秋葉原で大容量VRAM搭載GPUの中古が出たときけばすぐに飛んでいき買ってくるということを繰り返している。ちなみに最近、秋葉原の大容量中古GPUは飛ぶように売れているので、見つけたら即買いしないとすぐに無くなる。

なぜ大容量GPUが欲しいかと言えば、それがAIの全ての力の源だからである。

しかし、これは我々のようなガチ研究勢の話であって、ビジネスパースン的には「いやあこれからはローカルLLMでセキュリティバッチリでしょう」みたいな話をしても、「んで、おいくら万円ですか?」と聞かれて、「えーと、GPU単体で500万円くらいでやす」とか言うと客も裸足で逃げていく。そもそもそれだけの価格を払っても、買えるかどうかわからない。

こないだデンバーの学会で、NVIDIAのジェンスン・ファンとMetaのマーク・ザッカーバーグが対談した時に、マークはずっと「とにかくGPUを売ってくれ」とジェンスンに懇願していた。今やGPUは砂の惑星DUNEにおけるスパイスであり、世界が渇望する全ての力の源なのだ。

どんな大資本が買おうと思っても、売ってくれない。
世界的に欠乏している資源、それが大容量GPUなのである。

これがクソ高い。だいたい80GBで500万円くらい。去年の今頃、僕が清水の舞台から飛び降りるつもりで買った時は同じものが300万円だった。馬鹿げている。

例えば8GBのGPUはAI界隈ではもはや冗談みたいな低スペックだが、8GBのVRAMを積んだ4060は5万円未満で買える。10個買っても50万円であり、500万円あれば100個買える。つまり、500万円あれば800GBのVRAMがあるのと同じなのだ。

GPUとCPUは車とエンジンのように例えられる。昔は80GBのGPU一つと8GBのGPU10個では全く同列に語ることは不可能だったが今はもう違う。この二つはほとんど同列である。それどころかCPUが10倍あることを考えると、後者の方が高性能にできる可能性すらある。

そのためのオープンソース技術が先般発表された。

TPI-LLMという。

これは、複数のノードに跨って推論を実行できる。
3GBのマシンを8台繋げて70Bのフル推論ができるらしい。
この手のやつは実際に動かしてみないと嘘松の可能性があるので動かしてみた。

Tanuki-8B-dpoをとりあえずローカル環境で4分割して動かしてみると、確かに動いた。ただ、推論速度は分割しない場合よりも遅い。なぜだろうと思うと、実はデフォルトのコマンドではGPUを使っていなかった。

GPUを使うには「--use_gpu」というフラグを渡す必要があるのだが、--use_gpuを渡したらエラーで止まった。

仕方ないので自分で修正することにしたところ、どうにかGPUで動くは動いた。VRAM消費量は合計で8GBくらい。優秀と言える。

$ python examples/run_multiprocess.py --world_size 4 --model_type llama --model_path ~/.cache/huggingface/hub/models--weblab-GENIAC--Tanuki-8B-dpo-v1.0/snapshots/0ac7ea5d3dc440df4c138c5db887d9332ea6cd80/ --prompt "吾輩は" --length 200 --memory_window 4 --use_gpu --torch_dist
猫である。名前はまだ無い。どこで生れたかとんと見当がつかぬ。何でも薄暗いじめじめした所でニャーニャー鳴いて居た事だけは記憶している。吾輩はここで初めて人間というものを見た。しかも後で聞くとそれは書生という人間中で一番獰悪な種族であったそうだ。この書生というのは時々吾輩を捕へて鼻でも引っかけやしないかと思ふと、とんと不安になってくる。食事を出す者

Tanukiがちゃんと青空文庫を学習していることがわかる。

どこを修正したかというと、--use_gpuを使うとところどころエラーで落ちるので、色々調べるとそもそもtransformersの実装にバグがあった。

~/.pyenv/versions/anaconda3-2023.09-0/envs/tpi-llm/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.pyを一部変更する。

 197     @torch.no_grad()
 198     def forward(self, x, position_ids):
 199         if "dynamic" in self.rope_type:
 200             self._dynamic_frequency_update(position_ids, device=x.device)
 201
 202         # Core RoPE block
 203         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shap>
 204         position_ids_expanded = position_ids[:, None, :].float()
 205         # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
 206         device_type = x.device.type
 207         device_type = device_type if isinstance(device_type, str) and device_type != "mps>
 208         inv_freq_expanded = inv_freq_expanded.to(device_type) #ここを追加
 210         with torch.autocast(device_type=device_type, enabled=False):
 211             freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose>
 212             emb = torch.cat((freqs, freqs), dim=-1)
 213             cos = emb.cos()
 214             sin = emb.sin()
 215             
 216         # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equival>
 217         cos = cos * self.attention_scaling
 218         sin = sin * self.attention_scaling
 219         
 220         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

これはTransformerのバグとも言い切れないのだが、まあ動かないことは動かないのでバグとも言える。

これでGPUを使って分散して動かすことができた。
ちなみにGPUを使わなくてもCPUだけでも推論そのものはできる。分散すればGPUすらいらないということになるのかもしれない。まあこの辺りのトレードオフはまだかかりそう。

最近、ヤン・ルカン先生が推してた分散推論フレームワークが他にもある。

こっちは分散で推論できると主張されているがどうやって分散して推論できるのかという方法があまりきちんと説明されていないので試していない。

TPIの方が分散メカニズムが見えやすいし改造しやすい。

ただ、分散することによって生まれるボトルネックの存在は如何ともしがたく、405Bとかを推論するには12GBくらいの民生用GPUが30台くらい必要になって果たしてそれは楽になったと言えるのかはよくわからない。まあラズパイが一つ8GBくらいのメモリがあるから、ラズパイ10個で80GBモデルが動いたらアツいかもしれないが実用的とはまだ言えない。

今月出るMacの新型に期待してるのだが、まだ分散で推論するのは専用ハードや大容量GPUが必要そうだ。