ハクソク

世界を動かす技術を、日本語で。

SwiftでのLLMトレーニング、パート1：Gflop/sからTflop/sへの行列乗算の最適化

1日前原文(cocoawithlove.com)

概要

  • Swiftによる独自の 行列積（Matrix Multiplication）最適化 の試み
  • Apple Silicon上の CPU、SIMD、AMX、GPU の性能比較
  • C実装（llm.c）と Swift実装の性能差 とその原因分析
  • Swift 6.2のMutableSpan や Relaxed演算 による最適化手法
  • 今後の記事ではAppleの MLフレームワーク も検証予定

SwiftでLLMトレーニング用の行列積を最適化する

  • 目的 ：Swiftで手書きの行列積コードを高速化し、LLMトレーニングに活用する方法の解説
  • 背景 ：PythonのMLライブラリは計算を自前で行わないため、Swiftで「すべて自作」したい動機
  • 参考実装 ：Andrej Karpathyの llm.c （GPT2互換C実装）をSwiftに移植
  • 計測対象 ：行列積カーネルを含む 前向き・後向き伝播の全体反復時間 で性能比較

llm.cの行列積の構造

  • matmul_forward 関数が前向き伝播のコア
    • 入力inp、重みweight、バイアスbiasを使い、出力outに計算結果を格納
    • 4重ループ構造だが、本質は val += inp[...] * weight[...] の繰り返し
  • 1回の学習反復で約0.2兆回の浮動小数点演算 が必要

C実装とSwift実装の性能比較

  • C実装（llm.c） ：Releaseビルドで 1トークン/秒未満、1反復7秒
  • 基本Swift実装 ：Cに忠実に書き換え、配列境界チェックも除去
    • 約15〜20倍遅い （1トークン19秒、20反復で30分以上）
    • 2.8 Gflop/s 程度の低性能

Swift配列のボトルネック

  • ArrayBuffer.beginCOWMutation() が最大の性能コスト
    • Swift配列の「ユニーク性」判定が毎回発生し、最適化を阻害
    • MutableSpan （Swift 6.2）を使うことで ほぼゼロオーバーヘッド で回避可能
      • out = out.mutableSpanで配列を直接書き換え
    • 前向き伝播には効果小 だが、 全体の学習反復は3倍以上高速化

Cの最適化とSwiftの違い

  • Cは-ffast-math による FMA（fused-multiply-add）命令 を利用
    • 1命令で掛け算＋足し算を同時に実行
    • アセンブリレベルで8回のfmadd命令に展開される
  • SwiftはFMAを自動で使わない （-ffast-math相当がない）
    • SIMDで4回の掛け算（fmul.4s）は使うが、加算が分離・mov命令も多い
    • Relaxed演算（Swift-Numerics） で FMA利用を明示的に許可 できる
      • Relaxed.multiplyAddやRelaxed.sumを活用
      • 精度より速度を優先したい場面で有効

まとめと今後

  • Swift標準配列のユニーク性判定 が大きな性能低下要因

  • MutableSpanやRelaxed演算 の活用で C並みの最適化 が可能

  • 今後はAppleの MLフレームワーク や Metal実装 も検証予定

    • Karpathyの動画やPython実装 で理論を学び、本記事で Swiftでの実践的高速化手法 を習得できる流れ
    • 「フレームワークなし」 の自作アプローチでLLMトレーニングを探求

Hackerたちの意見

この記事は本当に素晴らしいね。LLMの使い方に興味がない人でも、Swiftのパフォーマンス最適化についての素晴らしい記事だと思う。残念ながら、こういうテーマの書かれた資料はあまりないからね。AMX命令が本当に秘密なのか気になるな。理論的にはM4以上を使えばSME経由で取得できると思うけど、実際にSwiftからintrinsicを試したことがないから、ただの推測だよ。

└

SME経由で取得するって、何のことか全然わからない。AMXはM4でSMEに置き換わったんだよね。これは単なる「抽象的なintrinsic」じゃなくて、新しいユニットなんだよ。

1.1 Tflop/sって良いの？理論的には、M3 MaxのGPUは約15 Tflop/sの能力があるんだ。でも、この種のタスクの実際の上限は3-5 Tflop/sになると思う。これ、めっちゃ同意だわ。だから、基本的なGPUベンチマークを真剣に受け止めるべきじゃないんだよね。GPUからピークパフォーマンスを引き出すのは、CPUよりもずっと複雑なんだ。それが、Nvidiaが他のGPUメーカーと比べてまだソフトウェアの強みを持っている理由の一つだよ。CUDAには、データセットのピークパフォーマンスを引き出すために調整された小さなカーネルがたくさんあるからね。

└

このリンクはお気に入りに入れてて、たまに見返してる。ナイーブなカーネルとよく調整されたカーネルの違いについて、今まで見た中で一番良い書き方だと思うよ。https://siboehm.com/articles/22/CUDA-MMM

Matt GallagherとCocoaWithLoveは、iOS開発を学び始めた頃の大きなハイライトだよ。今でもこんな高品質な情報を発信しているのを見ると嬉しいね！

└

同感だわ。Matt GallagherのSwiftの並行処理やCocoaのメモリ所有権についての投稿は、どこよりもわかりやすいよ。この投稿は、まだ質問してないStack Overflowの質問に対する正しい答えだと思う。

TFAでは、Cバージョンがコンパイラに「-ffast-math」を使って融合乗算加算（FMA）を生成させていると言われているね。ML/AIは「-ffast-math」の使用が許容される数少ないアプリケーションの一つだけど、一般的にはFMAを得るために「-ffast-math」を使うべきじゃない。コンパイラによるFMA生成を有効にするためには、gccとclangの両方で正しいフラグは「-ffp-contract=fast」なんだ。「-ffast-math」は「-ffp-contract=fast」を有効にするけど、数値の正確性が重要なアプリケーションでは非常に望ましくない他のコード変換も有効にしちゃうから、目立ったパフォーマンス向上はほとんどないんだ。ML/AIやグラフィックス/ゲーム以外では、「-ffast-math」はその影響を完全に理解している専門家だけが使うべきだよ。実際、専門家であっても「-ffast-math」が役立つことは少なくて、むしろ「-ffast-math」にまとめられている多くのオプションの中から一部だけを選んで有効にする方がいいと思う。2026年になっても、ほとんどのコンパイラがデフォルトでFMAを生成しないのは、IBMでこの操作が発明されてから36年も経っているのに、ちょっとおかしいよね。圧倒的に多くのケースで、FMAを使うことでより正確な結果が得られるんだ。（これが当てはまらないのは、FMAなしで計算された特定の式で、丸めが相殺される場合だけだよ。）デフォルトオプションになっていない理由は、数値結果がFMAなしのレガシーコンピュータで得られた結果と異なるからで、それが素人ユーザーには驚きだったんだよね。だから、古い結果があまり正確でなくても、同じ結果を保証するためにFMAが無効にされたんだ。このレガシーシステムを模倣する方針は、ユーザーの混乱を避けるためだけに、ずっと前に廃止されるべきだったと思う。

ちなみに、Xcodeに付属しているclangのバージョンは、正しいコンパイラフラグを指定すればopenMPコードをコンパイルできるよ。libomp.dylibといくつかのヘッダーファイルを提供しなきゃいけないけど、Rを運営している人たちがコンパイルして配布用に提供しているから（彼らが必要としているからね）。Rを信頼できるなら、これも信頼できると思うよ。https://mac.r-project.org/openmp/