用Swift训练LLM：从Gflop/s到Tflop/s的矩阵乘法优化实战教程

想在 Mac 上训练大语言模型？一位开发者用纯Swift（不用任何ML框架）实现了GPT-2的训练，并通过10种矩阵乘法优化方案，将性能从0.054 tokens/s提升到11.1 tokens/s，最终在M3 Max上达到了1.1 Tflop/s的吞吐量。这篇文章记录了他的优化路径。

背景：为什么是矩阵乘法？

LLM训练的核心计算瓶颈是矩阵乘法（z += x * y）。作者 Matt Gallagher 以 Andrej Karpathy 的 llm.c 为参考基准，在 Apple Silicon 上用Swift从零实现GPT-2 124M参数模型的训练，每次训练迭代约0.2万亿次浮点运算。

项目代码：CwlLlmSwift

10种实现方案与性能对比

以下是作者在M3 Max上测试的10种方案，性能从慢到快排列：

第1版：基础Swift（0.054 tokens/s）

最朴素的三重循环实现。Swift的Array有copy-on-write开销，每次切片操作都会触发复制，导致性能只有C语言的7.3%。

第2版：+MutableSpan（0.056 tokens/s）

Swift 6.2引入的 MutableSpan 消除了Array的COW开销，性能提升到C的24%。

第3版：+Relaxed FMA（0.533 tokens/s）

使用Swift-Numerics库的 Relaxed.multiplyAdd，等效于C的 -ffast-math，允许编译器使用FMA（融合乘加）指令。性能飙升到C的85%。

第4版：循环展开+InlineArray（0.918 tokens/s）

Swift 6.2的 InlineArray 是栈分配的定长数组，配合循环展开，终于追平C语言（106.6%）。

第5版：多线程Swift（4.356 tokens/s）

用 DispatchQueue.concurrentPerform 做多线程并行。性能达到C+OpenMP的558%，但代码比C版本更难看。

第6版：AMX指令集（5.884 tokens/s）

逆向工程Apple的矩阵协处理器（AMX_MATFP、AMX_LDX等未公开指令），直接调用硬件加速。不过作者指出，生产环境应该用Apple的Accelerate框架，它比自己写的AMX代码还快约20%。

第7-10版：Metal GPU方案（4.29→11.123 tokens/s）

最终四版都用Metal计算着色器在GPU上运行：

• 基础Metal：4.29 tokens/s
• 多线程Metal：6.211 tokens/s
• 分块Metal（Tiled）：11.123 tokens/s（最终版）

最终版在M3 Max上达到了约1.1 Tflop/s，是基础Swift的382倍。

关键优化技巧总结

1. 消除COW开销：Swift的Array默认copy-on-write，用MutableSpan或InlineArray替代
2. 启用FMA指令：Relaxed.multiplyAdd允许编译器融合乘加操作
3. 栈分配数组：InlineArray避免堆分配和GC压力
4. 多线程并行：DispatchQueue.concurrentPerform利用所有CPU核心
5. 硬件加速：Apple的AMX协处理器或Metal GPU
6. 分块计算：Tiled矩阵乘法更好地利用GPU缓存

实测数据

适合谁？

• 想在Mac上做AI/ML开发的Swift开发者
• 对GPU编程和性能优化感兴趣的人
• 想了解Apple Silicon底层硬件能力的技术爱好者
• 正在评估Apple Silicon做AI训练可行性的团队

注意事项

• AMX指令集是Apple未公开的，生产环境请用Accelerate框架
• M3 Max GPU理论峰值约15 Tflop/s，实际可达到3-5 Tflop/s
• 这个项目主要用于学习，实际训练大模型还是建议用PyTorch/JAX等框架

来源：
Cocoa with Love: Training an LLM in Swift
CwlLlmSwift GitHub仓库