PaPoo
cover
technews
Author
technews
世界の技術ニュースをリアルタイムでキャッチし、日本語でわかりやすく発信。AI・半導体・スタートアップから規制動向まで、グローバルテックシーンの「今」をお届けします。

SwiftでLLMを速く動かすには? マトリクス乗算をGflop/sからTflop/sへ押し上げる試行錯誤

この記事のキーポイント

何をやっている記事なのか

Cocoa with Love の Matt Gallagher 氏が書いたこの記事は、ひとことで言うと ​「SwiftでLLMを訓練するとき、行列積をどこまで速くできるか」​ を追いかけた実験記事です。

LLM(Large Language Model、大規模言語モデル)は、ざっくり言うと大量の重み(weights)を使って入力を計算し、次の出力を予測する仕組みです。その中でも、いちばん重い処理の一つが matrix multiplication。要するに、

z += x * y

みたいな掛け算と足し算を、ものすごい回数くり返す世界です。

この「地味だけど重い」計算をどれだけ速くできるかで、学習全体の速度がかなり変わります。だから著者は、LLMの中核をなす行列積を、​C → 素のSwift → Swiftの新機能 → SIMD/AMX/Metal といった順で最適化していきます。

率直に言うと、こういう記事はかなり好きです。AIの話って派手な未来っぽさに寄りがちですが、実際は「ひたすら足し算と掛け算を速く回す」世界なんですよね。その生々しさがちゃんと見える。

背景:なぜSwiftでLLMをやるのか

著者はもともと、昔の研究用コードを掘り起こしたことがきっかけで、機械学習の世界に戻ってきたそうです。その後、PythonのPyTorchやTensorFlowも触ったものの、Pythonは自分で計算するというより、裏で動く別エンジンを操作している感覚が強く、​​「自分で全部コントロールしている感じがしない」​ と感じたと書いています。

そこで登場したのが Andrej Karpathy の llm.c です。これはGPT-2互換のモデルを、約1000行のCで書いた実装で、かなり読みやすく、しかも「何が起きているか隠されていない」のが魅力です。

著者はそれを見て、​すぐSwiftに書き換えた とのこと。ここ、かなりエンジニアっぽくて好きです。普通は「参考にする」で終わるのに、すぐ自分の手元に持ってきて動かし始める。その行動力が記事全体の熱量につながっています。

image_0002.svg

まずはベースライン:llm.c と素のSwift

記事の中心は、llm.c の matmul_forward という関数です。これは forward pass(推論側の計算)で使う行列積の処理で、入力と重みを掛け合わせて出力を作ります。

著者はこの処理を Swift に移植し、ほぼ同じ形のコードにしました。ところが、結果はかなり厳しいものでした。

ベース比較

実装 Tokens/s Training iterations/s llm.c比
llm.c 0.926 0.175 100%
Basic Swift 0.054 0.014 7.3%

これ、かなり衝撃的です。​Swift版はC版の15〜20倍遅い。LLMが1トークンを出すのに19秒近くかかり、20回の訓練反復には約30分。さすがにこれは実用的とは言いがたいです。

著者はこれを 2.8 Gflop/s 程度だと述べています。1999年にAppleが「PowerMac G4は1 Gflop/s級だ」と宣伝していたことを引き合いに出しつつ、今の状況としては到底満足できない、としています。こういう皮肉混じりの比較、ちょっと好きです。技術の進歩って、本来は夢があるはずなのに、最適化をサボると簡単に昔の広告レベルまで落ちるんですよね。

最初の改善:Arrayをやめて MutableSpan を使う

遅さの原因を Instruments で追うと、_ArrayBuffer.beginCOWMutation() が大きなコストになっていたそうです。

ここで軽く補足すると、Swift の Array は Copy-on-Write(COW)​ という仕組みを使っています。これは「本当に書き換える瞬間までコピーを遅らせる」便利な設計なのですが、今回のような超高頻度ループでは、その所有権チェックや安全確認のオーバーヘッドが積み上がってしまうわけです。

そこで著者は、Swift 6.2 の MutableSpan を使います。Span は、配列の一部を「生の連続メモリ」として扱うための仕組みで、要するに Arrayよりも低レベルで、余計な確認を減らせる道具 です。

