0%

Mac mini M4 上实测 Apple Neural Engine 训练 Transformer

最近在 Mac mini M4 上跑了一下 maderix/ANE。这个项目的意思很直接:不用 CoreML 训练接口,不走 Metal,也不靠 GPU,而是通过逆向出来的私有接口,把计算图直接送到 Apple Neural Engine 上执行。

这事儿听起来很酷,但我一开始跑完以后也有点困惑:能跑是能跑,跟我平时用 Ollama、Dify、本地知识库有什么关系?折腾完以后,我的判断是:它现在不是一个生产工具,更像一个能摸到硬件底层的学习样本。

先说结论

它对日常使用帮助不大,至少现在不能指望它让 Ollama 或 Dify 变快。

但它很适合用来理解几件事:

  • Apple Neural Engine 不是只能做推理,硬件本身可以参与训练;
  • INT8 在 ANE 上确实能换来接近 2 倍吞吐;
  • 片上缓存、内存带宽、算子拆分这些东西,不是书上的抽象概念,跑一次 benchmark 就能看到;
  • 对学习 Machine Learning Systems、LLM from scratch 这类内容,它比单纯看公式更有帮助。

环境

我的机器是 Mac mini M4,系统为 macOS 15+。项目 README 里也写明作者是在 M4 Mac mini 上测试的,所以硬件上没有绕路。

仓库克隆后,先编译几个测试程序:

1
2
3
4
5
6
7
8
cd ~/projects
git clone https://github.com/maderix/ANE.git
cd ANE

clang -fobjc-arc -framework Foundation -framework CoreML -framework IOSurface -o inmem_peak inmem_peak.m
clang -fobjc-arc -framework Foundation -framework CoreML -framework IOSurface -o ane_int8_bench ane_int8_bench.m
clang -fobjc-arc -framework Foundation -framework CoreML -framework IOSurface -o inmem_bench inmem_bench.m
clang -fobjc-arc -framework Foundation -framework CoreML -framework IOSurface -o sram_bench sram_bench.m

FP16 峰值测试

先跑 inmem_peak

1
./inmem_peak

部分结果如下:

配置 权重 单次耗时 实测算力
32x conv 512ch sp64 16 MB 0.158 ms 6.78 TFLOPS
64x conv 512ch sp64 32 MB 0.208 ms 10.34 TFLOPS
96x conv 512ch sp64 48 MB 0.272 ms 11.83 TFLOPS
128x conv 512ch sp64 64 MB 0.333 ms 12.90 TFLOPS

M4 ANE 的 FP16 理论峰值大约 15.8 TFLOPS。这个测试里最高跑到 12.9 TFLOPS,已经能看出 ANE 不是摆设。

INT8 量化测试

然后跑 ane_int8_bench

1
./ane_int8_bench

结果更有意思:

配置 FP16 INT8 W8A8 加速比
128x conv 512ch 64x64 18.55 TOPS 34.20 TOPS 1.84x
64x conv 512ch 64x64 17.95 TOPS 33.24 TOPS 1.85x
256x conv 256ch 64x64 16.80 TOPS 31.19 TOPS 1.86x
128x conv 256ch 64x64 16.71 TOPS 30.39 TOPS 1.82x
128x conv 384ch 64x64 17.76 TOPS 32.84 TOPS 1.85x

这个结果比较直观。权重从 FP16 变成 INT8,内存占用差不多减半,吞吐接近翻倍。以前看量化,总觉得是模型部署里的一个技巧;在 ANE 上实测一遍以后,能更清楚地看到它为什么重要。

SRAM 和内存带宽

sram_bench 用来观察权重和激活变大以后性能怎么掉:

1
./sram_bench

部分结果:

配置 总数据量 单次耗时 实测算力
2048ch x 64sp 8.5 MB 0.229 ms 2.35 TFLOPS
3072ch x 64sp 18.8 MB 0.392 ms 3.09 TFLOPS
4096ch x 64sp 33.0 MB 0.959 ms 2.24 TFLOPS
8192ch x 32sp 129.0 MB 3.672 ms 1.17 TFLOPS

这里能看到一个性能断崖。数据规模上来以后,片上缓存装不下,计算开始受内存访问拖累。机器学习系统里常说“瓶颈不一定在算力,常常在内存”,这个测试就是很好的例子。

真正在 ANE 上训练 Transformer

项目里带了动态训练管道。我用 Stories110M 跑了一次完整训练。

先下载 TinyStories 预分词数据:

1
2
cd ~/projects/ANE/training
bash download_data.sh

下载完成后得到:

1
2
3
tinystories_data00.bin
20658981 tokens
约 41.3 MB

然后编译训练程序:

1
2
cd ~/projects/ANE/training/training_dynamic
make MODEL=stories110m

开始训练:

1
./train --scratch

最后完整跑完 10000 步:

1
2
3
4
Total steps:  10000
Compile: 590ms
Train time: 1087826ms
Wall time: 1539.5s

训练中后段的典型日志:

1
2
step 9990 loss=2.4622  lr=3.00e-05  104.0ms/step
timing: ane_fwd=14ms 左右,ane_bwd=24ms 左右

也就是说,109M 参数级别的 Transformer,在这套实验代码里能以大约 100ms 一步的速度跑完训练。这里的重点不是 loss 多少,也不是模型能不能拿来生成好文本,而是前向和反向传播确实跑到了 ANE 上。

它现在不能解决什么

需要冷静一点看。

它现在不能直接替代:

  • Ollama;
  • MLX;
  • llama.cpp;
  • CoreML;
  • Dify 的本地模型推理后端。

它也不是给普通用户训练大模型的工具。很多算子还在 CPU 上,ANE 利用率也没有完全打满。作者自己也说得很清楚:这是研究项目,不是生产框架。

所以,如果目标是“让本地大模型马上更快”,这个项目暂时帮不上太多。

它真正有用的地方

我觉得它最有价值的地方,是把一些平时很抽象的词变成了可以观察的东西。

比如:

  • 算子怎么切分;
  • 前向、反向分别耗时多少;
  • CPU、ANE、Accelerate 怎么配合;
  • 量化为什么会提高吞吐;
  • 片上缓存为什么会影响性能;
  • 编译图和运行时之间有什么关系。

这些内容在 Machine Learning Systems 里都会出现。单看书容易停留在概念层,跑一次这样的项目,印象会深很多。

我的判断

这个项目现在对日常工作没有直接帮助。它不会让我的知识库问答更快,也不会让 Ollama 自动用上 ANE。

但它值得保留。原因很简单:它是一个能跑起来的本地 AI 硬件实验样本。以后学到量化、算子融合、缓存、NPU、模型编译这些内容时,可以回来对照代码和日志。

对我来说,这次实验最大的收获不是“我训练了一个模型”,而是确认了一件事:Mac mini M4 里的 ANE 不是一个黑盒宣传词,它有实际计算能力,也有很现实的工程限制。理解这些限制,比单纯追参数、追模型名字更有用。