如何用 MLX 在苹果芯片上跑通推理与训练：3 条命令上手指南

如何用 MLX 在苹果芯片上跑通推理与训练：3 条命令上手指南

【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx

MLX 是苹果机器学习研究团队为 Apple Silicon 芯片打造的机器学习数组框架，Python API 与 NumPy 高度一致，支持自动微分、向量化和计算图优化。这篇文章用一条命令装好环境，讲清延迟计算与统一内存这两个核心机制，然后跑通一个完整的线性回归训练循环，最后演示如何用 Metal 调试器定位性能瓶颈。

为什么是 MLX：数据不再需要"搬家"

这一节回答"它和别的框架差在哪"，关键在内存模型的差异。

苹果芯片的 CPU 和 GPU 直接访问同一个内存池。多数框架要求你先把数组挪到 GPU 再计算，MLX 的数组从诞生起就住在统一内存里，任何设备都能直接操作它。这相当于两个人共用同一块案板，谁接手谁接着切，不用把菜打包端来端去。

第二个特点是延迟计算：a + b不会立刻算，只记录"一步菜谱"，等你调用mx.eval、打印或转 NumPy 时才真正下锅。类似面馆取号点单——先把单子记下来，你说"现在要吃"厨房才开工。好处是整张菜谱还没做，就可以对它做变换（自动微分、向量化），这为后面的mx.grad、mx.compile打下了基础。

3 条命令装好环境并验证

这一节解决"30 秒内有一个能跑的环境"的问题。

第一条，安装（macOS）：

pip install mlx

如果要在 Linux 机器上跑，按硬件加后缀：pip install mlx[cuda]或pip install mlx[cpu]。

第二条，验证安装，顺便看延迟计算的实际行为：

import mlx.core as mx a = mx.array([1, 2, 3, 4]) b = mx.array([1.0, 2.0, 3.0, 4.0]) c = a + b # 此时只记录"菜谱"，没有计算 mx.eval(c) # 到这里才真正执行 print(c) # array([2, 4, 6, 8], dtype=float32)

注意print那行：打印数组会自动触发求值，所以不用 eval 也能看到结果。NumPy 风格 API 的完整对照见仓库里的快速入门文档。

底层机制：调度器决定谁在什么时候算

这一节解释统一内存背后的分工规则：延迟计算管"记菜谱"，统一内存管"共享案板"，调度器是中间的主管。

没有依赖时，CPU 和 GPU 上的操作可以并行执行；有依赖时，调度器会自动插入等待——像工地施工，不用你叮嘱，"钢筋完工才能浇筑"这个顺序是调度自己保证的。看一个故意把计算拆到两个设备的例子：

a = mx.random.uniform(shape=(4096, 512)) b = mx.random.uniform(shape=(512, 4)) x = mx.matmul(a, b, stream=mx.gpu) # 大矩阵乘，交给 GPU for _ in range(500): b = mx.exp(b, stream=mx.cpu) # 512×4 的小 exp，留在 CPU mx.eval(x, b)

大矩阵乘计算密集适合 GPU；而 512×4 的 exp 太小，在 GPU 上启动开销比计算本身还大。官方文档unified_memory.rst在 M1 Max 上实测：全程 GPU 跑要 2.8 ms，按上面的方式拆到两个设备约 1.4 ms，接近一倍。

上图是 MLX 多设备张量并行推理中的列-行切分，底层逻辑和上面一致：数据不动，把计算拆开分派。

实战：一个最小线性回归

这一节把"前向 → 求梯度 → 更新 → 保存"的完整循环跑一遍，代码是仓库 examples/python/linear_regression.py 的精简版。

先生成一份"有标准答案"的假数据，方便最后检验学得准不准：

