8.0
热点
AI SCORE
技术实践2026-05-11 01:05
Swift 训练 LLM：从 Gflop/s 到 Tflop/s
Hacker News · cocoawithlove.com#Swift#性能优化#机器学习
Editor brief · 编辑速览
详细展示矩阵乘法优化如何实现 1000 倍性能提升。对 ML 工程师和性能优化者有深度参考价值。
在这篇文章中，我尝试让自己手写的矩阵乘法代码在 Swift 中运行得尽可能快，用于训练大语言模型 (LLM)。目的是深入了解在 Swift 中优化数学代码的关键步骤。我还希望这些例子能给出一个规模感，展现 Apple Silicon 上不同计算单元的能力——CPU、SIMD、AMX 和 GPU。
这将是我系列文章的第一篇，探讨在 Apple Silicon 上用 Swift 训练神经网络。后续文章将关注苹果为 Mac 上的机器学习提供的众多框架。那些经验丰富的框架才是你真正应该用于矩阵乘法和机器学习的工具（相比之下，他们花了更多年去优化矩阵乘法核心）。
但在那之前，我正在享受用"无框架、无库"的纯代码方式自己写一切的乐趣。
而且我不仅在写矩阵乘法核心。示例应用会将这些核心用作完整 LLM 实现的一部分，我引用的数字将是整个前向和后向训练迭代的结果。这个系列的参考实现是 Andrej Karpathy 的 llm.c（一个纯 C 实现的 GPT2 兼容模型）。这是一个相当基础的模型，但包含了所有必要的组件，代表了真实世界的工作负载。
这意味着是时候进行我最喜欢的游戏了：优化 Swift 让它比 C 更快。
大约两年前，我翻出了我 2000 年代初的工程学位论文。这是一个用 C++ 编写的图像识别器，使用神经网络来分类图像。我想让我的旧代码重新运行，但我已经很久没做过机器学习代码了。结果变得很烦人，我放弃了。
在 2024 年初关于 LLM 的所有讨论中，感觉没有人在 Mac 上训练神经网络。至少不是用 Swift 这样的语言。我尝试了一些 Python 库，如 PyTorch 和 TensorFlow，但 Python 本身从不执行计算——它更像是底层另一个计算引擎的编排器——这种分离让我感觉不到自己有掌控力。
一个月后，Andrej Karpathy 发布了 llm.c。这以一种其他机器学习内容没有的方式吸引了我，因为没有什么是隐藏的。它大约 1000 行纯 C 代码，虽然充满了一些相当神秘的变量名，但相对来说是可读的。
所以当然，我立刻用 Swift 重写了它。玩这些代码非常有趣。
当然，玩代码需要一些工作来让它快速运行。有个剧透在这里：最初的 Swift 实现非常慢。但优化是一个持续的过程：总有更多东西可以尝试。
这终于引出了这篇文章：我将走过我当时写的不同探索（以及过去一周添加的几个），让 LLM 训练相当快速，而不需要借助库。大部分代码将是 Swift（虽然我最后会展示一个 Metal 实现）。
顺便说一下，我不会解释神经网络或 LLM 如何工作。如果你感兴趣，Karpathy 的视频《Let's build GPT: from scratch, in code》基本上是学习 GPT 类 LLM 如何工作的决定性指南，他之前的系列从《The spelled-out intro to language modeling: building makemore》开始，如果你想要更入门的课程，5 个视频系列涵盖了许多介绍性概念。当然，两者都是用 Python 写的，所以当你准备好看我们如何在 Swift 中做这些事情时，请回到这里。
机器学习本质上是将模型权重应用于输入数据（称为前向传播，也就是推理），然后计算误差梯度和那些权重的更新（后向传播）。
我们通常会把这些计算打包在一起，并尽可能快地让它们运行。这些操作包可能被称为："线性张量投影"、"矩阵乘法"，甚至是一系列"向量点积"（取决于你如何划分工作单元的大小）。本质上都是一个执行 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;
}
}
}
}
这四层嵌套循环增加了一些视觉上的复杂性，但实际上，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 版本尽可能接近 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 个 token。在这个引擎上运行 20 次训练迭代需要近 30 分钟。到底发生了什么？
这个性能代表大约 2.8 GFLOP/s。1999 年苹果为 PowerMax G4 做广告，声称其 1 GFLOP/s 的能力使其在美国军方眼中成为一种武器。现在 2.8 GFLOP/s 是完全无法接受的。
在 Instruments 中检查，之前运行中迄今为止最大的性能成本是 _ArrayBuffer.beginCOWMutation()。Swift 认为其他人可能在使用我们的数组，即使它们是唯一的（所以我们没有得到数组副本），仅有的唯一性检查就是我们最大的开销。
什么？有时你会遇到可能只是一个 bug 的问题。这可能就是其中之一。当我在 2024 年第一次处理这些代码时，我的记忆是这不是一个问题。我不知道是否有性能回退，或者安全漏洞被修复导致 _ArrayBuffer.beginCOWMutation() 堵住了性能。当我用 inline(none) 禁用函数内联时，这个问题也会移动，所以感觉优化器就是无法正常工作。
无论如何，我们不能使用数组并获得我们需要的性能。幸运的是，Swift 6.2 给了我们一个基本没有开销的可靠修复：MutableSpan。
我只是在顶部添加了 var out = out.mutableSpan 这一行，用自己的 mutableSpan 遮蔽了 out。我也在整个文件中应用了相同的模式。
有趣的是，虽然这个改变对前向传播的影响不大，但它确实使训练迭代（前向 + 反向 + 更新）快了 3 倍多。
但我们需要把注意力转向为什么前向传播很慢。Instruments 证实了我们已经知道的：前向传播中最热点的一行是循环中心的 value += inp[bt * C + i] * weight[o * C + i]。
是时候面对一个难以接受的事实了：C 有一些 Swift 没有的编译器优化标志。在这个具体情况下，C 有 -ffast-math，它允许 C 使用融合乘加（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 内循环主要只是 8 个展开的 fmadd（融合乘加指令）应用。
有趣的事实：我们生活在未来，所以如果像汇编语言这样的东西对你没有意义，你可以把它们放入你最喜欢的 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 所做的那样。
幸运的是，我们有 Swift-Numerics，它为我们自己的快速数学版本提供了 Relaxed。就像《银河系漫游指南》中的"不要惊慌"一样，我相信"Relax"这个词只是在这里让我们都变得更加平静，享受放松的氛围。或者它让我们放松关于舍入结果的规则，以便可以进行 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 选择了 fmadd 的 SIMD 向量化版本（fmla），但前提是相同的。
这是每秒令牌数的近乎 10 倍的速度提升。但我们的训练性能仍然比 C 慢 15%。我们如何缩小最后的差距？
是时候承认一些事情了：我一直在展示来自 C 的"朴素"版本的矩阵乘法。真正的 C 函数稍微有点丑，因为它一次跨越外循环 8 步，希望编译器会展开最内层的循环。而 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 在这里为我们提供了另一个有用的功能：InlineArray，它终于与 C 栈分配的数组相匹配。
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（开放多处理）的。OpenMP 通常需要修改后的 C 编译器，但这是 llm.c 代码的预期运行方式。
我认为这是 Swift 在可读性方面最终输给 C 的地方。不仅没有简单的方法来标记一个循环以进行并行化，而且没有 Swift 6 安全的方法来切片数组并并发处理单独的切片。MutableSpan 之类的方法会抱怨你正在并发访问变量（即使它们是由单独的切片组成的）。
这意味着我们所有的可变数组都需要采用 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) / chunkSize
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]
}
}
}
}
}
这不是字面意义上最糟糕的代码，但是：
out.withUnsafeMutableBufferPointer 作用域和 DispatchQueue.concurrentPerform 作用域肯定会增加代码的视觉混乱度在这一步，我想：C 代码看起来更好。如果能在 RangeReplaceableCollection 上有一个"切片并并发执行"的操作就好了，这样就可以像替换 for 循环一样简洁。
我已经将这个模式应用到了四个最热的训练循环中：matmul_forward、matmul_backward、attention_forward 和 attention_backward。
得到了 5.4 倍的性能提升。相当不错，虽然我的 CPU 有 16 个核心，但仅仅 5 倍的多线程改进不算完美利用率。我们现在可能受到内存遍历的限制。
我看到的最大问题是，在循环展开、不安全指针、未检查的 sendable 包装器和跟踪并发执行迭代之间，代码现在已经如此繁琐，很难操作。
我认为我们已经做得很好了。但现在是时候接受现实了：如果我们想要 Mac 上最快的基于 CPU 的矩阵乘法，我们需要的不仅仅是 Relaxed.multiplyAdd 和底层 fmla.4s 指令所能提供的基本 SIMD 性能优化。
那么……Apple Silicon 上最快的 CPU 指令是什么呢？
这就是我们到达的奇特位置：这是个秘密。
Apple Silicon 包含一个叫 AMX（Apple Matrix Coprocessor，与英特尔的 AMX 无关）的单元。我不确定这是否是这个单元的官方名称，因为 Apple 从未在官方场合把它称为除了"机器学习加速器"以外的任何东西。唯一的公开访问 AMX 指令的方式是通过 Accelerate 框架的 BLAS 实现，但使用库与本文的前提（无框架实现）相冲突。
幸运的是，人们已经逆向工程了 AMX 单元的工作方式，所以基于其中的一些指令，让我们看看是否能写出更快的矩阵乘法实现。
我在这里玩得很开心，但应该很清楚：不要直接使用 Apple 的 AMX 指令。在你自己的应用中使用 Accelerate 框架。Apple 将其保持为"未文档化"状态，所以他们可以随时破坏二进制兼容性。Accelerate 框架将继续工作，但这段代码会失败。而且为下一篇文章剧透一下：Apple 的实现大约比我的快 20%。
我们需要的核心指令是 AMX_MATFP。这条指令将一个 16 元素向量中的每个元素乘以另一个 16 元素向量中的每个元素，生成一个 16 × 16 的瓦片，该瓦片在输出瓦片中累积。用适当的输入向量做这个 16 次可以乘以整个 16 × 16 矩阵。我们需要加载输入，所以我们还有 AMX_LDX 和 AMX_LDY，在最后，我们使用 AMX_STZ 发出 16 × 16 的累积瓦片。加上 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))))
}
}
}
}
再一次，它的形状与我称之为"快速 Swift"的 LOOP_UNROLL 实现没有太大区别。
在训练中又快了 1.67 倍。分瓦片需要大量的数据打包和分散，以将我们的行优先矩阵转换为所需的