この記事では、私は Swift で大規模言語モデル（LLM）をトレーニングする際、自力で記述した行列積乗算コードを可能な限り高速化することを試みました。その目標は、Swift における数学計算のコードを最適化する上で重要なステップへの洞察を提供することです。また、Apple Silicon の異なるユニット（CPU, SIMD, AMX, GPU）の性能規模感について感覚的な理解を得ることを願っています。

これは、Apple Silicon 上で Swift を用いてニューラルネットワークをトレーニングするというシリーズの第一弾となります。今後の記事では、Mac 向けの機械学習向けに Apple が提供している数多いフレームワークについても取り上げる予定です。行列積乗算や機械学習には、これらの確立されたフレームワークを使用するのが最も賢明です（彼らは行列演算のカーネル開発において私よりもさらに多くの年数を費やしています）。しかし、今しばらくの間、私は「ライブラリもゼロ」という方針で全てを自分自身のために記述することに楽しみを感じています。これには行列積乗算のカーネルだけでなく、サンプルアプリはこれらのカーネルを用いて完全な LLM の実装の一部として使用し、私が引用する数値はすべての前向きおよび逆伝播トレーニングイテレーション全体のものであり、このシリーズの基準実装は Andrej Karpathy の llm.c（GPT2 互換モデルの純粋な C 実装）です。これは比較的単純なモデルですが、必要なすべての構成要素を含んでおり、実際の仕事負荷を代表しています。つまり、Swift を C よりも速く最適化するという私のお気に入りのゲームの時間です。

目次

• 前略（Backstory）
• llm.c
• ベースの Swift
• Span, Egg and Span
• Chill vibes
• They see me rollin'
• マルチスレッド化
• 最高機密のハック
• 最大電力消費
• テスト環境
• 結論

前略（Backstory）

約2年前、私は2000年代初頭に書かれた工学論文から古いコードを掘り起こしました。それは画像分類にニューラルネットワークを使用する C++ で記述された画像認識システムでした。旧コードを再稼働させることを目指しましたが、長い間 ML 関連のコードに触れておらず、次第に面倒になり諦めてしまいました。

2024年初頭の LLM に関する議論の数々においても、Mac 上でニューラルネットワークがトレーニングされている様子はほとんど見られませんでした。少なくとも Swift という言語においては同様です。PyTorch や TensorFlow といった Python ライブラリを触りましたが、Python は計算自体を行わず、背後にある別の計算エンジンのオーケストレーターとして機能に過ぎず、その分離感から自分がコントロールできていないような感覚を抱きました。

その1ヶ月後、Andrej Karpathy が llm.c をリリースしました。これにより、従来の機械学習コンテンツとは異なる形で到達した理由は、何も隠されておらず、約1000行の純粋な C コードで構成されていることです。（ただし、かなり謎めいた変数名は使用されていますが）比較的読みやすいです。そのため当然のように、私はすぐにこれを Swift で再実装しました。そしてこれとの付き合いは大変楽しかったです。

もちろん、このコードとの遊戯には高速化させるための作業が必要でした。ここで少し伏線を引きます：初期の Swift 実装は本当に非常に遅かったのです。しかし最適化とは常に続くプロセスであり、常に試してみることが可能な余地があります。それが本書を導くことになりました：私が以前（そして過去1週間で追加した数件を含む）行った異なる探索を通じて walkthrough を行い、ライブラリに頼らず LLM を比較的高速にトレーニングする方法を示します。コードの大部分は Swift で記述されます（ただし最後に Metal 実装も示します）。

ところで、ニューラルネットワークや LLM の仕組みを解説することはないので注意してください。関心のある場合は、Karpathy の動画「Let's build GPT: from scratch, in code, spelled out.」が GPT 類似の LLM の仕組みを理解するためのほぼ definitivo なガイドです。さらに以前のシリーズ「The spelled-out intro to language modeling: building makemore」は5つのビデオシリーズで導入概念をカバーしており、より入門的なレッスンをお求めの場合はこちらをご覧ください。もちろん両方とも Python なので、Swift でどのように物事を進めるかを準備ができたらここで戻ってきてください。

llm.c

機械学習本質的には、モデルの重みを入力データに適用する（前向き伝播、通称推論）、そして誤差勾配を計算し、その重みを更新する（逆伝播）という作業です。通常、これらの計算はまとめてパッケージ化され、可能な限り高速実行しようと試みます。これらの演算のパッケージは、「線形テンソル射影」、「行列積乗算」、あるいは「ベクトル内積の連続」などと呼ばれるかもしれません（仕事単位をどれだけ大きくも小さくスライスするかによります）。究極的には、非常に多くの回

z += x * y

を実行するループです。

これらの行列積乗算は機械学習における作業の大部分を表すため、このコードに焦点を当てます。他の実装部分についても順次更新しますが、行列積乗算で示す同じ改善のみを使用します。

まず llm.c の

matmul_forward