import mlx.core as mx X = mx.random.normal((1000, 100)) w_star = mx.random.normal((100,)) # 真实权重 y = X @ w_star + 1e-2 * mx.random.normal((1000,)) # 带噪声的标签 w = 1e-2 * mx.random.normal((100,)) # 随机初始猜测 def loss_fn(w): return 0.5 * mx.mean(mx.square(X @ w - y)) grad_fn = mx.grad(loss_fn) # 一行搞定自动微分 for _ in range(10_000): w = w - 0.01 * grad_fn(w) mx.eval(w) # 每个外循环迭代求值一次 print(mx.sum(mx.square(w - w_star)).item()) # 越接近 0 越成功

mx.grad自动把loss_fn包成一个可求导版本，不用手推公式。循环结束后输出值越小，说明学到的w越接近w_star。

想把结果存下来，两行就够：

mx.savez("model", w=w) # 落盘为 model.npz w_loaded = mx.load("model.npz")

mx.load按扩展名自动识别.npy、.npz、.safetensors、.gguf几种格式，加载大模型权重文件时尤其有用。

⚠️ 避坑：隐式求值和 eval 放哪儿

这一节讲三个最容易"变慢"或"行为反直觉"的场景。

第一，很多操作会悄悄触发求值：打印数组、转 NumPy、调用.item()、用mx.save保存。如果在训练循环里mx.eval之前先print(loss)，会先只算前向，再补算反向，一步迭代被劈成两次。

第二，标量数组做控制流同样触发求值：

def fun(x): h, y = first_layer(x) if y > 0: # 对标量数组比较，这里发生一次求值 z = second_layer_a(h) else: z = second_layer_b(h) return z

这段代码能跑、也能参与求导，但如果这种分支出现在高频循环里，求值开销会滚雪球。

第三，eval 的位置：官方文档lazy_evaluation.rst的建议是每个外循环迭代放一次，而不是循环内每操作一次。每次求值有固定开销，图太大又占内存，几十到几千个操作一次求值都在合理区间。

🔍 再提速：compile 与 Metal 抓取剖析

这一节给两个工具：一个让代码跑更快，一个帮你找到慢在哪。

前者是mx.compile，它会合并、融合计算图，生成更小的代码。首次调用慢（建图、优化、编译），之后同形状、同类型的调用直接命中缓存；输入形状或类型变了会重编译，所以只包会被反复调用的函数：

def fun(x, y): return mx.exp(-x) + y compiled_fun = mx.compile(fun) compiled_fun(x, y) # 首次调用：完整编译 compiled_fun(x, y) # 命中缓存

后者是 Metal 调试工具链：先用CMAKE_ARGS="-DMLX_METAL_DEBUG=ON"构建（会记录 Metal 编译源码、给对象打标签），运行时设置MTL_CAPTURE_ENABLED=1，在代码里抓取 GPU 工作：

mx.metal.start_capture("mlx_trace.gputrace") for _ in range(10): mx.eval(mx.add(a, b)) mx.metal.stop_capture()

生成的.gputrace文件用 Xcode 打开，Dependencies 视图一览所有操作及其依赖关系：

如果你用 CMake 生成了 Xcode 工程，可以跳过落盘文件直接在 Xcode 里抓取：

往哪儿深入

这一节给一张入口地图，免得翻遍整个仓库。

• docs/src/usage/下：lazy_evaluation.rst（求值细节）、compile.rst（编译边界情况）、function_transforms.rst（grad/vmap/vjp 的组合）、unified_memory.rst（stream 与调度语义）；
• examples/下：python/logistic_regression.py、python/qqmm.py（量化矩阵乘）、cpp/tutorial.cpp（C++ API 教程）；
• benchmarks/下是各算子的官方基准脚本，自己写压测时可以参考其中的计时方式。

做多设备张量并行时，还有两张图值得对照：

下一步可以直接做：运行一次python examples/python/linear_regression.py，记下它输出的吞吐量数字；然后把你自己的代码里重复执行的计算函数用mx.compile包一层，对比前后耗时，这就是你第一次 MLX 调优实验。

【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx