MLX 有多快？在 8 个苹果硅芯片和 4 个 CUDA GPU 上的全面基准测试
原文：
towardsdatascience.com/how-fast-is-mlx-a-comprehensive-benchmark-on-8-apple-silicon-chips-and-4-cuda-gpus-378a0ae356a0
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/802ead39f0c1351a622b8d14ccd5011b.png
作者提供的图片：softmax 操作的基准示例
自从首次发布以来不到两个月，苹果的 ML 研究团队最新成果 MLX 已经在 ML 社区取得了显著进展。新框架迅速获得关注令人瞩目，正如 GitHub 上超过 12k 个星标和 Hugging Face 上超过 500 名成员的社区增长所证明的那样 🤗。
在之前的文章中，我们展示了 MLX 在训练简单的图卷积网络（GCN）中的表现，将其与包括CPU，PyTorch 的**MPS和CUDA** GPU 在内的各种设备进行了基准测试。结果令人启发，显示了 MLX 在高效运行模型方面的潜力。
在这次探索中，我们深入挖掘，旨在基准测试神经网络中常用的一些关键操作。
测试平台
在我们的基准测试中，每个操作都是基于各种实验进行评估的，这些实验在输入形状和大小上有所不同。我们已经按顺序多次在不同的进程上运行这些实验，以确保稳定可靠的运行时间测量。
在开放协作的精神下，我们已经将基准代码开源并易于运行。这使得贡献者可以轻松地根据他们的设备和配置添加自己的基准。
GitHub – TristanBilot/mlx-benchmark: 苹果 MLX 操作的基准测试
注意：非常感谢所有贡献者，没有他们，这个基准就不会包含如此多的基线芯片。
我们成功地在10个不同的苹果硅芯片和3个高效CUDAGPU 上运行了这个基准。
苹果硅：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 的操作。
线性层
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/7611189e70665c3464aa1ead41537644.png
作者提供的图片：线性操作基准
MLX 比 MPS 快约 2 倍，但在 M3*芯片上表现较差
性能最佳的 M2 Ultra 芯片比两个 Tesla V100 GPU 快！RTX4090 的速度仍然是 M2 Ultra 的 3 倍。
Softmax
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/d2833cd7cd7ff5617cc8e7555b961360.png
图片由作者提供：Softmax 操作基准
在某些芯片上，MLX 比 MPS 快达 2 倍
CUDA GPU 比 M2 快约 10 倍，比 M2 Ultra 快约 4 倍
Sigmoid
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e98a136551ac4a0b3b8b7531faa094a2.png
图片由作者提供：Sigmoid 操作基准
MLX 在所有芯片上的速度比 MPS 快约 30%
M2 Max 在这里比 M2 Ultra 快
Concatenation
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/c0d7e74f37886457355e60a90f45a8a3.png
图片由作者提供：Concatenation 操作基准
MLX 和 MPS 几乎相当
CUDA GPU 和 M2 Ultra 之间没有太大差异，而 M3 在这里特别慢
二元交叉熵（BCE）
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/2b510dc79941402769f6353d70ac9150.png
图片由作者提供：BCE 操作基准
BCE 在 M1 和 M2 上的 MPS 实现似乎非常慢
M2 Max、M2 Ultra 和 M3 Max 比 CUDA GPU 慢约 3 倍
Sort
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/ad33fe85cca507edf79979f7ca8a7721.png
图片由作者提供：Sort 操作基准
MLX 的排序操作真的很快。
在 M2 Max、M2 Ultra 和 M3 Max 上，它实现了比所有 CUDA GPU（包括最新的 RTX4090）更好的性能
Conv2D
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e220cadb837cc0f31ba0051845522a32.png
图片由作者提供：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 之间的数据传输。这听起来可能是一个小细节，但在机器学习项目的现实世界中，这些数据传输是延迟的知名来源。我们的上一篇文章强调了这一点，展示了在训练期间考虑数据传输时间时的基准测试结果非常不同：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/f5879a736277afa5f30b9241f7c0536d.png
图片由作者提供：带有设备数据传输的 GCN 基准测试
这些结果显示了图卷积网络模型每轮平均训练时间，该模型基本上由两个线性层组成。与本文中的基准测试不同，这个特定的基准测试评估了完整训练循环的平均运行时间，包括从 CPU 到 GPU 的数据传输时间。
这些结果揭示了一个有趣的见解：当包含实际数据传输时间时，CUDA GPU 的性能明显减慢。这强调了数据传输对整体速度的显著影响，这是我们当前基准测试中不能忽视的一个因素。
结论
这个基准测试让我们清楚地了解了 MLX 与在 MPS 和 CUDA GPU 上运行的 PyTorch 相比的表现。我们发现，MLX 对于大多数操作通常比 MPS 快得多，但比 CUDA GPU 慢。
使 MLX 真正脱颖而出的是其统一的内存，它消除了 CPU 和 GPU 之间耗时数据传输的需要。
在未来的文章中，我们计划在我们的基准测试中包含数据传输时间。这可能会提供一个更全面的了解 MLX 的整体效率。
— – – – – – – –
感谢阅读！
— 特别感谢 Ian Myers 在之前的 CUDA 基准测试中发现了问题。