から見ていきましょう。これは前向き伝播で使用される中心的な行列積乗算です。入力 (

inp

) をイテレートし、モデルの重み (

weight

) で乗算し、結果を実行中の合計 (

val

) に加算します。

void matmul_forward(float* out, const float* inp, const float* weight, const float* bias, int B, int T, int C, int OC) {
    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            int bt = b * T + t;
            for (int o = 0; o < OC; o++) {
                float val = (bias != NULL) ? bias[o] : 0.0f;
                for (int i = 0; i < C; i++) {
                    val += inp[bt * C + i] * weight[o*C + i];
                }
                out[bt * OC + o] = val;
            }
        }
    }
}

4 つのネストしたループは視覚的な複雑さをもたらしますが、実際にはその

val += inp[bt * C + i] * weight[o*C + i];

の一行がニューラルネットワークの心臓部です。まさに先ほどの通り：非常に多くの回

z += x * y

です。

どれだけの規模なのか？

val

行は2つの浮動小数点演算を含みますが、Karpathy は全トレーニングイテレーションにおける浮動小数点演算数は約

6 x N x D

と述べています（N はモデルの重みの数で私たちの場合は124,439,808、D は B * T = 4 * 64 = 256 が今回のアプリ）。つまり、トレーニングイテレーションあたり約

6 x 124,439,808 x 256 ≈ 1.911×10¹¹ ≈ 2兆

の浮動小数点演算を扱うことになります。

したがって、非常に速く実行する必要があります。

純粋な C コードは簡単に Swift パッケージで動作します。私は C 実装を常に

-O3

オプティマイゼーションレベルで実行するように修正しました（Xcode の設定に関わらず）。このオプティマイゼーションレベルであっても、C 実装はトレーニングイテレーション1つあたり約7秒、推論は1トークン未満の速度で動作します。素晴らしい概念証明ですが、有用になるはずの速さの10倍も遅いです。

ベースの Swift

基本的な Swift バージョンを C バージョンに忠実に保つよう最善を尽くしました：

static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) {
    for b in 0..<B {
        for t in 0..<T {
            let bt = b * T + t
            for o in 0..<OC {
                var value = bias?[o] ?? 0
                for i in 0..<C {
                    value += inp[bt * C + i] * weight[o * C + i]
                }
                out[bt * OC + o] = value
            }
        }
    }
}

C コードは本質的に「非安全」であるため、Swift コードにも同じ利点を与えるため

-remove-runtime-asserts

で実行し（アレイインデックスのランタイムチェックを削除）、常にアプリを「Release」設定で実行するようにしました。つまり Swift と C の実装は比較的同値なのではありませんか？

Debug モードで実行しないでください。 私が引用する数値はすべて Release 設定のものであり、このコードの一部を Debug で実行したことはありますが、20回に及ぶ全トレーニング走行を Debug で待って見たことなどはありません。私は通常 Xcode の Scheme を「Release」に設定し、デバッグ中もそのまますべて実行します。

前略を読んでいただければ、「非常に遅かった」とすでに述べています。Swift コードは15〜20倍も遅いです。LLM が19秒ごとに1トークン生成する速度です。このエンジンで20回に及ぶトレーニング走行を走破するにはほぼ30分必要になります。いったい何が起きているのでしょう？

この性能は約 2.8 Gflop/s を表します。1999年、Apple は PowerMax G4 の広告を出し、その 1 Gflop/s の性能が米国軍の眼中に「武器」であることを主張しました。さて今では 2.8 Gflop/s は全く許されざる低速です。

Span, Egg and Span

Instruments で確認したところ、前回の実行における最も大きなパフォーマンスコストは

_ArrayBuffer.beginCOWMutation()

でした。Swift は他者が我々のアレイを使用する可能性があると考え、それらがユニークである（したがってアレイコピーが発生しない）にもかかわらず、ユニークネスチェックのみで最大のオーバーヘッドとなっています。

え？時々、単なるバグかもしれない問題に遭遇します。これがその例かもしれません。私が2024年に初めてこのコードに取り組んだときは、これは問題ではなかったと記憶していますが、リグレッションが発生したか、何らかの安全穴が塞がれて

_ArrayBuffer.beginCOWMutation()

がパフォーマンスを妨害している可能性もあります。また、

inline(none)

で関数内联化を無効にするとこの問題は移動するため、最適化器が正しく機能していないようにも感じられます。

いずれにしても、アレイを使用することはできず、必要な性能は得られません。幸いにも Swift 6.2 では、ほぼオーバーロードなしで信頼性の高い解決策

MutableSpan

を提供しています。

static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) {
    var out = out.mutableSpan
    for b in 0..<B {
        for t in 0..<T {
            let bt = b * T + t
            for o in 0..<OC {
                var value = bias?[o] ?? 0
                for i in 0..<C {
                    value += inp[bt * C + i] * weight[o * C + i]
                }
                out[bt * OC + o] = value
            }
        }
    }
}

追加したのは最初の行にある

var out = out.mutableSpan

だけで、その独自の

mutableSpan

で

out

をシャドウしています。ファイル全体を通して同じパターンを適用しました。興味深いことに、この変更は前向き伝播には大きな影響を与えなかったものの、トレーニングイテレーション（前向き＋逆伝播＋更新）は3倍以上高速化されました。

Chill vibes

しかし、前向き伝播がなぜ遅いのかという点にも目を向けなければなりません。Instruments はすでに我々の知っていることを確認しています：前向き伝播のホットな行は、ループの中心にある

value += inp[bt * C + i] * weight[o * C + i]

です。

厳しい真実と向き合う時です：C には Swift が持っていないコンパイラ最適化フラグがあります。この特定のケースでは、C は

-ffast-math

を有しており、これにより C は浮動小数点乗算と加算を単一の命令で実行する fused-multiply-addition (FMA) コマンドを使用でき、一般的には正確な正しさについてはあまり気にしません。

+0xa34   fmadd               s0, s17, s16, s0
+0xa38   ldr                 s17, [x20, #0x4]
+0xa3c   fmadd               s7, s17, s16, s7
+0xa40   ldr                 s17, [x21, #0x4]
+0xa44   fmadd               s4, s17, s16, s4
+0xa48   ldr                 s17, [x22, #0x4]
+0xa4c   fmadd               s6, s17, s16, s6
+0xa50   ldr                 s17, [x23, #0x4]
+0xa54   fmadd               s1, s17, s16, s1
+0xa58   ldr                 s17, [x24, #0x4]
+0xa5c   fmadd               s2, s17, s16, s2
+0xa60   ldr                 s17, [x25, #0x4]
+0xa64   fmadd               s3, s17, s16, s3
+0xa68   ldr                 s17, [x26, #0x4]
+0xa6c   fmadd               s5, s17, s16, s5

C のインナー ループは、主に

fmadd

（FMA 命令）の8回のアンロール適用です。

面白い事実は：我々は未来に生きているため、アセンブリ言語のように物事が理解できなくても、お気に入りの LLM にドロップして翻訳してくれます。

Swift には

-ffast-math

がありませんので、代わりに別個の乗算と加算を行います。

+0x164   fmul.4s             v1, v1, v5
+0x168   mov                 s5, v1[3]
+0x16c   mov                 s17, v1[2]
+0x170   mov                 s18, v1[1]
+0x174   fmul.4s             v2, v2, v6
+0x178   mov                 s6, v2[3]

/*...*/

+0x1d0   fadd                s0, s0, s7
+0x1d4   fadd                s0, s0, s4
+0x1d8   fadd                s0, s0, s24
+0x1dc   fadd                s0, s0, s23
+0x1e0   fadd                s0, s0, s16

Swift は4倍のループアンロールを試みていますが（

fmul.4s

は4乗の SIMD 演算）、しかしそのすべての

mov

命令と最後の別個の加算が我々を妨げています。C が使用しているような FMA を使用する必要があります。

幸运的是、Swift-Numerics があり、我々の版のファスト数学のための「Relaxed」を提供しています。「ギャラクシーガイドへの挨拶」の"Don't panic"同様、「Relax」という単語は単に皆を穏やかにし、chill vibes を楽しむためにここにあるのかと思うか、それとも丸め結果周りのルールを緩めて FMA が可能にするのか。我々のすべての

a += b * c

および

x = y + z

操作は

Relaxed.multiplyAdd

と

Relaxed.sum

で改善できますが、私は

gelu_backward

関数にこれらの関数を適用しないように特に注意します（C 実装はこの関数で明示的に

-ffast-math

を無効化しているため）。

static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) {
    var out = out.mutableSpan
    for b in 0..<B {
        for t in 0..<T {
            let bt = b * T + t
            for o in 0..<OC {
                var val = bias?[o] ?? 0
                for i in 0..<C {
                    val = Relaxed.multiplyAdd(inp[bt * C + i], weight[o * C + i], val)
                }
                out[bt * OC + o] = val
            }
        }
    }
}

変化の唯一箇所は真ん中にあります：

val = Relaxed.multiplyAdd(inp[bt * C + i], weight[o * C + i], val)

。あまり綺麗ではありませんが、パフォーマンスを見てみましょう。

これでアセンブリは以下のように見えます：

+0x178   fmla.4s             v1, v16, v4
+0x17c   fmla.4s             v0, v17, v5
+0x180   fmla.4s             v2, v18, v6
+0x184   fmla.4s             v3, v19, v7
+0x188   add                 x6, x6, #0x40
+0x18c   add                 x30, x30, #0x40
+0x190   subs                x13, x13, #0x10
+0x194   b.ne                "specialized static LLMBasicSwift.matmul_forward(out:inp:weight:bias:B:T:C:OC:)+0x168"
+0x198   fadd.4s             v0, v0, v1
+0x19c   fadd.4s             v0, v2, v0
+0x1a0   fadd.4s             v0, v3, v0
+0x1a4   faddp.4s            v0, v0, v0
+0x1a8   faddp.2s            s0, v0

興味深いことに、Swift は SIMD ベクトル化バージョンの

fmadd

（

fmla

）を選択していますが、前提条件は同じです。これはトークンあたりほぼ10倍の速度向上です。しかしトレーニング性能は C よりも15%遅いままです。最終的なギャップを閉じるにはどうすればよいでしょうか？

They see me rollin'

何かを認めるときです：私は C の「素朴な」行列積乗算バージョンを示してまいりました。実際の C 関数は少し醜く、外ループを8ステップずつstrideし、コンパイラがインナー ループを展開することを期待しています。そして C の

-O3

はこれを承諾し、8倍のループアンロールを提供します。

for (int obt = 0; obt < B * T; obt += LOOP_UNROLL) {
    for (int o = 0; o < OC; o++) {
        float result[LOOP_UNROLL];
        for (int ibt = 0; ibt < LOOP_UNROLL; ibt++) {
            result[ibt] = (bias != NULL) ? bias[o] : 0.0f;
        }
        for (int i = 0; i < C; i++) {
            float w = weight[i + o * C];
            for (int ibt = 0; ibt < LOOP_UNROLL; ibt++) {
                int bt = obt + ibt;
                result[ibt] += inp[bt * C + i] * w;
            }
        }
        for (int ibt = 0; ibt < LOOP_UNROLL; ibt++) {
            int bt = obt + ibt;
            out[bt * OC + o] = result[ibt];
        }
    }
}

明白でないように見えますが、C 実装の重要な行は

float result[LOOP_UNROLL];

ブファーであり、8つの結果値を格納します。以前は Swift でこれを行うことができませんでした。ループ内で

Array<Float>

を割り当てることは単にコストが高すぎました。私の2024年実装では、私はただループを手動で8回アンロールするしかなかった（読むのが非常に醜いものです）。

しかし、Swift 6.2 はここで別の有用な機能を提供しました：C スタック配列とほぼ一致する

InlineArray

です。

for obt in stride(from: 0, to: BT, by: LOOP_UNROLL) {
    for o in 0..<OC {
        var result = InlineArray<8, Float>(repeating: bias?[o] ?? 0)
        let bt = inp.span.extracting(droppingFirst: obt * C)
        let w = weight.span.extracting(droppingFirst: o * C)

        for i in 0..<C {
            for r in result.indices {
                result[r] = Relaxed.multiplyAdd(bt[r * C + i], w[i], result[r])
            }
        }
      
        for r in result.indices {
            out[(obt + r) * OC + o] = result[r]
        }
    }
}

アセンブリは最後の例と同じ

fmla.4s

と

fadd.4s

命令ですが、より多くあるので分岐や他のオーバーヘッドが少し減少しています。重要なのは：我々は C 実装とできるだけ類似しています。

C と Swift は推論では本質的に同様の速度であり、Swift は現在トレーニングでわずかに速くなっています。Swift はここで速いとはいえ、C コンパイラにいくつかのフラグを振ることで SIMD 命令を出させ、C が再度先走らせることができるでしょう。重要な点は、両者は今や概ね同等であることです。

マルチスレッド化

llm.c

コードには多くの

#pragma

アノテーション（例：

#pragma omp parallel for

）が含まれています。これらは Swift パッケージマネージャーで使用する通常の clang コンパイル出力では動作しませんでしたが、これらのアノテーションは OpenMP（Open Multi-Processing）を意図しています。OpenMP は通常修正された C コンパイラを必要としますが、それが

llm.c

コードの意図する実行方法です。

ここから Swift が最終的に C にリーダビリティ面で負けるという認識が生まれました。ループに並列性をタグ付けする方法がないだけでなく、アレイをスライスして別々のスライス上で同時に作業する Swift 6 安全な方法もありません。

MutableSpan

のようなアプローチは、変数を同時アクセスしているとして文句を言います（たとえそれらが別々のスライスから作られていても）。

なぜか Swift 同期ではありませんか？Swift

TaskGroup

が存在し、部分的に同じ役割を果たすことができますが、

DispatchQueue.concurrentPerform

は完了までブロックするため重要で、我々の

Span

インスタンスは脱出不可能だからです。つまり、すべての変更可能なアレイは

withUnsafeMutableBufferPointer

と

@unchecked Sendable

ラッパーを使用して、コンパイラが我々が試みていることに文句を言うのを防ぐ必要があります。

我々にとっての同期選択は

DispatchQueue.concurrentPerform

です。これは重要です。なぜならそのクロージャーが

@escaping

としてマークされていないため、コピーや参照カウントチェックなどのオーバーヘッドをかけずに不変アレイを

Span

として渡せるからです。

let tileCount = BT / LOOP_UNROLL
let workerCount = max(1, ProcessInfo.processInfo.activeProcessorCount)
let chunkSize = max(1, (tileCount + workerCount - 1) / workerCount)
let chunkCount = (tileCount + chunkSize - 1) / chunkCount
let bias = bias?.span
let inp = inp.span
let weight = weight.span

out.withUnsafeMutableBufferPointer { outBuffer in
    let outStorage = SendableUnsafeMutableBuffer(baseAddress: outBuffer.baseAddress!)

    DispatchQueue.concurrentPerform(iterations: chunkCount) { chunk in
        let startTile = chunk * chunkSize
        let endTile = min(tileCount, startTile + chunkSize)

        for tile in startTile..<endTile {
            let obt = tile * LOOP_UNROLL

            for o in 0..<OC {
                var result = InlineArray<8, Float>(repeating: bias?[o] ?? 0)
                let bt = inp.extracting(droppingFirst: obt * C)
                let w = weight.extracting(droppingFirst: o * C)

                for i in 0..<C {
                    for r in result.indices {
                        result[r] = Relaxed.multiplyAdd(bt[r * C + i], w[i], result[r])
                    }
                }

                for r in result.indices {
                    outStorage[(obt + r) * OC + o] = result[r]
                }
            }
        }
    }
}

これは文字として最悪のコードではありませんが：

• タイルやチャンクの計算、Span によるシャドウイングなど多くのオーバーヘッドがある

• out.withUnsafeMutableBufferPointer

  スコープと

  DispatchQueue.concurrentPerform

  スコープは確かに視覚的な乱雑さをもたらす

これは「C コードの方が良く見える」と私が考える時点です。

RangeReplaceableCollection

のスライスして同時に実行する操作があれば、

for

ループ行を置換するほどシンプルにできるでしょう。

私はこのパターンをトレーニングイテレーションの4つの hottest ループ：

matmul_forward

,

matmul_backward

,

attention_forward

, および

attention_backward

に適用しました。

llm.c

を OMP 経由で実行しなかったため、トレーニング性能は llm.c GitHub ページで Karpathy が引用した数値から取られています（我々はどちらも M3 Max CPU を使用しているため、大まかに比較可能です）。5.4倍の改善です。まあいいほうですが、私は16コア持っているCPUなので、マルチスレッド化でのわずかな5倍の改善は完全な利用ではありません。現在、メモリトラバーサルによって制限されている可能性があります。

私が最も問題と感じるのは、ループアンロール、非安全ポインター、チェックなし Sendable ラッパー、同時実行イテレーション追跡の間、コードが非常に醜くて扱いにくいことです。もっと良くできます。

最高機密のハック

我々はおそらく十分に頑張りました。しかし、Mac で最も速い CPU ベースの行列積乗算を得るためには、

Relaxed.multiplyAdd

と基本的な

fmla.4s

命令が提供する基本 SIMD パフォーマンス最適化以上のものが必要になることを受入れなければならない時です。

では… Apple Silicon 上で行列積乗算に最も速い CPU 命令は何でしょうか？そこが奇妙な立場に到達します：それは秘密です。

Apple Silicon は AMX（Apple Matrix Coprocessor）と呼ばれるユニットを含んでいます（Intel の AMX とは無関係）。これがこのユニットの正式名称かはわからないですが、Apple はこれをこれまで公式に「機械学習アクセラレーター」としか呼んできたことはありません。AMX 命令にアクセスする唯一の公開方法は

Accelerate

フレームワークの BLAS 実装経由ですが、ライブラリを使用することはこの記事のフレームワークフリー実装の前提と衝突します。幸運にも、人々は AMX ユニットがどのように機能するかをリバースエンジニアリングしたので、いくつかの命令に基づいて、より高速な行列積乗算実装を書くことができるか見てみましょう。

ここで楽しみつつも明確にするべきです：Apple の AMX 命令を直接使用しないでください。自分のアプリには

Accelerate

フレームワークを通してください。Apple はこれを「非公式」に保ち、いつでもバイナリ互換性を崩すことができます。Accelerate フレームワークは引き続き動作しますが、このコードは失敗します。そして次の記事のためのスプویلヤー：Apple の実装はどうしても私のものより約20%速いです。

必要な核心命令は

AMX_MATFP

です。この命令は16要素ベクトルの各要素を別の16要素ベクトルの各要素で乗算し、16x16 タイルを生成して出力タイルに累積します。これを適切な入力ベクトルで16回繰り返すと全体の 16x16 マトリクスを乗算できます。入力をロードするためには

AMX_LDX

と

AMX_LDY

も必要で、最後に

AMX_STZ

を使用して累積した 16x16 タイルをエミュートします。さらに

AMX_LDZ

（累積タイルを定義されたゼロタイルにリセットするために使用）を使って、タイル化行列積のインナー ループを処理できます。

つまりこのアルゴリズムは全体ループを示していません：それは単にインナー ループであり、データを準備し結果を統合するための外ループが必要です。ここではインナー ループのみを示します。

private static func amxF32_16x64(
    outTiles: UnsafeMutablePointer<Float>,
    lhsPanel: UnsafePointer<Float>,
    rhsPanels: UnsafePointer<Float>,
    innerCount: Int
) {
    zeroTileRow.withUnsafeBufferPointer { zeroBuffer in
        guard let zeroBase = zeroBuffer.baseAddress else { return }

        for tile in 0..<accumulatorCount {
            for row in 0..<tileRows {
                amx_ldz(zeroBase.amxZOperand(row: UInt32(tile + (row * accumulatorCount))))
            }
        }

        for k in 0..<innerCount {
            let lhsBase = lhsPanel + (k * tileRows)
            amx_ldx(lhsBase.amxXYOperand)
            for tile in 0..<accumulatorCount {
                let rhsBase = rhsPanels + (tile * innerCount * tileRows) + (k * tileRows)
                amx_ldy(rhsBase.amxXYOperand)
                amx_matfp(amxMatFPF32 | (UInt64(tile) << 20))
            }
        }

        for tile in 0..<accumulatorCount {
            let tileBase = outTiles + (tile * tileRows * tileRows)
            for row in 0..<tileRows {
                let rowBase = UnsafePointer<Float>(tileBase + (row * tileRows))
                amx_stz(rowBase.amxZOperand(row: UInt32(tile + (row * accumulatorCount))))
            }
        }
    }
}

再び、これは「Fast Swift」と呼んだ

LOOP_UNROLL

実装とは形状がそれほど異なるわけではありません。トレーニングではさらに1.67倍高速です。タILING はデータを必要なタイル形状にするために多くのパッキングと散列を必要とし、私はそこでより良い仕事をできると思います。これがより効率的であれば簡単に2倍以上も速いです。

最大電力消費

この実装の Metal コードは James Thompson の llm.metal から派生した実装に基づいています。ただし、彼らは行列積乗算にライブラリを使用したので、フレームワークフリーアプローチを続けるために独自の Metal 行列積乗算コードを書きました。

前節では「Apple Silicon 上で最も速い CPU 命令」と注意深く述べましたが、もちろん GPU もあります。Metal コードでの行列積乗算はどう見えますか？C と Swift コードとは異なり、2つの部分に分けられます：インナー カーネル（我々 が Metal/C++ で書く）と外側の呼び出しマシーン（Swift 側で残る）。

まず、インナー カーネル：

kernel void matmul_forward_kernel(
    device float* out [[buffer(0)]],
    const device float* inp [[buffer(1)]],
    const device float* weight [[buffer(2)]],
    const device float* bias [[buffer(3)]],
    constant uint& BT [[buffer(4)]],
    constant uint& C [[buffer(5)]],
    constant uint& OC [[buffer(6)]],
    uint2 gid [[thread_position_in_grid]]
) {
    uint oc = gid.x;
    uint bt = gid.y;
    if (bt >= BT || oc >= OC) {
        return;
    }

    float sum = bias[oc];
    for (uint i = 0; i < C; i++) {
        sum += inp[bt * C + i] * weight[oc * C + i];
    }
    out[bt * OC + oc] = sum;
}

やや重いパラメータブロックをスキップすると、これは実質的に C と Basic Swift の

matmul_forward

の4つのループのインナー ループのみが含まれていることがわかります。B, T, OC, C に対するループではなく、ここでは C に対するループのみです。B*T と OC をイテレートすることは外側のマシーンで処理されます：

func matmul_forward(out: MTLBuffer, inp: MTLBuffer, weight: MTLBuffer, bias: MTLBuffer, B: Int, T: Int, C: Int, OC: Int, ctx: MetalCommandContext) {
    ctx.compute.compute2D(
        pipelines.matmulForward,
        width: OC,
        height: B * T,
        threadsPerThreadgroup: MTLSize(width: 1, height: 1, depth: 1),
        buffers: [out, inp, weight, bias],
        uints: [UInt32(B * T), UInt32(C), UInt32(OC)]
    )
    ctx.endCompute()
}

ここで

compute2D

をヘルパーとして実装し、以下のように見えます：

func compute2D(
    _ pipeline: MTLComputePipelineState,
    width: Int,
    height: Int,
    threadsPerThreadgroup: MTLSize,
    buffers: [MTLBuffer],
    uints: [UInt32] = []
) {
    setComputePipelineState(pipeline)
    for (i, buf) in buffers.enumerated() {
        setBuffer(buf, offset: 0, index: i)
    }
    var uintsCopy = uints
    for i in 0..<uintsCopy.count {
        setBytes(&uintsCopy[i], length: MemoryLayout<UInt32>.stride, index: buffers.count + i)
    }
    dispatchThreads(
        MTLSize(width: width, height: height, depth: 1),
        threadsPerThreadgroup: threadsPerThreadgroup
    )
}

マルチスレッド化された Swift について、「スライスして同時に実行する操作があればもっと良かった」曾说过しました。まさにこれがそれを表しています：バッファをタイルセットに切り分けて、同時に処理できるようにすることです。ここでは幅と高さパラメータがカーネルコードの

gid.x

と

gid.y

パラメータに変換され、

inp

と

weight

バッファへのインデックスにほぼ同様にアクセスできます。

我々は

matmulForward

パイプラインを2次元の「スレッド」セット（CPU スレッドではなく GPU ワーカー）として呼び出していることがわかります。Basic Swift 実装の B と T ループを単一の B*T にフラット化すると想像すれば、文字通り同じ作業を行っており、単に電力を大量消費する GPU でなく CPU 上であります。

AMX 実装に対する大きな改善ではありません。トレーニングは少し速いですが推論は遅いです。高電力消費の高級 GPU の場合、大きな改善とは言えません。他のすべて同様、最も単純で素朴なアプローチでも自動的にトップクラスの性能を得るわけではありません。もちろん、GPU に「スレッド追加」は CPU よりもずっと簡単です。単に

threadsPerThreadgroup

パラメータを

MTLSize(width: 16, height: 16, depth: 1)

に変更するだけで、時間を容易な勝利に上げることができます（巨大な改善ではありませんが）。

我々が本当にやるべきことは、AMX ソリューションと同じようにマトリクスをタILING することです。つまり、インナー ループごとに長い行全体をトラバーサルするのではなく、よりローカルな方法でメモリにアクセスします。緊密なカーネルとタイルパッキング・散列を書くことを愛する人々いますが、それは私の強みではありません。ここに私が Metal タILING カーネルでの最善の試みがあります：

kernel void matmul_forward_kernel(
    device float* out [[buffer(0)]],
    const device float* inp [[buffer(1)]],
    const device float* weight [[buffer(2)]],
    constant uint& BT [[buffer(3)]],
    constant uint& C [[buffer(4)]],
    constant uint& OC [[buffer(5)]],
    uint2 gid [[thread_position_in_grid]],
    uint2 lid [[thread_position_in_threadgroup]]
) {
    threadgroup float inpTile[MATMUL_TILE][MATMUL_TILE];
    threadgroup float weightTile[MATMUL_TILE][MATMUL_TILE];

    uint oc = gid.x;
    uint bt = gid.y;
    bool inBounds = bt < BT && oc < OC;
    float sum = 0.0f;

    for (uint kBase = 0; kBase < C; kBase += MATMUL_TILE) {
        uint inpK = kBase + lid.x;
        uint weightK = kBase + lid.y;
        inpTile[lid.y][lid.x] = (bt < BT && inpK < C) ? inp[bt * C + inpK] : 0.0f;
        weightTile[lid.y][lid.x] = (oc < OC && weightK < C) ? weight[oc * C + weightK] : 0.0f;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint k = 0; k < MATMUL_TILE; k++) {
            sum += inpTile[lid.y][k] * weightTile[k][lid.x];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    if (inBounds) {
        out[bt * OC + oc] = sum;
    }
}

これまで見てきた最適化の試みと同様、これは読むのがずっと難しくなりますが、追加的な漸進的改善を提供します。これが私がこれで行く最上端なので、完全な表を見てみましょう：

我々がついに1 TFLOP/s の計算性能を超えました。1.1 Tflop/s は良いでしょうか？理論的には私の M3 Max GPU は約 15 Tflop/s の能力があります。しかし、この種類のタスクの実際の天井は数ある理由から 3-5 Tflop/s です—もちろん私が手書きした Metal マトリクスコードで得られる性能よりはるかに多いですが、桁違いにはありません。

ここで興味深い質問があります：私は高級 GPU を持っていますが、ベーシック M シリーズチップでは AMX ユニットと GPU の間の性能比率は変わるでしょうか？

テスト環境

CwlLlmSwift で TinyShakespeare トレーニングセット上のすべてのエンジン間の比較を示しています。

コードとテスト環境アプリは CwlLlmSwift からダウンロードできます。Andrej Karpathy の llm.c で使用される

llmc-starter-pack

もダウンロードする必要があります。

警告: 私は常にアプリを「Release」設定で実行するようすでに推奨しましたが、たとえ「Release」でも、「Basic Swift」の20回に及ぶトレーニング比較を実行するには約30分がかかります（そのスクリーンショット上の「Total」時間を参照してください）。大部分の時間は「Basic Swift」エンジンを無効にして他のものを見るだけで良いでしょう。

アプリを初めて起動すると、オープンダイアログが表示されますが、「New document」ボタンを押す必要があります。これは空のドキュメントを表示し、以下のドキュメントを選択する必要があることを示します：

• gpt2_124M.bin (チェックポイント)
• tiny_shakespeare_train.bin (トレーニングデータ)
• tiny_shakespeare_val.bin (バリデーションデータ)
• gpt2_tokenizer.bin (トークナイザー)

その後「Create」ボタンを押してデータセットを作成します。データセットが作成されると、トレーニングや推論、比較を実行し、データセットをドキュメントとして保存できます。

明確にするために：実際の GPT-2 スタイルの LLM はかなり中程度の性能です（1.2億パラメータのモデルで、多くの LLM は数十億パラメータから始まります）。これは

gpt2_124M.bin

チェックポイントでのモデルの例（どのトレーニングも実行していない前）です：

プロンプト： The quick brown fox
生成： would register the decision to slowdown his steps as he held out his right paw to hold up place. As Cas mornarye yelped as the fox clutched at his paw his temper ground up. "You sure you're not scared of birds?"

そして Shakespeare 上500回に及ぶポストトレーニング後：

プロンプト： The quick brown fox
生成： charged, Barking as he did, To fight withanna in his bosom, On one occasion dying: No fault is his life, for there blood being out. 'tis I; I did mean death, though it be by the water.

両方の例は遠くから見るとそれっぽいが、詳細にはやや不整合です。全く AGI ではありません。比較のため、ここで Apple の Foundation モデル「LanguageModelSession」（iOS/macOS に組み込まれ、唯一必要な「トレーニング」は与えるプロンプトのみ）から得られる結果を示します：

プロンプト： Turn the prompt "The quick brown fox" into a few lines of verse in Shakespearean language and meter.
生成： Upon the heath, a swift fox doth roam, His coat of brown, like autumn's golden hue, With nimble steps, he darts through the foam, A creature swift, both bold and true.

このモデルはイアンビックペンタメーターを正確に捉えていないと思われますが、少なくとも整合性があります。macOS でのすぐに利用可能な迅速な LLM を必要とする場合は、Apple の Foundation モデルを使用してください。しかし、このシリーズの目的は最高のモデルを作成する方法を示すことではなく、モデルを実行しポストトレーニングを実行する方法を探ることだけです。

結論

この記事はコードを最適化する方法を見て、機械学習ライブラリの行列積乗算を最適化するのに必要な作業への理解を深めることでした。これらのコードの一部も生産用途での良い選択とは考慮しないでください。次の記事では、macOS に組み込まれている BLAS, BNNS, CoreML, MPSGraph および他の高性能ライブラリについて見ていきます。これらのライブラリは、ここで示した選択肢よりもすべて効率的またはパフォーマンスが良いです。

基本的な Swift 実装の行列積乗算を取り、2.8 Gflop/s から 1.1 Tflop/s に変えるまで382倍も高速化しました。その過程で、最適化されていないコードから、少し改善されたコード、SIMD 最適化されたコード、マルチスレッド化されたコード、専用ユニット上で動作するコード、GPU コンピューティングシェーダーに至るスケール能力を確認しました。

最初の72倍の増加は、すべてのアルゴリズムで考慮すべき基本的な Swift 最適化によるものです：

• アレイのコピー-on-rite/参照チェックオーバーヘッドを回避する（Swift がこれを自動的に適用するはず；それは単なるバグかもしれません）
• 「Relaxed」を使用することで制約を緩め、Numerics ライブラリで SIMD 最適化された fused-multiply-add 命令を使用する
• トラバーサルがより効率的になり SIMD パイプライン動作がより良く機能するようにループを再構築する

• DispatchQueue.concurrentPerform

  を使用して並列化する

いつも通り、高性能 Swift は曲球を投げることが多々ありますが、Swift を C と同様に速くするのは常に可能です。今回、Swift がわずかに優れた SIMD 命令を選択したため、Swift が C よりも速い結果となりました（C も同じ FMA の SIMD バージョンを使用する方法はあるでしょうが）。問題は Swift が速いことができるかどうかではなく、むしろ：そこに到達したとき Swift はより良く見えるでしょうか？Swift はマルチスレッド化された実装までよく耐えました。

withUnsafeMutableBuffer

を使用すると任意の関数の審美性を台無しにします。Swift にアレイのサブレンジ上で変異を実行するために優雅で安全かつ良好な性能を持つ関数があると願います。

通常の Swift 最適化を超えて、AMX 命令を使用すると、より大きな行列操作のために設計された Apple Silicon の演算によりさらに70%も速くなりました。最終的なブーストはグラフィックカードに移動することでした。両方のケースで、タイル化マトリクス操作に熟練した開発者からより良い結果が得られます。私の努力は「まあ」でしたが、本当に何をすべきかを知っていれば性能は30%向上しても不思議ではありません。しかし382倍も速いというまだ十分ではありません。たとえこの例でも、秒あたり12トークン未満—これはいくつかのケースでわずかに受け入れられる速度かもしれませんが、この小さいモデルにとって笑い程度の遅さです。

明らかに、さらに進める必要があります。Apple はニューラルネットワークトレーニングのための多くの他のフレームワークを持っており、このシリーズでの最終実装は、この記事内の大部分の実装が最初のものを出したよりも前に、20以上のトレーニングイテレーションを完了させます。