ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

如何在 Apple Silicon 上做机器学习:MLX 新手实战指南

如何在 Apple Silicon 上做机器学习:MLX 新手实战指南 如何在 Apple Silicon 上做机器学习MLX 新手实战指南【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlxMLX 是 Apple 机器学习研究团队推出的、面向 Apple silicon 的数组框架array framework。它的 Python API 与 NumPy 高度相似内置延迟计算、统一内存模型和可组合的函数变换自动微分、向量化可以在 Mac 上直接训练和推理模型。如果你熟悉 PyTorch 的写法、想把代码搬到 Apple Silicon 设备或者想在 MacBook 上跑本地模型这个项目值得花一小时上手。安装命令怎么选三种场景根据你的运行环境三选一即可pip install mlx # Apple Silicon Mac需 macOS 14.0、原生 Python 3.10 pip install mlx[cuda12] # Linux NVIDIA GPUSM 7.5 及以上架构 pip install mlx[cpu] # Linux 纯 CPU 版本在 Mac 上如果pip找不到匹配版本多半是 Python 不是原生的 arm 版本运行python -c import platform; print(platform.processor())输出应为arm如果是i386请换成原生 Python比如通过 Conda 安装。5 分钟跑通第一个数组先理解延迟计算下面这段代码演示了 MLX 最核心的行为——操作先记录、后执行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 # 此时 c 只是计算图上的一个配方尚未真正计算 mx.eval(c) # 显式触发计算 print(c) # array([2, 4, 6, 8], dtypefloat32)你会发现c a b之后并没有立即产生数值MLX 只是在记录一张动态计算图只有在mx.eval、打印数组、调用.item()或转成 numpy 时才会真正落盘计算。这个设计的直接好处是图可以在执行前被变换求导、向量化而且最终没被使用的输出不会白算。更多细节可以看仓库里的说明docs/src/usage/lazy_evaluation.rst统一内存在操作时指定设备而不是搬移数组Apple Silicon 的 CPU 和 GPU 直接共享同一个内存池MLX 完全利用了这一点创建数组时不指定位置而是在调用操作时通过stream参数指定设备a mx.random.normal((100,)) b mx.random.normal((100,)) mx.add(a, b, streammx.cpu) # 让 CPU 算 mx.add(a, b, streammx.gpu) # 让 GPU 算与上一条并行执行这两个操作互不依赖MLX 会并行调度如果第二条依赖第一条的结果调度器会自动在两个 stream 之间插入依赖不需要你手动同步。官方文档里有个实测例子M1 Max矩阵乘法放 GPU、几百次小的exp放 CPU总耗时约 1.4ms比全部堆在 GPU 上2.8ms快一倍。可组合的函数变换求导和向量化各一行MLX 提供了mx.grad、mx.vmap、mx.jvp、mx.vjp、mx.value_and_grad等变换并且可以任意嵌套组合。比如求二阶导就是mx.grad(mx.grad(mx.sin))(x)mx.grad(mx.vmap(mx.grad(fn)))这样的组合也是合法的。两个最常用的模式训练时用mx.value_and_grad(fn)一次拿到函数值和梯度避免重复计算前向图批量推理时给模型外层套一个mx.vmap(model)自动沿新维度向量化。保存与加载模型4 种序列化格式mx.load会根据文件扩展名自动识别格式支持.npy单数组、.npz多数组、.safetensors和.gguf后两者也常用于交换模型权重a, b mx.array([1.0]), mx.array([2.0]) mx.savez(arrays, a, bb) # 存为 arrays.npz mx.save_safetensors(arrays, {a: a, b: b}) # 存为 safetensors state mx.load(arrays.npz) # 返回 {名称: 数组} 的字典动手示例跑仓库自带的训练脚本仓库里有一个自包含的线性回归训练脚本examples/python/linear_regression.py直接用python examples/python/linear_regression.py运行即可。它的流程值得逐行读一遍生成随机设计矩阵X和带噪声的标签y→ 定义损失0.5 * mx.mean(mx.square(X w - y))→ 用mx.grad(loss_fn)拿到梯度函数 → 循环 10000 次做w w - lr * grad并mx.eval(w)固化结果 → 最后打印损失、与真实参数的距离和吞吐it/s。你可以把它当模板替换成自己的数据就得到一个最小训练闭环。排坑与性能eval、内存、GPU 分析忘记 eval 会导致计时失真训练循环里如果不 eval计算图会一直累积time.perf_counter()测到的只是建图时间。每步或每几步eval 一次参数既是正确性也是性能问题。峰值内存可以省一半得益于延迟计算先以 float32 实例化大模型、再换成 float16 权重中间不会真的算出 float32 那份数据峰值内存接近减半。用 Metal 调试器分析 GPU以-DMLX_METAL_DEBUGON构建后调用mx.metal.start_capture()/mx.metal.stop_capture()需配合MTL_CAPTURE_ENABLED1环境变量就能捕获 GPU 追踪文件在 Xcode 里回放。指南在 docs/src/dev/metal_debugger.rst。多设备张量并行要把大模型分片到多个设备时mlx.nn提供了AllToShardedLinear和ShardedToAllLinear等现成分层两者前后衔接可以省掉中间通信。完整示例见 examples/python/distributed_tensor_parallel.py。去哪继续学快速上手docs/src/usage/quick_start.rst更多示例examples/python/含逻辑回归、分布式数据/张量并行、量化推理建议你先pip install mlx然后原样跑一遍examples/python/linear_regression.py把延迟计算和mx.eval的节奏跑熟再动手替换成自己的数据。等你熟悉了这条主线函数变换和张量并行这两块会顺理成章地接上。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表