砂嵐から文章が現れる — 拡散型LLMを手元のミニPCで速くした記録

ChatGPTのようなAIが文章を書くところを見たことがあると思う。左から右へ、一語ずつ吐き出していく。あれは演出ではなく本当にそう動いていて、次の一語が決まるまで、その次の語はまだこの世に存在しない。

DiffusionGemmaというモデルは、そこが違う。

まず256語分の、でたらめな語の並びを用意する。砂嵐のようなものだ。そこからモデルは、自信のある場所だけを少しずつ正しい語に置き換えていく。何度か繰り返すと、砂嵐の中から文章が浮かび上がってくる。写真の現像に近い。

一度に256語をまとめて相手にするのだから、一語ずつ書くより速いに違いない。私はそう思っていた。手元のミニPCで動かしてみたら、毎秒26語だった。

26語という数字の内訳

使ったのはAMDのRyzen AI MAX+ 395というCPUと、そこに内蔵されたRadeon 8060Sというグラフィックス機能である。単体のグラフィックスカードは挿していない。メモリは96GBあって、CPUとGPUがそれを分け合って使う。モデルのファイルは15.6GBあるので、ふつうのゲーム用GPUだと積んでいるメモリに入りきらない。この構成なら丸ごと乗る。

動かすと、1回の「磨き」に435ミリ秒かかっていた。そして磨きは22回繰り返される。256語を作るのに約9.6秒。割り算すると毎秒26語になる。

つまり遅さの正体は、一度に256語を作ることではなく、それを22回やり直すことのほうにあった。速くする道は二つある。1回あたりを短くするか、回数を減らすか。

まず1回の435ミリ秒を疑った。

犯人はCPUに違いない、と思った

拡散型のモデルは、磨きのたびに「各位置でどの語がどれくらい有力か」という点数表を作る。この点数表がとにかく大きい。位置が256個あって、語の候補が262144種類ある。掛けると6700万個の数値になる。

そのあと、位置ごとに「どれくらい迷っているか」を計算する。迷いの度合いが小さい位置から順に確定させていくためだ。この計算のためにプログラムは6700万個の数値を3回なめる。しかも指数関数という重めの計算を、合わせて1億3400万回ほど呼ぶ。

紙の上で見積もった。16スレッドで分担しても100ミリ秒から150ミリ秒はかかるはずだ。435ミリ秒のうち3割がCPUで、GPUは3割ほど手を止めて待っている——そういう絵を私は描いていた。

描いたところで、測っていなかった。

そこで、プログラムに時計を入れた。GPUに計算を投げてから結果が返るまでと、その後CPUが点数表をなめている間と、二つの区間をストップウォッチで挟むだけの、10行に満たない追加である。

diffusion_profile: 22 steps, 435.05 ms/step
  | gpu  373.51 ms (85.9%)
  | host  61.50 ms (14.1%)

CPUは62ミリ秒だった。見積もりの半分以下である。

外れた理由自体は、あとから考えれば当たり前だった。指数関数の計算は今どきのCPUなら何個かまとめて同時に処理できるし、16スレッドが本当に16倍近く効いていた。それでも、2倍以上外していたという事実は残る。この見積もりを信じたまま次の一手を選んでいたら、私は「CPUの計算を減らす」ことに時間を使っていた。実際にはそこには62ミリ秒しか置かれていない。

時計を入れる作業は10分で終わった。

天井の測り方

CPUが14%なら、残る86%、373ミリ秒のGPUが本丸である。ではこれは削れるのか。

削れるかどうかを言うには、比べる相手がいる。「このGPUが256語を一度に処理するとき、どれだけ速くできるのが上限か」を知りたい。

幸い、llama.cppにはllama-benchという測定用のプログラムが付いている。これに同じモデルを渡して、256語をまとめて処理させると毎秒969語と出た。裏返すと、256語で264ミリ秒。これがこのハードの実力である。

実測は373ミリ秒だった。差は約110ミリ秒。

この110ミリ秒が無駄なのかというと、そうではなかった。拡散型のモデルは、ふつうのモデルがやらない計算を2つ余分にやっている。

一つめは、262144種類の語それぞれの点数を出す計算を、1か所ではなく256か所すべてでやること。ふつうのAIは次の1語だけ決めればいいので、この計算は1回で済む。拡散型は全部の位置で同時にやる。

二つめは、前回の磨きで「こう思った」という情報を次の磨きに渡す仕組みがあること。これも大きな掛け算になる。

二つの計算量を数えると、それぞれ3780億回のかけ算とたし算だった。このGPUがこのモデルで実際に出している速度から逆算すると、合わせて90〜105ミリ秒。差の110ミリ秒とほぼ一致する。

つまりGPUの373ミリ秒は、264ミリ秒の「このハードの限界」と、100ミリ秒ほどの「拡散型である以上どうしても必要な計算」で、ほぼ埋まっていた。

念のため確かめたことが二つある。一つは、行列演算を速く回すための専用回路がちゃんと使われているか。起動時のログにmatrix cores: KHR_coopmatと出ていたので、使われていた。もう一つは、注意機構を高速化する「フラッシュ注意」という仕組みが効いているか。こちらはオンで376.08ミリ秒、オフで377.42ミリ秒。差は誤差の範囲で、効いていなかった。ただし注意機構はそもそもボトルネックではないので、直しても得はない。

1回あたりの435ミリ秒は、ほぼ限界だった。

名簿に載っていない

一つだけ、構造的に損をしている場所が見つかった。

さきほどのCPUの62ミリ秒は、本来GPUの中で片付けられる計算である。llama.cppにはそのためのコードが実際に用意されていた。使われていないだけだった。

コードを読むと、理由はあっけないものだった。

ggml_backend_reg_t reg = ggml_backend_reg_by_name("CUDA");
if (!reg) { reg = ggml_backend_reg_by_name("ROCm"); }
if (!reg) { reg = ggml_backend_reg_by_name("MUSA"); }
if (!reg) { return false; }

GPUを動かすための仕組みには何種類かあって、それぞれ名前が付いている。NVIDIA向けの「CUDA」、AMD向けの「ROCm」、そして今回使っている「Vulkan」。上のコードは、CUDAかROCmかMUSAという名前を探しに行って、見つからなければ諦める。Vulkanという名前はどこにも書かれていない。

性能が足りないのではない。名簿に載っていないから呼ばれない。

だからこのモデルは、Vulkanで動かしているかぎり毎回必ずCPUに点数表を送り返す。62ミリ秒はそのぶんである。AMD向けの「ROCm」に乗り換えれば、この関数はそのまま動きはじめる。

とはいえ、それで消えるのは62ミリ秒と、点数表を送り返す10〜20ミリ秒。合わせて2割ほどだ。残る355ミリ秒は、どの仕組みに乗り換えても必要な計算である。半日かけて環境を作り直して2割、というのが正直な期待値になる。

1回あたりで戦うのは、ここまでだと思った。

回しても何も起きないツマミ

「entropy_boundが速度と品質の主ダイヤル」——手元のドキュメントにはそう書いてあった。私が以前そう書いた。

磨きをやめる判断に関わるパラメータはいくつかあって、これはその一つである。回してみた。既定値0.1を中心に、0.02から3.2まで160倍。

entropy_bound	長文: 磨きの回数	計算問題: 磨きの回数
0.02	21	10
0.05	25	9
0.1（既定）	20	26
0.4	17	8
1.6	17	10
3.2	17	9

増えたり減ったりしているだけで、傾向がない。出てくる文章の質も変わらない。160倍振って何も起きないというのは、そもそも別のものを動かしている証拠である。

停止条件のコードを読み直した。磨きが止まるのは、次の二つが同時に成り立ったときだった。

• モデルの答えが前回と変わらなかった回数が、決めた回数に達した
• 全体の迷いの平均が、決めたしきい値を下回った

entropy_boundはこのどちらにも出てこない。このパラメータが決めているのは「1回の磨きで何か所を確定させるか」だけで、いつやめるかには関わっていなかった。

では二つめの「迷いの平均のしきい値」が本命か。こちらも振ってみた。既定は0.005である。

しきい値	長文	計算問題
0.001	41	39
0.005（既定）	39	36
0.02	35	36
0.05〜1.0	35	35

0.02から上は完全に横ばいになった。しきい値を1.0にするというのは「迷いの平均がどんな値でも合格」という意味なので、この条件は事実上外れている。それでも回数は0.02のときと変わらない。

