如何用 MLX 在 Apple 芯片上训练与部署模型【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx在 Apple 芯片上跑深度学习时你多半遇到这种情况模型在 CPU 和 GPU 之间来回搬数据或者干脆只能在 Linux CUDA 环境里训练。MLX 是 Apple 机器学习研究团队推出的一个数组框架专为 Apple Silicon 优化让你用接近 NumPy 的语法完成模型训练、推理与部署。它的延迟计算、统一内存和可组合函数变换三个设计正是为在 Mac 上高效跑模型而做的。这篇笔记带你完成安装、跑通一个完整的训练加推理示例并给出几个常见坑的处理方式。环境与安装先确认设备满足要求以 macOS 版本为例macOS 14.0且为 Apple Silicon 芯片原生 Python 3.10架构为 arm非 RosettaLinux 版需要 glibc 2.35CUDA 版额外要求 CUDA 12.0、Nvidia 架构 SM 7.5最小安装就一条命令pip install mlx这样装是因为 PyPI 包已经编译好 Metal 后端无需你本地准备 CMake 或 Xcode。安装特定后端可选需要在 Linux 上用 CUDA、只跑 CPU或从源码定制时才走下面的路径# Linux CUDA 后端 pip install mlx[cuda12] # Linux CPU 专用 pip install mlx[cpu] # 从源码构建需要 C20 编译器、cmake 3.25 git clone https://gitcode.com/GitHub_Trending/ml/mlx cd mlx pip install -e .[dev]核心能力拆解延迟计算先把算子记下来需要时再执行你写的每个操作并不会立刻执行而是被记录成一张计算图只有真正需要结果时才落地。好处是未使用的中间结果不会被计算也方便后续做微分和图优化。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], dtypefloat32)打印数组时会自动触发计算所以print(c)也会得到一个具体的结果。统一内存同一份数据CPU 与 GPU 随用随取Apple 芯片的 CPU 和 GPU 共享同一块内存因此 MLX 里创建数组不用指定它在哪个设备而是把设备写在操作上。矩阵乘法交给 GPU、小算子交给 CPU数据无需来回拷贝调度器会自动处理依赖。x mx.random.normal((4096, 512)) y mx.random.normal((512, 4)) out mx.matmul(x, y, streammx.gpu) # 密集的矩阵乘法跑在 GPU mx.eval(out) # 形状为 (4096, 4)这条输出是一个(4096, 4)的数组不同设备上的算子可以并行推进。可组合函数变换自动微分与批量向量化grad、vmap这些变换可以任意嵌套你不需要手写反向传播也不需要为批处理复制代码。import mlx.core as mx x mx.array(0.0) print(mx.grad(mx.sin)(x)) # 一阶导array(1, dtypefloat32) print(mx.grad(mx.grad(mx.sin))(x)) # 二阶导array(-0, dtypefloat32) batched mx.vmap(mx.sin) # 对整批数据同时计算 print(batched(mx.array([0.0, 1.0])))最后一行输出[0, 0.841471]vmap把原本标量函数扩展成了批量函数。当模型大到单设备放不下时MLX 的分布式层能把一张Linear拆到多个设备上做张量并行再聚合输出。端到端实战下面是一个完整的训练 → 保存权重 → 重新加载推理小例子训练一个 4 分类的线性分类器然后把权重存成 safetensors再用一个全新模型加载并做预测。import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim # 1) 合成一个 4 分类数据集 X mx.random.normal((1000, 64)) w_star mx.random.normal((64, 4)) y mx.argmax(X w_star, axis1) # 2) 定义模型与优化器 model nn.Linear(64, 4) opt optim.Adam(learning_rate1e-2) def loss_fn(model, X, y): return nn.losses.softmax_cross_entropy(model(X), y) loss_and_grad mx.value_and_grad(loss_fn, model.parameters()) for step in range(300): loss, grads loss_and_grad(model, X, y) opt.update(model, grads) mx.eval(model.parameters()) # 3) 保存训练好的权重 mx.save_safetensors(clf.safetensors, model.parameters()) # 4) 推理全新模型加载权重后预测 fresh nn.Linear(64, 4) fresh.load_weights(mx.load(clf.safetensors)) print(最终损失, loss.item(), 预测, mx.argmax(fresh(X[:3]), axis1))训练结束会打印出收敛的损失值和前三个样本的预测类别。把输入维度从 64 换成 3072、类别数换成词表大小这个结构就变成一个可以直接接在文本编码器后面的分类头。常见问题与调优pip 提示找不到 mlx 包现象是系统版本明明达标pip install mlx却报 No matching distribution。原因是你用的 Python 是 Rosetta 下的 x86 版本装不上为 arm 编译的包。先用下面命令确认架构python -c import platform; print(platform.processor())输出应为arm若是i386换用原生 arm 的 Python例如 Conda再装即可。加载大模型时内存峰值翻倍现象是模型刚建好就占用接近两倍的显存。原因是模型先以 float32 完成初始化随后才被权重覆盖中间会留下两份。解法是利用延迟计算让权重在替换阶段才落地model Model() # 此时尚未分配显存 model.load_weights(weights.safetensors) # 加载 fp16 权重峰值减半想看清 GPU 上每个算子在做什么现象是性能问题难以定位。原因是 Metal 的 GPU 轨迹默认不记录源码和对象标签。解法是构建时打开调试开关运行时启用捕获CMAKE_ARGS-DMLX_METAL_DEBUGON pip install .再配合MTL_CAPTURE_ENABLED1运行用mx.metal.start_capture(trace.gputrace)与mx.metal.stop_capture()把 GPU 轨迹导出即可在 Xcode 里回放。捕获到的依赖视图能让你看到Arange、Add等算子在各 Stream 上的排布用来判断瓶颈在哪一步。下一步与资源官方文档里与本文对应的几篇按由浅入深的顺序阅读即可快速入门数组、eval 与函数变换的基础保存与加载.npy、.npz、safetensors 等格式逻辑回归示例纯 core 的训练循环写法MLP 导出示例把训练步骤导出为可复用函数掌握延迟计算和统一内存这两点后你就可以把现有的 NumPy 或 PyTorch 代码逐步迁到 MLX 上先在 Mac 上把流程跑通再决定要不要扩展到其他设备。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/4 9:18:45 · 阅读 1.2k 如何用 MLX 在 Apple 芯片上训练与部署模型
如何用 MLX 在 Apple 芯片上训练与部署模型【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx在 Apple 芯片上跑深度学习时你多半遇到这种情况模型在 CPU 和 GPU 之间来回搬数据或者干脆只能在 Linux CUDA 环境里训练。MLX 是 Apple 机器学习研究团队推出的一个数组框架专为 Apple Silicon 优化让你用接近 NumPy 的语法完成模型训练、推理与部署。它的延迟计算、统一内存和可组合函数变换三个设计正是为在 Mac 上高效跑模型而做的。这篇笔记带你完成安装、跑通一个完整的训练加推理示例并给出几个常见坑的处理方式。环境与安装先确认设备满足要求以 macOS 版本为例macOS 14.0且为 Apple Silicon 芯片原生 Python 3.10架构为 arm非 RosettaLinux 版需要 glibc 2.35CUDA 版额外要求 CUDA 12.0、Nvidia 架构 SM 7.5最小安装就一条命令pip install mlx这样装是因为 PyPI 包已经编译好 Metal 后端无需你本地准备 CMake 或 Xcode。安装特定后端可选需要在 Linux 上用 CUDA、只跑 CPU或从源码定制时才走下面的路径# Linux CUDA 后端 pip install mlx[cuda12] # Linux CPU 专用 pip install mlx[cpu] # 从源码构建需要 C20 编译器、cmake 3.25 git clone https://gitcode.com/GitHub_Trending/ml/mlx cd mlx pip install -e .[dev]核心能力拆解延迟计算先把算子记下来需要时再执行你写的每个操作并不会立刻执行而是被记录成一张计算图只有真正需要结果时才落地。好处是未使用的中间结果不会被计算也方便后续做微分和图优化。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], dtypefloat32)打印数组时会自动触发计算所以print(c)也会得到一个具体的结果。统一内存同一份数据CPU 与 GPU 随用随取Apple 芯片的 CPU 和 GPU 共享同一块内存因此 MLX 里创建数组不用指定它在哪个设备而是把设备写在操作上。矩阵乘法交给 GPU、小算子交给 CPU数据无需来回拷贝调度器会自动处理依赖。x mx.random.normal((4096, 512)) y mx.random.normal((512, 4)) out mx.matmul(x, y, streammx.gpu) # 密集的矩阵乘法跑在 GPU mx.eval(out) # 形状为 (4096, 4)这条输出是一个(4096, 4)的数组不同设备上的算子可以并行推进。可组合函数变换自动微分与批量向量化grad、vmap这些变换可以任意嵌套你不需要手写反向传播也不需要为批处理复制代码。import mlx.core as mx x mx.array(0.0) print(mx.grad(mx.sin)(x)) # 一阶导array(1, dtypefloat32) print(mx.grad(mx.grad(mx.sin))(x)) # 二阶导array(-0, dtypefloat32) batched mx.vmap(mx.sin) # 对整批数据同时计算 print(batched(mx.array([0.0, 1.0])))最后一行输出[0, 0.841471]vmap把原本标量函数扩展成了批量函数。当模型大到单设备放不下时MLX 的分布式层能把一张Linear拆到多个设备上做张量并行再聚合输出。端到端实战下面是一个完整的训练 → 保存权重 → 重新加载推理小例子训练一个 4 分类的线性分类器然后把权重存成 safetensors再用一个全新模型加载并做预测。import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim # 1) 合成一个 4 分类数据集 X mx.random.normal((1000, 64)) w_star mx.random.normal((64, 4)) y mx.argmax(X w_star, axis1) # 2) 定义模型与优化器 model nn.Linear(64, 4) opt optim.Adam(learning_rate1e-2) def loss_fn(model, X, y): return nn.losses.softmax_cross_entropy(model(X), y) loss_and_grad mx.value_and_grad(loss_fn, model.parameters()) for step in range(300): loss, grads loss_and_grad(model, X, y) opt.update(model, grads) mx.eval(model.parameters()) # 3) 保存训练好的权重 mx.save_safetensors(clf.safetensors, model.parameters()) # 4) 推理全新模型加载权重后预测 fresh nn.Linear(64, 4) fresh.load_weights(mx.load(clf.safetensors)) print(最终损失, loss.item(), 预测, mx.argmax(fresh(X[:3]), axis1))训练结束会打印出收敛的损失值和前三个样本的预测类别。把输入维度从 64 换成 3072、类别数换成词表大小这个结构就变成一个可以直接接在文本编码器后面的分类头。常见问题与调优pip 提示找不到 mlx 包现象是系统版本明明达标pip install mlx却报 No matching distribution。原因是你用的 Python 是 Rosetta 下的 x86 版本装不上为 arm 编译的包。先用下面命令确认架构python -c import platform; print(platform.processor())输出应为arm若是i386换用原生 arm 的 Python例如 Conda再装即可。加载大模型时内存峰值翻倍现象是模型刚建好就占用接近两倍的显存。原因是模型先以 float32 完成初始化随后才被权重覆盖中间会留下两份。解法是利用延迟计算让权重在替换阶段才落地model Model() # 此时尚未分配显存 model.load_weights(weights.safetensors) # 加载 fp16 权重峰值减半想看清 GPU 上每个算子在做什么现象是性能问题难以定位。原因是 Metal 的 GPU 轨迹默认不记录源码和对象标签。解法是构建时打开调试开关运行时启用捕获CMAKE_ARGS-DMLX_METAL_DEBUGON pip install .再配合MTL_CAPTURE_ENABLED1运行用mx.metal.start_capture(trace.gputrace)与mx.metal.stop_capture()把 GPU 轨迹导出即可在 Xcode 里回放。捕获到的依赖视图能让你看到Arange、Add等算子在各 Stream 上的排布用来判断瓶颈在哪一步。下一步与资源官方文档里与本文对应的几篇按由浅入深的顺序阅读即可快速入门数组、eval 与函数变换的基础保存与加载.npy、.npz、safetensors 等格式逻辑回归示例纯 core 的训练循环写法MLP 导出示例把训练步骤导出为可复用函数掌握延迟计算和统一内存这两点后你就可以把现有的 NumPy 或 PyTorch 代码逐步迁到 MLX 上先在 Mac 上把流程跑通再决定要不要扩展到其他设备。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考