用 Swift 训练大语言模型，第一部分：将矩阵乘法从 Gflop/s 提升到 Tflop/s

# 在 Swift 中训练 LLM，第 1 部分：将矩阵乘法从 Gflop/s 提升到 Tflop/s | Cocoa with Love 来源：https://www.cocoawithlove.com/blog/matrix-multiplications-swift.html 在本文中，我将尝试让手写的矩阵乘法代码在 Swift 中以最快的速度运行，以便在 Apple Silicon 上训练一个大语言模型（LLM）。目的是为了深入了解优化数学代码的关键步骤。我也希望这些示例能让读者对 Apple Silicon 上不同单元（CPU、SIMD、AMX 和 GPU）的性能有一个直观的感受。 这是关于在 Apple Silicon 上用 Swift 训练神经网络系列的第一篇。后续文章会介绍苹果为 Mac 提供的机器学习框架——那些框架*可能有点多*。对于矩阵乘法和机器学习，你真正应该使用的是那些成熟的框架（它们在内核优化上花的时间比我多得多）。但在那之前，我正乐在其中地用“无框架、无库”的纯代码方式自己实现一切。我不只是写矩阵乘法内核，示例应用还会把这些内核用于一个完整的 LLM 实现，而我所引用的数字也将来自完整的正向和反向训练迭代。 本系列的参考实现是 Andrej Karpathy 的 [llm.c](https://github.com/karpathy/llm.c)（一个兼容 GPT2 模型的纯 C 实现）。这虽然是一个相当基础的模型，但它包含了所有必要的组件，并且能够代表真实的工作负载。这意味着，是时候来玩我最喜欢的游戏了：把 Swift 优化到比 C 还快。 ## 目录 - [背景故事](#backstory) - [llm.c](#llmc) - [基本 Swift](#basic-swift) - [Span、Egg 和 Span](#span-egg-and-span) - [放松氛围](#chill-vibes) - [他们看到我在滚动](#they-see-me-rollin) - [多线程](#multi-threaded) - [绝密黑客手法](#top-secret-hacks) - [最大功耗](#maximum-power-draw) - [测试工具](#test-harness) - [结语](#conclusion) ## 背景故事 {#backstory} 大约两年前，我翻出了自己 21 世纪初的工程毕业论文。那是一个用 C++ 编写的图像识别器，利用神经网络对图像进行分类。我本想重新跑起这些老代码，但已经很久没有碰过 ML 代码了。过程很令人恼火，最后我放弃了。 尽管 2024 年初大家都在讨论 LLM，但感觉没有人在 Mac 上*训练*神经网络，至少不用 Swift 这类语言。我试过 PyTorch 和 TensorFlow 等 Python 库，但 Python 本身从不亲自计算——它更像是幕后计算引擎的调度器——这种分离让我觉得自己失去了掌控。 一个月后，Andrej Karpathy 发布了 [llm.c](https://github.com/karpathy/llm.c)。它触动我的方式和其他机器学习内容不同，因为其中毫无隐藏。它大约 1000 行纯 C 代码，虽然变量名有些晦涩，但相对容易读懂。于是，很自然地，我立刻用 Swift 把它重写了一遍。玩起来非常有趣。 当然，要让代码跑得快，还需要做些工作。先剧透一下：最初的 Swift 实现真的超级慢。但优化是一个持续的过程：总有更多的东西可以尝试。这最终引出了这篇文章：我将逐一讲述我当时（以及上周补充的几个）所做的各种探索，目的是让 LLM 训练变得*相当快*，而**无需**依赖任何库。大部分代码会用 Swift 编写（最后会展示一个 Metal 实现）。 顺便说一句，**我不会解释神经网络或 LLM 是如何工作的**。如果你感兴趣，Karpathy 的视频 [Let’s build GPT: from scratch, in code, spelled out.](https://www.youtube.com/watch?v=kCc8FmEb1nY) 几乎是学习 GPT 类 LLM 工作原理的权威指南，而他更早的系列视频从 [The spelled-out intro to language modeling: building makemore](https://youtu.be/PaCmpygFfXo?si=SQmmakwYc2B3BfMv) 开始，用 5 个视频介绍了许多基础概念，如果你想要更入门的内容。当然，两者都是 Python，所以等你准备好了解如何在 Swift 中实现时，再回到这里。 ## llm.c {#llmc} 机器学习本质上是对输入数据应用模型权重（称为前向传播，又称推理），接着计算误差梯度并更新这些权重（反向传播）。我们通常把一系列计算打包在一起，并使其运行得尽可能快。这些操作包可能被称为“线性张量投影”、“矩阵乘法”，甚至是“向量点积”（取决于你拆分工作单元的大小）。归根结底，它是一个循环，多次执行 `z += x * y`。 由于矩阵乘法在机器学习中占如此大的比重，我将重点介绍执行此操作的代码。我会同步更新其他部分的实现，但只使用与矩阵乘法相同的优化方法。 先来看看 llm.c 中的 `matmul_forward`，这是前向传播中使用的核心矩阵乘法。它遍历输入（`inp`），乘以模型权重（`weight`），并将结果累加到求和变量（`val`）中。 ```c 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; } } } } ``` 四层循环增加了一些视觉复杂性，但实际上，`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¹¹ ≈ 0.2 万亿次浮点运算，每次训练迭代。因此必须运行得很快。 纯 C 代码很容易在 Swift Package 中运行。我已经将 C 实现固定为始终以 `-O3` 优化级别编译（不管 Xcode 设置如何）。即使在这个优化级别下，C 实现每次训练迭代也需 7 秒，推理速度不到每秒 1 个 token。这只能算一个优秀的概念验证，但比能用的程度慢了 10 倍。 ## 基本 Swift {#basic-swift} 我尽可能让基本 Swift 版本保持与 C 版本一致： ```swift 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 val = bias?[o] ?? 0.0 for i in 0..<C { val += inp[bt * C + i] * weight[o*C + i] } out[bt * OC + o] = val } } } } ``` **不要在 Debug 模式下运行。** 我只会引用 Release 配置下的运行数据。虽然我曾在 Debug 下运行过*部分*代码，但我从没等过完整的 20 轮训练在 Debug 下完成。我通常把 Xcode 的 Scheme 设置为“Release”——即使在调试时。 如果读过背景故事，我已经提过：这“极其缓慢”。Swift 代码慢了 15 到 20 倍。这意味着 LLM 每 19 秒才产生 1 个 token。在这个引擎上跑 20 个训练迭代需要近 30 分钟。究竟发生了什么？ > 这个性能大约为 2.8 Gflop/s。1999 年苹果为 PowerMac G4 做广告，声称其 1 Gflop/s 的能力使其成为美国军方眼中的武器（[链接](https://www.youtube.com/watch?v=OoxvLq0dFvw)）。如今 2.8 Gflop/s 完全不能接受。 ## Span、Egg 和 Span {#span-egg-and-span} 用 Instruments 检查后发现，上次运行中最大的性能耗费是 `_ArrayBuffer.beginCOWMutation()`。Swift 认为其他人可能也在使用我们的 `Array`，即便它们是唯一的（所以我们并没有发生数组拷贝），但仅唯一性检查本身就造成了最大的开销。 > **哈？** 有时你会遇到可能只是 bug 的问题。这可能就是一个。我在 2024 年第一次编写这份代码时，记得这并非问题。我不知道是否出现了性能回退，还是某个安全漏洞被堵上了，导致 `_ArrayBuffer.beginCOWMutation()` 阻塞了性能。当我用 `inline(none)` 禁用函数内联时，这个问题也会转移位置，感觉像是优化器无法正常工作。无论如何，我们不能用 `Array` 来获得所需性能。 幸运的是，Swift 6.2 提供了一个可靠的几乎零开销的解决方案：`MutableSpan`。 ```swift 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.0 for i in 0..<C { val += inp[bt * C + i] * weight[o*C + i] } out[bt * OC + o] = val } } } } ``` 代码几乎相同。唯一的区别是 `var out = out.mutableSpan`，它返回一个 `MutableSpan<Float>`，可以完全取代 `Array`。事实上，在许多情况下，`MutableSpan` 的行为类似 `Array`，但它并不提供写时复制（COW）语义，因此我们为回避该特性所做的所有工作都将获得回报。 这导致了 4.5 倍的性能提升。仍然只是 4.5 Gflop/s，但现在至少比 1999 年的武器还要快那么一点点了。 ## 放松氛围 {#chill-vibes} 有时你会看着反汇编代码，感觉自己仿佛走进了西部荒野。 **有趣的事实**：我们生活在未来，所以如果有些像汇编语言的东西让你摸不着头脑，你可以把它丢给你最喜欢的 LLM，它会为你解释清楚。Swift 里没有 `-ffast-math`，所以我们看到的是分开的乘法和加法指令。 ```asm +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](https://github.com/apple/swift-numerics)，它提供了 `Relaxed` 模块，让我们拥有自己的“快速数学”。就像《银河系漫游指南》里的“别慌”一样，我猜“Relax”这个词就是让我们更镇静，享受轻松的氛围。或者它允许我们放宽舍入结果的规则，从而使 FMA 成为可能。 所有 `a += b * c` 和 `x = y + z` 操作都可以用 `Relaxed.multiplyAdd` 和 `Relaxed.sum` 来改进，但我会特别避免在 `gelu_backward` 函数上应用这些函数，因为 C 实现明确在该函数周围禁用了 `-ffast-math`。 ```swift 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.0 for i in 0..<C { val = Relaxed.multiplyAdd(inp[bt * C + i], weight[o*C + i], val) } out[bt * OC + o] = val } } } } ``` 这带来了 1.8 倍的提升。运行效果开始像有了点动力。 ## 他们看到我在滚动 {#they-see-me-rollin} 我以前在 C 级性能对比中看到过这个把戏，但一直没能找到一种不用手动展开就能使其工作的方式。通常我们期望编译器能为我们进行循环展开，但检查汇编代码后，我发现它并没有。在 `Relaxed.multiplyAdd` 内部的跳跃指令（`b.ne`）是一个太大的开销。在 2024 年的实现中，我只能手动将循环展开 8 次（读起来相当丑陋）。不过，Swift 6.2 为我们提供了另一个有用的特性：`InlineArray`，终于匹配了 C 的栈分配数组。 ```swift for obt in stride(from: 0, to: BT, by: LOOP_UNROLL) { for o in 0..<OC { var vals = InlineArray<Float, LOOP_UNROLL>(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 idx in 0..<LOOP_UNROLL { vals[idx] = Relaxed.multiplyAdd(bt[obt + idx * C + i], w[i], vals[idx]) } } for idx in 0..<LOOP_UNROLL { out[obt + idx * OC + o] = vals[idx] } } } ``` 哦，我还把最外层的循环改成了以 `LOOP_UNROLL` 为步长的 `stride`，这样每次迭代计算多个 `obt`。 再获得 2 倍的提升。此时我们终于达到了与 C 相当的性能：每次训练迭代不到 7.5 秒。 ## 多线程 {#multi-threaded} 虽然我们赶上了 C 的单线程性能，但 C 在 llm.c 中是用 OpenMP 进行多线程的。我们也要多线程。Swift 有 GCD（Grand Central Dispatch），它提供的 `DispatchQueue.concurrentPerform` 非常适合这种利落小队式的工作。 ```swift DispatchQueue.concurrentPerform(iterations: threadCount) { threadIndex in // ... threadIndex 贴出了本次迭代处理的 ob 范围 } ``` 这几乎是一个线性的性能提升，乘以 14（我的 M1 Max 上有 14 个能用的核心吗？是 8 个性能核心，再加上其他一些……总之大约 14）。不错。再获得 14 倍的提升。现在比 C 快 1.15 倍。 ## 绝密黑客手法 {#top-secret-hacks} **注意**：我在这里玩得很开心，但应该清楚：**不要直接使用 Apple 的 AMX 指令。通过 Accelerate 框架来使用**，为你自己的应用着想。Apple 保持这些“未公开”的指令，是为了可以随时破坏二进制兼容性。Accelerate 框架会继续工作，但这段代码会失效。而且下一篇文章的剧透：Apple 的实现比我的大约快 20%。 我们需要的核心指令是 `AMX_MATFP`。该指令将一个 16 元素向量的每个元素与另一个 16 元素向量的每个元素相乘，生成一个 16 x 16 的小块，并将其累加到输出小块中。通过对适当的输入向量执行 16 次此操作，可以乘以整个 16 x 16 矩阵。我们需要加载输入，因此还有 `AMX_LDX`、`AMX_LDY`，最后我们用 `AMX_STZ` 发出累加后的 16 x 16 小块。配合 `AMX_LDZ`（用于将累加器块重置为预定义的零块），就能处理分块矩阵乘法的内层循环。 这意味着该算法并未展示整个循环：它只是内层循环，需要外层循环来准备数据并统一结果。我只展示内层循环： ```c 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 { amx_ldz(zeroBase.amxZOperand(row: UInt32(tile))) } for i in 0..<innerCount { for tile in 0..<accumulatorCount { amx_ldx(lhsPanel.amxXOperand(row: UInt32(tile * 16))) amx_ldy(rhsPanels.amxYOperand(row: UInt32(tile * 16))) amx_matfp(lhsPanel.amxXOperand(row: UInt32(tile * 16)), rhsPanels.amxYOperand(row: UInt32(tile * 16))) } } for tile in 0..<accumulatorCount { let tileBase = outTiles.advanced(by: tile * tileRows * tileCols) for row in 0..<tileRows { amx_stz(tileBase.amxZOperand(row: UInt32(tile + (row * accumulatorCount)))) } } } } ``` 再强调一次，它在形状上与那个被我称为“快速 Swift”的 `LOOP_UNROLL` 实现没有*那么*大的不同。训练速度又快了 1.67 倍。分块需要大量的数据打包和散布，才能将行主序矩阵转换成所需的分块形式，我确信自己在这方面可以做得更好。如果更高效，很容易能再快 2 倍以上。 ## 最大功耗 {#maximum-power-draw} > 本实现中的 Metal 代码**源自 James Thompson 在 [llm.metal](https://github.com/regrettable-username/llm.metal) 中的实现**。不过，他们使用了矩阵乘法库，所以我写了属于自己的 Metal 矩阵乘法代码，以保持无框架的方式。 在上一节中，我小心地表述为“Apple Silicon 上用于矩阵乘法的最快 *CPU* 指令”，因为当然，我们还有 GPU。Metal 代码中的矩阵乘法长什么样？与 C 和 Swift 代码不同，它分为两部分：内部内核（用 Metal/C++ 编写）和外部调用机制（留在 Swift 侧）。首先是内部内核： ```cpp 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 和基本 Swift 版 `matmul_forward` 四层循环中的最内层循环。这里我们不再循环 `B`、`T`、`OC` 和 `C`，而只保留了对 `C` 的循环。遍历 `B * T` 和 `OC` 的任务则交给了外部调用机制： ```swift 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,