要点テキストから画像を生成する…
解説
ねえねえ智也くん!この「Beyond GEMM-Centric NPUs」っていう論文、タイトルがかっこいいね!GEMMを超えろ!みたいな?
ああ、それはDiffusion LLMっていう、最近注目されてる新しいAIを高速に動かすためのハードウェアの研究だよ。GEMMっていうのは行列演算のことで、今のAIチップはそればっかり得意なんだけど、それだけじゃダメだって話だね。
えっ、行列演算が得意ならAIはバッチリなんじゃないの?
普通のLLMならそうなんだけど、Diffusion LLMはちょっと特殊なんだ。これは言葉を一度にドバっと並列で作るんだけど、その後の「どの言葉にするか選ぶ作業」、つまりサンプリングがめちゃくちゃ重いんだよ。GPUで動かしても、時間の70%がその作業に使われちゃうこともある。
70%!?それって、テストで計算はすぐ終わったのに、名前を書くのに時間のほとんどを使っちゃうみたいな感じ?
……例えは微妙だけど、まあ効率が悪いのは確かだね。サンプリングでは、膨大な単語のリストから確率が高いものを選んだり、メモリをあちこち読み書きしたりするんだけど、今のチップはそういう「こまごました作業」が苦手なんだ。
なるほどねー。じゃあ、この論文はどうやってそれを解決したの?
「d-PLENA」っていう新しい設計を提案してるんだ。まず、サンプリングの計算をハードウェアが処理しやすい「Stable-Max」っていう形に作り替えた。それから、上位の単語を素早く見つける「Top-k」専用の回路とか、新しい命令をたくさん追加したんだよ。
専用の道具を揃えたってことか!メモリの使い方も工夫してるの?
そう。データの種類に合わせて、ベクトル用、浮動小数点用、整数用ってメモリを完全に分けたんだ。これでデータの混雑を防いで、効率よく処理できるようになった。あと、メモリを上書きして再利用する戦略も取ってるね。
すごそう!それで、実際にどれくらい速くなったの?
NVIDIAのRTX A6000っていう、すごく強くて高いGPUと比べても、最大で2.53倍も速くなったんだ。同じ製造プロセスで比較してこの結果だから、かなり画期的だよ。
2.53倍!智也くんがカップラーメン作る間に、私は2杯半食べ終わっちゃうってことだね!
……君の食欲の話じゃないんだけど。でも、これが普及すれば、スマホとかの小さなデバイスでも、爆速で賢いAIが動かせるようになるかもしれないね。
未来だねー!でも、何か難しいところはないの?
課題としては、まだ特定のサンプリング手法に特化している部分があることかな。もっと色んな種類のAIモデルに対応できるように汎用性を高めるのが、これからの研究課題だね。
ふむふむ。じゃあ、私もサンプリングの練習しなきゃ!デパ地下の試食コーナーで、一番美味しいのを並列で選ぶ練習!
それはただの食いしん坊だろ。さっさと研究室に戻るぞ。
要点
- Diffusion LLM (dLLM) は、従来の逐次的な生成(自己回帰型)とは異なり、トークンを並列に生成できる新しいタイプの言語モデルである。
- しかし、GPUなどの既存ハードウェアでは、単語の候補から最適なものを選ぶ「サンプリング」工程が全体の処理時間の最大70%を占める大きなボトルネックになっていることが判明した。
- 既存のAIチップ(NPU)は行列演算 (GEMM) に特化しすぎており、サンプリングで必要な複雑なメモリ操作や比較演算が苦手である。
- 本論文が提案する「d-PLENA」は、サンプリング専用の命令セットと、メモリを効率的に使い分けるアーキテクチャを導入した新しいNPU拡張である。
- 評価実験の結果、d-PLENAはハイエンドGPUであるNVIDIA RTX A6000と比較して、最大2.53倍の高速化を達成した。