苹果公司开源机器学习框架 MLX，针对 Silicon 芯片进行了优化

苹果公司的机器学习框架MLX结合了开发者熟悉的 API、可组合的函数转换和惰性计算，部分灵感源于 NumPy 和 PyTorch，并针对苹果的 Silicon 进行了优化。该框架使用 Python 和 C++实现，旨在为在苹果 Silicon 上训练和部署机器学习模型提供用户友好且高效的解决方案。

根据苹果公司的说法，MLX 是由机器学习研究人员为机器学习研究人员设计的，并基于 MIT 发布许可，可以很容易地被扩展和改进。它支持转换语言模型训练、使用 Mistral 进行大规模文本生成、使用 Stable Diffusion 进行图像生成以及使用 Whisper 进行语音识别。

MLX 提供了受 NumPy 启发的底层 Python API 和一个完整的与之密切对应的 C++ API。此外，它还提供了一个高级 API，可用于根据 PyTorch API 创建更复杂的模型。

该框架支持自动微分、自动向量化和计算图优化，可组合的函数使得构建复杂数组转换变得更加容易。MLX 还支持惰性计算，这意味着它可以只在必要时才计算数组，以提高计算效率。同样，计算图是动态构建的，因此修改函数参数并不会触发缓慢的编译过程。

MLX 的一个独有的特性是使用了苹果 Silicon 的统一内存，这让它有别于其他的 ML 框架。这意味着数组位于共享内存中，可以在 CPU 或 GPU 上执行数组操作，无需在内存之间传输数据。例如，在创建一个数组时，你不需要指定位置，因为它位于统一内存中，而在执行操作时可以选择在 CPU 或 GPU 上执行转换：

a = mx.random.normal((100,))
b = mx.random.normal((100,))
mx.add(a, b, stream=mx.cpu)
mx.add(a, b, stream=mx.gpu)

MLX 可在任意的苹果 Silicon CPU 上运行，包括 M1，并可以利用集成的 GPU，因此研究人员可以选择最适合其需求的硬件。

MLX 的代码库中包含了一些针对不同模型的示例，包括 BERT、Llama、Mistral、Stable Diffusion 等。每个示例都在requirements.txt文件中列出所需的依赖项，并提供了现成的 CLI 工具。例如，要使用 Stable Diffusion 生成图像，首先安装所有必需的依赖项，然后运行txt2image.py命令：

pip install -r requirements.txt

python txt2image.py "A photo of an astronaut riding a horse on Mars." --n_images 4 --n_rows 2

苹果尚未公开发布基准测试，因此我们目前不知道它与PyTorch/MPS或 Georgi Gerganov 的Llama.cpp相比表现如何。

不过，Stable Diffusion 示例中包含了使用 PyTorch 和 MLX 运行 UNet 的性能比较。MLX 在批次大小为 16 时的吞吐量比 PyTorch 高约 40%，最佳批次大小大 15%左右。

然而，PyTorch 在较小的批次大小时表现更好，批次大小为 1 时吞吐量高约 50%，批次大小为 4 时高约 10%。根据苹果公司的说法，PyTorch 在这些情况下的优势要归因于在模型还没有被加载到内存中且 PyTorch 的 MPS 图内核未被缓存时的编译速度。

如果你有兴趣体验 MLX，请参阅其快速入门指南或完整文档。

原文链接：

https://www.infoq.com/news/2023/12/apple-silicon-machine-learning/