この変更だけで、訓練全体の速度がかなり改善しました。

MutableSpan導入後

image_0003.svg

実装 Tokens/s Training iterations/s llm.c比
llm.c 0.926 0.175 100%
Basic Swift 0.054 0.014 7.3%
Basic Swift + mutableSpan 0.056 0.042 24.0%

面白いのは、​forward pass自体はそこまで変わらないのに、training iterations全体では3倍以上速くなった ことです。つまり、学習処理全体を通すと、見えにくいメモリ操作のコストがかなり効いているということです。

ここは個人的にかなり重要だと思います。性能問題って、計算そのものより「安全性のための便利機能」が足を引っ張ることがあるんですよね。もちろん安全性は大事です。でも、数値計算では「どこまで安全チェックを残すか」のバランスが本当に難しい。

次の壁:なぜCのほうがまだ速いのか

改善しても、まだCに届きません。著者が次に注目したのは、loopの中心にあるこの式です。

value += inp[...] * weight[...]

ここはまさに「掛け算して足す」を何百万回もやる場所です。C側には -ffast-math があり、これによって FMA(fused multiply-add)​ が使えるようになります。FMAは「掛け算と足し算を1命令でまとめる」機能で、速い一方で、厳密な浮動小数点の振る舞いには少し寛容になります。

つまり、Cは「細かい正確さを少しゆるめる代わりに速くする」ことができる。こういう、実務での性能と精度のトレードオフは、いかにも低レベル最適化らしくて面白いです。

Swift側がここでどこまで追いつけるか、記事はさらに深掘りしていきます。タイトルにもある通り、最終的には Gflop/s から Tflop/s の世界まで狙っていくようです。Tflop/s は、ざっくり言えば 1兆回規模の浮動小数点演算を1秒で処理する という、かなりとんでもない領域です。

この先に見えてくるもの:CPU、SIMD、AMX、GPU

元記事の説明文には、​plain C / Swift / Metal まで含めて10種類の実装 があると書かれています。つまり、この記事は単なる「Swiftは遅い」話ではなく、

という、かなり本格的な性能比較の入り口になっています。

image_0004.svg

個人的には、この構成がすごく良いと思います。なぜなら、AIの計算性能って「GPUが速い」で終わりがちだけど、実際には CPU側の実装品質 が体感や開発体験にすごく効くからです。Swiftで書いたコードが、どこまで素直にハードウェア性能を引き出せるのか。そこに著者の興味があるわけです。

この記事が面白い理由

この元記事の面白さは、単なるベンチマークではなく、​​「Swiftで数値計算を本気でやると何が起きるか」​ を、かなり具体的に見せてくれるところです。

特に印象的なのは次の3点です。

1. 理論ではなく、実測ベースで殴っている

「Swiftは遅い」と雑に言うのではなく、実際に同じ処理を移植して、どこで何が詰まるかを測っています。こういう記事は信頼できます。

2. 新しいSwift機能が実戦で効く

MutableSpan のような機能が、ただの文法追加ではなく、​性能のボトルネック解消に直結する ことが示されています。これは言語進化の実例としても興味深いです。

3. AI記事なのに、主役は「地味な計算」

LLMというと華やかですが、結局は高速なループの塊です。その現実をちゃんと見せてくれるのがいい。夢のある話を、地面に引きずり下ろしてくれる感じがあります。

まとめると

この記事は、SwiftでLLMを訓練するシリーズの導入編として、​matrix multiplication をどう最適化していくか を追ったものです。

最初は素のSwiftがかなり遅く、C版に大きく水をあけられました。けれど、MutableSpan のような低レベル寄りの機能を使うことで改善の余地が見えてきます。そして、この先には fast-math、SIMD、AMX、Metal といったさらに強力な手段が控えています。

率直に言えば、​​「SwiftでLLM?」と聞くと最初はかなり無謀に見える のですが、だからこそ面白いです。無謀に見える挑戦を、ちゃんと計測しながら一段ずつ詰めていく。こういう地道な記事は、読んでいてすごく刺激になります。


参考: Training an LLM in Swift, Part 1: Taking matrix multiplication from Gflop/s to Tflop/s | Cocoa with Love

同じ著者の記事