つまり磨きが止まっているのは、迷いが小さくなったからではなく、答えが変わらなくなったからだった。モデルは勝手に収束していて、既定値はその収束点にぴったり張り付いていた。二つのツマミには、はじめから余地がなかったことになる。

打ち切りではなく、火加減だった

残ったのはmax_stepsという、磨きの回数の上限を決めるパラメータである。既定は48。

上限を下げれば回数は減る。当然だ。ただしそれは途中で強制的に止めるということで、文章が途中で壊れるはずだ——と思いながらコードを見たら、そうではなかった。

t = t_min + (t_max - t_min) * cur_step / S      （S = max_steps）

磨きには「温度」という調整値があって、最初は高く、進むにつれて下がっていく。高いうちは大胆に書き換え、低くなるにつれて慎重になる。鉄を熱してから徐々に冷ますのに似ている。

このSがmax_stepsだった。つまりmax_stepsは「何回で打ち切るか」ではなく「何回かけて冷ますか」である。16に下げると、48回かけて下げていた温度を16回で下げきることになる。回数が減っても、冷まし方としては最後まで完結している。

512語を生成させて、時間と、出てきた文章の質を見た。

max_steps	長文: 回数 / 時間	質	計算問題: 回数 / 時間	質
4	4 / 1.8s	崩壊	4 / 1.8s	崩壊
6	6 / 2.7s	崩壊	12 / 5.6s	ほぼ良好・正解
8	8 / 3.6s	崩れ	15 / 7.1s	良好・正解
12	24 / 11.3s	軽微な崩れ	17 / 7.7s	良好・正解
16	31 / 14.6s	良好	18 / 8.1s	良好・正解
24	38 / 17.5s	良好	19 / 8.6s	良好・正解
48（既定）	39 / 17.9s	良好	36 / 15.7s	良好・正解

16まで下げても文章は崩れなかった。速度は長文で1.2倍、計算問題で1.9倍になった。1回あたり435ミリ秒の壁と半日がかりの環境構築の話をしていたのに、設定を一つ変えるだけでこれである。

質の判定には少し工夫をした。「良い文章かどうか」を目で見て決めると、速くしたい気持ちのぶんだけ甘くなる。そこで、答えが一つに決まる問題を混ぜた。

電車が120kmを1.5時間で走り、その後60km/hで2時間走った。全体の平均速度は？

正解は240km ÷ 3.5時間で、およそ68.57km/h。これなら合っているか間違っているかを目で確かめられる。表の「正解」はその結果である。max_stepsが6以上なら、モデルはすべての条件でこれを正しく解いた。

崩れた文章には、崩れ方に特徴があった。max_stepsを8にしたときの長文はこうなる。

Core Innovation: Attention mechanism (re of RNNs/Lss). ** components:* Encoder Encoder,,,,--

語が重複し、途中で欠け、記号が散らばる。砂嵐を磨ききる前に冷めてしまった状態だ、と考えると納得がいく。

そしてもう一つ、危ない罠があった。上の表で、長文のmax_steps 4から8は時間が2秒前後しかない。速いのではない。512語作るはずが、途中で生成そのものが止まって256語しか出ていない。壊れた砂嵐の中にたまたま終端の記号が現れると、モデルはそこで書き終えたと判断してしまう。

毎秒何語という数字だけ見ていたら、この3行は「圧倒的に速い設定」に見える。実際には仕事を半分で投げ出しているだけである。

速くする作業の半分は、速くすることではない

一度に256語を作るのだから速いに違いない、と私は思っていた。実際には22回やり直していて、毎秒26語だった。

そこから先で私が触った数字を並べると、CPUの計算を減らす話は62ミリ秒しかなく、GPUの計算はほぼ限界で、環境を作り直しても2割で、主ダイヤルだと思っていたパラメータは160倍振っても何も起きなかった。効いたのは、温度の下げ方を48回から16回に変えることだけだった。

そして、そのどれも、測るまでは分からなかった。私は最初に「CPUが3割」という絵を描いていて、それは倍以上外れていた。10行の時計を入れる10分を惜しんでいたら、間違った場所を半日かけて掘っていたことになる。

残っている問いが一つある。ROCmに乗り換えれば2割は確実に取れて、そのうえでAMDの計算ライブラリがVulkanより速いかどうかは、まだ誰にも分からない。速ければもう一段、そうでなければ半日ぶんの授業料である。

これは、測ってみないと分からない。