作者提供的图片：softmax 操作的基准示例
自从首次发布以来不到两个月，苹果的 ML 研究团队最新成果 MLX 已经在 ML 社区取得了显著进展。新框架迅速获得关注令人瞩目，正如 GitHub 上超过 12k 个星标和 Hugging Face 上超过 500 名成员的社区增长所证明的那样 🤗。
在之前的文章中，我们展示了 MLX 在训练简单的图卷积网络（GCN）中的表现，将其与包括CPU，PyTorch 的**MPS和CUDA** GPU 在内的各种设备进行了基准测试。结果令人启发，显示了 MLX 在高效运行模型方面的潜力。
在这次探索中，我们深入挖掘，旨在基准测试神经网络中常用的一些关键操作。
测试平台
在我们的基准测试中，每个操作都是基于各种实验进行评估的，这些实验在输入形状和大小上有所不同。我们已经按顺序多次在不同的进程上运行这些实验，以确保稳定可靠的运行时间测量。
在开放协作的精神下，我们已经将基准代码开源并易于运行。这使得贡献者可以轻松地根据他们的设备和配置添加自己的基准。
GitHub – TristanBilot/mlx-benchmark: 苹果 MLX 操作的基准测试
注意：非常感谢所有贡献者，没有他们，这个基准就不会包含如此多的基线芯片。
我们成功地在10个不同的苹果硅芯片和3个高效CUDA GPU 上运行了这个基准。
苹果硅：M1，M1 Pro，M1 Max，M2，M2 Pro，M2 Max，M2 Ultra，M3，M3 Pro，M3 Max
CUDA GPU：RTX4090 128GB (笔记本电脑)，Tesla V100 32GB (NVLink)，Tesla V100 32GB (PCIe)。
基准测试结果
对于每个基准，运行时间以毫秒为单位测量。MPS 和 CUDA 基线都使用 PyTorch 中实现的操作，而苹果硅基线使用 MLX 的操作。
线性层
作者提供的图片：线性操作基准
-
MLX 比 MPS 快约 2 倍，但在 M3*芯片上表现较差
-
性能最佳的 M2 Ultra 芯片比两个 Tesla V100 GPU 快！RTX4090 的速度仍然是 M2 Ultra 的 3 倍。
Softmax
图片由作者提供：Softmax 操作基准
-
在某些芯片上，MLX 比 MPS 快达 2 倍
-
CUDA GPU 比 M2 快约 10 倍，比 M2 Ultra 快约 4 倍
Sigmoid
图片由作者提供：Sigmoid 操作基准
-
MLX 在所有芯片上的速度比 MPS 快约 30%
-
M2 Max 在这里比 M2 Ultra 快
Concatenation
图片由作者提供：Concatenation 操作基准
-
MLX 和 MPS 几乎相当
-
CUDA GPU 和 M2 Ultra 之间没有太大差异，而 M3 在这里特别慢
二元交叉熵（BCE）
图片由作者提供：BCE 操作基准
-
BCE 在 M1 和 M2 上的 MPS 实现似乎非常慢
-
M2 Max、M2 Ultra 和 M3 Max 比 CUDA GPU 慢约 3 倍
Sort
图片由作者提供：Sort 操作基准
-
MLX 的排序操作真的很快。
-
在 M2 Max、M2 Ultra 和 M3 Max 上，它实现了比所有 CUDA GPU（包括最新的 RTX4090）更好的性能
Conv2D
图片由作者提供：Conv2D 操作基准
-
MLX 比 MPS 慢 2-5 倍
-
RTX4090 在此操作中特别快
讨论
结果
我们的基准测试得出以下结论：
-
在大多数操作中，运行在 Apple Silicon 上的 MLX 在性能上始终优于使用 MPS 后端的 PyTorch。
-
CUDA GPU 在不可避免地快于 Apple Silicon。
-
显著的是，MLX 在某些操作（如Sort和Linear）中表现出色，即使与 CUDA GPU 基准测试相比也显示出强大的效率。
-
使用 MLX 进行卷积操作仍然特别慢。似乎 MLX 团队目前正在努力提高这些操作的速度。
注意：好奇的读者可能会对检查基准测试中可用的完整操作集感兴趣。
Apple Silicon 与 CUDA GPU
虽然 CUDA GPU 无疑更快，但退一步考虑更广泛的图景至关重要。英伟达多年来一直在完善其 GPU 技术，达到了成熟和性能的水平，目前无出其右。当你权衡 Nvidia GPU 和 Apple Silicon 芯片在硬件复杂度、成熟度和成本方面的差异时，这提供了一个重要的视角。
这就是 MLX 发挥其优势的地方。这个框架在需要在其个人 Mac 上运行大型模型的情况下特别有益。有了 MLX，你可以在不需要额外、昂贵的 GPU 硬件的情况下实现这一点。对于那些寻求在 Mac 设备上直接获得高效机器学习能力的用户来说，这是一个变革性的游戏。
显著的是，最近使用 MLX 在苹果硅上直接运行LLM 推理越来越受欢迎。
统一内存改变了游戏规则
苹果硅在 GPU 领域中脱颖而出，其统一的内存架构独特地利用了 Mac 的全部 RAM 来运行模型。这与传统的 GPU 设置相比是一个重大的变化。
但也许使用苹果硅最引人注目的优势是消除了 CPU 和 GPU 之间的数据传输。这听起来可能是一个小细节，但在机器学习项目的现实世界中，这些数据传输是延迟的知名来源。我们的上一篇文章强调了这一点，展示了在训练期间考虑数据传输时间时的基准测试结果非常不同：
图片由作者提供：带有设备数据传输的 GCN 基准测试
这些结果显示了图卷积网络模型每轮平均训练时间，该模型基本上由两个线性层组成。与本文中的基准测试不同，这个特定的基准测试评估了完整训练循环的平均运行时间，包括从 CPU 到 GPU 的数据传输时间。
这些结果揭示了一个有趣的见解：当包含实际数据传输时间时，CUDA GPU 的性能明显减慢。这强调了数据传输对整体速度的显著影响，这是我们当前基准测试中不能忽视的一个因素。
结论
这个基准测试让我们清楚地了解了 MLX 与在 MPS 和 CUDA GPU 上运行的 PyTorch 相比的表现。我们发现，MLX 对于大多数操作通常比 MPS 快得多，但比 CUDA GPU 慢。
使 MLX 真正脱颖而出的是其统一的内存，它消除了 CPU 和 GPU 之间耗时数据传输的需要。
在未来的文章中，我们计划在我们的基准测试中包含数据传输时间。这可能会提供一个更全面的了解 MLX 的整体效率。
— – – – – – – –
感谢阅读！
— 特别感谢 Ian Myers 在之前的 CUDA 基准测试中发现了问题。
扫一扫
554
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
