点击下方 卡片 ，关注“ 慢慢学AIGC ”
自首次发布以来不到两个月，苹果 ML 研究团队的最新创作 MLX 已经在 ML 社区取得了重大进展。 看到这个新框架如此迅速地获得关注令人印象深刻，这从 GitHub 上超过 12k 的星标和 Hugging Face 🤗 上超过 500 名成员的不断增长的社区可见一斑。
在之前的一篇文章中，我们展示了 MLX 在训练简单图卷积网络（GCN）方面的表现，将其与包括 CPU、PyTorch 的 MPS 和 CUDA GPU 在内的各种设备进行了基准测试。结果令人耳目一新，展示了 MLX 在高效运行模型方面的潜力。
在本次探索中，我们深入研究，着手对神经网络中常用的多个关键 op 进行基准测试。
测试平台
在我们的基准测试中，每个 op 都根据各种实验进行评估，这些实验在输入形状和大小上有所不同。我们在不同的进程中连续多次运行这些测试，以确保运行时间测量的稳定性和可靠性。
本着开放协作的精神，我们已经将基准测试代码开源并保证开箱即用。这允许贡献者根据自己的设备和配置轻松添加他们自己的基准测试。
代码：https://github.com/TristanBilot/mlx-benchmark
注意：非常感谢所有贡献者，没有他们，这个基准测试就不会包含如此多的基准芯片。
我们成功地在 10 种不同的苹果 M 芯片和 3 种高效的 CUDA GPU 上运行了这个基准测试：
苹果 M 芯片 : M1, M1 Pro, M1 Max, M2, M2 Pro, M2 Max, M2 Ultra, M3, M3 Pro, M3 Max
CUDA GPU : RTX 4090 128GB (笔记本), Tesla V100 32GB (NVLink), Tesla V100 32GB (PCIe)
基准测试结果
对于每个基准测试，运行时间以毫秒为单位测量。MPS 和 CUDA 基准都使用 PyTorch 中实现的 op，而苹果 M 芯片基准使用 MLX 的 op。
线性层
- MLX 比 MPS 快约 2 倍，但在 M3* 芯片上表现较差
- 性能最佳的 M2 Ultra 芯片比两款 Tesla V100 GPU 都快！RTX 4090 仍然比 M2 Ultra 快 3 倍。
Softmax
- MLX 在某些芯片上比 MPS 快达 2 倍
- CUDA GPU 比 M2 快约 10 倍，比 M2 Ultra 快约 4 倍
Sigmoid
- MLX 在所有芯片上比 MPS 快约30%
- M2 Max 在这里比 M2 Ultra 更快
连接（Concatenation）
- MLX 和 MPS 几乎相当
- CUDA GPU 和 M2 Ultra 之间没有太大差异，而 M3 在这里特别慢
二元交叉熵（BCE）
- MPS 对 BCE 的实现在 M1 和 M2 上似乎极其缓慢
- M2 Max、M2 Ultra 和 M3 Max 仅比 CUDA GPU 慢约 3 倍
排序
- MLX 的排序非常快
- 在 M2 Max、M2 Ultra 和 M3 Max 上，它的性能优于所有 CUDA GPU，包括最新的 RTX 4090
Conv2D
- MLX 比 MPS 慢 2-5 倍
- RTX 4090 在这个 op 上特别快
讨论
我们的基准测试得出了以下发现：
-
在大多数操作中，运行在苹果 M 芯片上的 MLX 始终优于使用 MPS 后端的 PyTorch。
-
CUDA GPU 不可避免地仍然比苹果 M 芯片快。
-
值得注意的是， MLX 在某些 op（如排序和线性）上表现出色 ，即使与 CUDA GPU 基准相比也显示出强大的效率。
-
MLX 的卷积 op 仍然特别慢 。似乎 MLX 团队目前正在努力提高这些 op 的速度。
注意：好奇的读者可能对查看基准测试中可用的全面 op 集感兴趣。
虽然 CUDA GPU 无疑更快，但重要的是要退一步考虑更广阔的图景。Nvidia 花了多年时间完善其 GPU 技术，达到了目前无与伦比的成熟度和性能水平。当你权衡 Nvidia GPU 和苹果 M 芯片之间的硬件复杂性、成熟度和成本差异时，它提供了一个重要的视角。
这就是 MLX 发挥作用的地方。在需要在个人 Mac 上运行大型模型的情况下，这个框架变得特别有益。使用 MLX，你可以在不需要额外昂贵 GPU 硬件的情况下实现这一目标。对于那些寻求直接在 Mac 设备上进行高效机器学习的人来说，这是一个改变游戏规则的方案。
值得注意的是，最近使用 MLX 直接在苹果 M 芯片上运行 LLM 推理越来越受欢迎。
统一内存改变了游戏规则
苹果 M 芯片在 GPU 领域以其 统一内存架构 脱颖而出，这种架构独特地利用 Mac 的全部 RAM 来运行模型。与传统的 GPU 设置相比，这是一个重大变化。
但使用苹果 M 芯片最显著的优势可能是消除了 CPU 和 GPU 之间的数据传输。这可能听起来像是一个小细节，但在机器学习项目的真实世界中，这些数据传输是延迟的臭名昭著的来源。我们之前的文章强调了这一方面，在训练期间考虑数据传输时间时显示了非常不同的基准测试结果：
这些结果显示了图卷积网络模型的每个 epoch 平均训练时间，该模型本质上由两个线性层组成。与本文中的基准测试不同，这个特定的基准测试评估了完整训练循环的平均运行时间，包括从 CPU 到 GPU 的数据传输时间。
这些结果带来了一个有趣的见解：当包括实际数据传输时间时，CUDA GPU 的性能明显变慢。这凸显了 数据传输对整体速度的重大影响 ，这是我们在当前基准测试中不能忽视的一个因素。
结论
这个基准测试让我们清楚地看到 MLX 与运行在 MPS 和 CUDA GPU 上的 PyTorch 相比表现如何。我们发现 MLX 在大多数操作中通常比 MPS 快得多，但也比 CUDA GPU 慢。
MLX 真正突出的地方是其统一内存，这消除了 CPU 和 GPU 之间耗时的数据传输需求。
在未来的文章中，我们计划在我们的基准测试中包括数据传输时间。这可能会提供对 MLX 整体效率更全面的理解。
感谢阅读！ 特别感谢 Ian Myers 发现了之前 CUDA 基准测试中的问题。
作者简介
作者：Tristan Bilot， Towards Data Science 专栏作家，基于 GNN 实现网络安全博士。 主要写作深度学习、编程和性能方面的内容。