苹果为自家芯片打造开源框架MLX，实现Llama 7B并在M2 Ultra上运行

2020 年 11 月，苹果推出 M1 芯片，其速度之快、功能之强大一时令人惊艳。2022 年苹果又推出了 M2，今年 10 月，M3 芯片正式登场。

苹果在发布芯片的同时，也非常注重在其芯片上训练和部署 AI 模型的能力。

苹果推出的 ML Compute 可用于在 Mac 上进行 TensorFlow 模型的训练。PyTorch 则支持在 M1 版本的 Mac 上进行 GPU 加速的 PyTorch 机器学习模型训练，使用苹果 Metal Performance Shaders (MPS) 作为后端来实现。这些使得 Mac 用户能够在本地训练神经网络。

现在，苹果宣布推出专门在 Apple 芯片上用于机器学习的开源阵列框架 ——MLX。

MLX 是专门为机器学习研究人员设计的，旨在有效地训练和部署 AI 模型。框架本身的设计在概念上也很简单。研究人员能够轻松地扩展和改进 MLX，以快速探索、测试新的想法。MLX 的设计灵感来自 NumPy、PyTorch、Jax 和 ArrayFire 等框架。

项目地址：https://github.com/ml-explore/mlx

MLX 项目贡献者之一、Apple 机器学习研究团队（MLR）研究科学家 Awni Hannun 展示了一段使用 MLX 框架实现 Llama 7B 并在 M2 Ultra 上运行的视频。

MLX 迅速引起机器学习研究人员的关注。TVM、MXNET、XGBoost 作者，CMU 助理教授，OctoML CTO 陈天奇转推表示：「苹果芯片又有新的深度学习框架了。」

有网友评价 MLX 称，苹果再次「重造了轮子」。

图源：https://twitter.com/ofervic/status/1732305883814596953

MLX 特性、示例

在该项目中，我们可以看到，MLX 有以下一些主要特性。

熟悉的 API。MLX 拥有非常像 NumPy 的 Python API，以及功能齐备的 C++ API（与 Python API 非常相似）。MLX 还有更高级的包（比如 mlx.nn 和 mlx.optimizers），它们的 API 很像 PyTorch，可以简化构建更复杂的模型。

可组合函数变换。MLX 拥有自动微分、自动矢量化和计算图优化的可组合函数变换。

惰性计算。MLX 中的计算是惰性的，阵列只有在需要时才被实例化。

动态图构建。MLX 中的计算图构建是动态的，改变函数参数的形状不会导致编译变慢，并且 debug 很简单、容易上手。

多设备。任何支持的设备上（如 CPU 和 GPU）都可以运行操作。

统一内存。MLX 与其他框架的显著差异在于统一内存，阵列共享内存。MLX 上的操作可以在任何支持的设备类型上运行，无需移动数据。

此外，项目中提供了多种使用 MLX 框架的示例，比如 MNIST 示例可以很好地让你学习如何使用 MLX。

图源：https://github.com/ml-explore/mlx-examples/tree/main/mnist

MLX 还有其他更多有用的示例，包括如下：

• Transformer 语言模型训练；

• LLaMA 大规模文本生成和 LoRA 微调；

• Stable Diffusion 生成图片；

• OpenAI 的 Whisper 语音识别。

更详细的文档可参阅：https://ml-explore.github.io/mlx/build/html/install.html#