Swift 训练大语言模型：速度从 2.8 Gflop/s 提升至 5.884 令牌/秒，超 200 倍增长！

背景故事

大约两年前，作者翻出 21 世纪初用 C++ 编写的图像识别器代码，想让其重新运行但失败。2024 年初，围绕 LLM 讨论多，但没人在 Mac 上用 Swift 训练神经网络。作者尝试 Python 库，觉得无法掌控。一个月后，Andrej Karpathy 发布 llm.c，作者将其重写为 Swift 代码，最初实现非常慢。

llm.c

机器学习是将模型权重应用于输入数据，计算误差梯度并更新权重。矩阵乘法在机器学习中占大量工作，作者关注 llm.c 中的 matmul_forward 函数，一次训练迭代约有 0.2 万亿次浮点运算。纯 C 代码在 Swift 包中运行，每 7 秒完成一次训练迭代，推理速度不到每秒 1 个令牌。

基础 Swift 实现

作者尽力让基础 Swift 版本与 C 版本保持一致，设置移除数组索引的运行时检查，在 Release 配置下运行。但 Swift 实现极其缓慢，比 C 代码慢 15 到 20 倍，性能约为 2.8 Gflop/s。

Span、Egg 和 Span

通过 Instruments 检查发现，性能开销最大的是 _ArrayBuffer.beginCOWMutation()。Swift 6.2 提供 MutableSpan 解决方案，使用后训练迭代速度提高 3 倍多。

轻松氛围

前向传播中最耗时的是 value += inp[bt * C + i] * weight[o * C + i]。C 有 -ffast-math 优化标志，可使用融合乘加指令。Swift 没有该标志，只能单独进行乘法和加法运算。使用 Swift-Numerics 的 Relaxed 实现快速数学运算后，每秒处理的令牌数几乎提高 10 倍，但训练性能仍比 C 慢 15%。

循环展开

C 语言中矩阵乘法实际函数会以 8 步为单位遍历外层循环，编译器实现 8 次循环展开。Swift 6.2 提供 InlineArray 特性，实现与 C 相似的循环展开，C 和 Swift 推理速度基本相同，Swift 训练速度略快。

多线程实现

llm.c 代码中有 OpenMP 注释，但 Swift 包管理器使用的常规 clang 编译输出中不会生效。Swift 没有简单方法标记循环进行并行处理，使用 DispatchQueue.concurrentPerform 进行并发处理。将此模式应用于训练迭代中最耗时的四个循环后，速度提高 5.4 倍。

绝密技巧

Apple Silicon 包含 AMX 单元，苹果未正式命名，访问其指令的唯一公开方式是通过 Accelerate 框架。有人逆向工程了 AMX 单元工作原理，作者基于此编写更快的矩阵乘法实现，训练速度又提高 1.67 倍。

最大功耗方案

用于矩阵乘法的 Metal 代码分为内层内核（用 Metal/C++ 编写）和外层调用机制（在 Swift 端）。