Benchmarking On-Device Machine Learning on Apple Silicon with MLX

TL;DR

This study benchmarks on-device transformer inference on Apple Silicon using the MLX framework and MLX-Transformers library, comparing against NVIDIA CUDA GPUs. It conducts two latency-focused benchmarks (operations and model inference) across M1 and M2 Macs and an NVIDIA A10, evaluating BERT, RoBERTa, and XLM-RoBERTa with varying input lengths and batch sizes. The findings show CUDA generally outperforms MLX on Apple devices, though the M2 Max reduces the gap, demonstrating practical on-device inference potential for edge deployments and research prototyping. The work highlights the MLX ecosystem's promise for democratizing on-device AI and outlines avenues for optimizing larger models and enabling on-device training in future work.

Abstract

The recent widespread adoption of Large Language Models (LLMs) and machine learning in general has sparked research interest in exploring the possibilities of deploying these models on smaller devices such as laptops and mobile phones. This creates a need for frameworks and approaches that are capable of taking advantage of on-device hardware. The MLX framework was created to address this need. It is a framework optimized for machine learning (ML) computations on Apple silicon devices, facilitating easier research, experimentation, and prototyping. This paper presents a performance evaluation of MLX, focusing on inference latency of transformer models. We compare the performance of different transformer architecture implementations in MLX with their Pytorch counterparts. For this research we create a framework called MLX-transformers which includes different transformer implementations in MLX and downloads the model checkpoints in pytorch and converts it to the MLX format. By leveraging the advanced architecture and capabilities of Apple Silicon, MLX-Transformers enables seamless execution of transformer models directly sourced from Hugging Face, eliminating the need for checkpoint conversion often required when porting models between frameworks. Our study benchmarks different transformer models on two Apple Silicon macbook devices against an NVIDIA CUDA GPU. Specifically, we compare the inference latency performance of models with the same parameter sizes and checkpoints. We evaluate the performance of BERT, RoBERTa, and XLM-RoBERTa models, with the intention of extending future work to include models of different modalities, thus providing a more comprehensive assessment of MLX's capabilities. The results highlight MLX's potential in enabling efficient and more accessible on-device ML applications within Apple's ecosystem.

Figures (4)

• Figure 1: Comparison of operation runtimes (in milliseconds) between MLX on Apple M1 GPU and PyTorch on the CUDA GPU. Lower bars indicate better performance. Key transformer operations such as MatMul, Linear, and Softmax are highlighted.
• Figure 2: Comparison of average inference times (in milliseconds) for different models across CUDA GPU, Apple M1, and Apple M2 Max. Lower bars indicate better performance.
• Figure 3: Influence of input length on inference times for different models on Apple M2 Max. The x-axis represents input lengths, and the y-axis shows inference times in milliseconds.
• Figure 4: Effect of batch size on inference times for different models on Apple M2 Max. The x-axis represents batch sizes, and the y-axis shows inference times in milliseconds.