苹果版CUDA来了！专为自家芯片打造，M3 Max可跑每秒迭代2.8次

　　苹果新发布适用于苹果芯片的机器学习框架 MLX，MLX 是一个类似于 NumPy 的阵列框架，旨在为使用苹果硬件的研究人员简化 ML 模型的训练和部署。

　　一直以来，英伟达 CUDA 称霸天下，使得它在 AI 领域的势力不可动摇。

　　对于许多人来说，都希望打破这个护城河，比如 PyTorch 支持更多 GPU，OpenAI 的 Triton 等等。

　　近日，苹果针对使用苹果芯片的机器学习研究人员，专门发布了一款机器学习框架 MLX。

　　苹果机器学习团队的 Awni Hannun 在X-note 中将该软件称为：「……专为苹果芯片设计的高效机器学习框架」。

　　MLX 使用起来类似 PyTorch、Jax 和 ArrayFire 等现有框架。

　　但是，MLX 针对 Apple 芯片进行了优化。以前受困于艰难使用M芯片 GPU 跑模型的研究人员，现在终于可以摆脱这个苦恼了!

　　同时，MLX 增加了对统一内存模型的支持，这也意味着阵列位于共享内存中，并且可以在任何支持的设备类型上执行操作，而无需执行数据复制。那么，一旦苹果的其他产品也能够开始跑模型……有没有很期待!

　　英伟达高级科学家 Jim Fan 表示，这可能是苹果公司迄今为止在开源人工智能方面的最大举措：

　　MLX 特别之处

　　先让我们来看看 MLX 的速度：

　　一位网友对 MLX 的速度进行了测试，他在推文中表示，使用 Apple MLX 框架进行计算的速度，M2 Ultra(76 个 GPU)可达到每秒 5 次迭代(It/sec)，而 M3 Max(40 个 GPU)可以达到每秒 2.8 次迭代。

　　这个速度是与 Pytorch 几乎持平的。

　　MLX 的另一大特点，像 Jim Fan 指出的那样，是它设计出色的 API。

　　对深度学习领域的开发人员来说，上手几乎零成本。

　　这是因为 MLX 的设计灵感来自于 PyTorch、Jax 和 ArrayFire 等现有框架。

　　研究团队解释道：「Python API 与 NumPy 高度相似，只有少数例外。MLX 还拥有一个功能齐全的 C++ API，它与 Python API 高度一致。」

　　研究团队对该设计的目的作了补充：「该框架旨在兼顾对用户操作友好，同时保持训练和部署模型的高效。我们的目标是使研究人员能够轻松扩展和改进 MLX，以便快速探索新的想法。」

　　除了可以轻松上手的 API 之外，MLX 的一些其他关键特性还包括：

　　可组合的函数变换：MLX 含有用于自动微分、自动向量化和计算图优化的可组合函数变换，这对于优化和加速机器学习模型的训练过程非常有用。

　　延迟计算：MLX 支持延迟计算，数组只在需要时才会被实例化。

　　动态图构建：MLX 中的计算图是动态构建的。改变函数参数的形状不会触发缓慢的编译，调试也简单直观。这使得模型的开发和调试更加灵活和高效。

　　多设备支持：操作可以在任何支持的设备上运行(目前支持 CPU 和 GPU)。

　　统一内存：MLX 与其他框架的显著区别在于统一的内存模型。在 MLX 中，数组存在于共享内存中，对 MLX 数组的操作可以在任何支持的设备类型上执行，而无需移动数据。

　　MLX Data

　　除了发布 MLX 以外，苹果机器学习团队还一起发表了 MLX Data。这是一个与 MLX 框架无关的数据加载库，可以与 PyTorch、Jax 或 MLX 一同使用。

　　MLX Data 可以用于加载大规模数据集，也可以独立用于数据预处理，以供机器学习训练使用。

　　该库的目标是允许用户利用多线程来加速数据处理流水线，而无需涉及复杂的多进程操作或使用符号语言编写代码。这有利于帮助提高数据加载和预处理的效率。

　　MLX 示例展示

　　目前，用户可以通过 PyPI 轻松安装 MLX，只需执行以下命令：pip install mlx

　　用户需要确保满足以下安装要求：

　　- 使用M系列芯片(苹果芯片)

　　- 使用本地 Python 版本不低于 3.8

　　- MacOS 版本需不低于 13.3

　　研究团队同时提供了一系列有关 MLX 功能的示例，包括 Transformer 语言模型训练、使用 LLaMA 进行大规模文本生成、使用 LoRA 进行微调、使用稳定扩散生成图像，以及用 OpenAI 的 Whisper 进行语音识别等。

　　这些例子也侧面证明了该公司在语言处理、图像生成和语音识别等方面的技术实力和工具的强大程度。

　　其他开发人员也对于 MLX 进行了其他有趣的尝试，比如一位来自塞浦路斯大学的教授，在使用 PyTorch 和新的 MLX 框架的情况下，成功在 Apple M2 Ultra GPU 上解决了一个板上的二维热传导方程问题，其中包含两个 Dirichlet 和两个 Neumann 边界条件。

　　有趣的是，MLX 的计算速度大约是之前的两倍，并且可以直接使用 Matplotlib 进行绘图，无需转移到 CPU。

　　从 MLX 一瞥苹果对于人工智能的展望

　　普惠机器学习

　　研究团队一再强调，「MLX 是由机器学习研究人员为机器学习研究人员设计的。」 这表明苹果意识到为机器学习开发开放、易于使用的开发环境的需求，以促进该领域的进一步发展。

　　芯片技术

　　因为苹果的处理器现在存在于其所有产品中，包括 Mac、iPhone 和 iPad。这种在这些芯片上使用 GPU、CPU 和(有可能在某个时刻)神经引擎的方法可能会转化为在设备上执行 ML 模型。其性能可能超过其他处理器，至少在边缘设备方面。

　　开发工具的重要性

　　苹果明确决定将重点放在为机器学习研究人员提供最佳工具上，包括强大的 M3 Mac，用于模型构建。现在，他们进一步希望将这种关注转化为对普通用户而言可行、以人为本的 AI 工具。

　　总体而言，苹果希望通过提供强大的工具和在不同设备上的性能优势，将机器学习和人工智能推向更广泛的用户。