在 Apple Silicon 平台上使用 MLX进行设备端机器学习基准测试

大型语言模型（LLM）及机器学习在日益普及的场景中广泛应用，催生了将其部署在笔记本和移动设备等小型设备上的研究热情。这为充分利用设备硬件的框架与方法提出了新需求。为此，MLX 框架应运而生，旨在针对 Apple Silicon 设备上的机器学习计算进行优化，简化研究、实验与原型设计。本文对 MLX 进行性能评估，聚焦于 transformer 模型的推理延迟。我们将 MLX 中不同 transformer 架构实现与 PyTorch 的性能进行比较。为此，我们构建了一个名为 MLX-Transformers 的框架，包含 MLX 中的各种 transformer 实现，并下载 PyTorch 模型检查点并转换为 MLX 格式。凭借 Apple Silicon 的先进架构与能力，MLX-Transformers 实现了对直接来自 Hugging Face 的 transformer 模型的无缝执行，无需频繁的模型检查点转换。我们的研究在两台 Apple Silicon MacBook 设备上对不同 transformer 模型进行基准测试，并与 NVIDIA CUDA GPU 进行比较。具体而言，我们比较了参数规模和模型检查点相同的模型的推理延迟性能。我们评估了 BERT、RoBERTa 和 XLM-RoBERTa 模型的性能，旨在扩展未来工作，纳入不同模态的模型，从而为更全面的评估 MLX 能力提供支持。结果表明，MLX 在实现高效且更易于 Apple 生态系统内的设备端 ML 应用方面具有潜力。