I Tested Apple’s New MLX Framework Against Torch on M2 Air

Author(s): Tim Cvetko

Originally published on Towards AI.

MLX vs Torch on BERT — Training, Inference, and CPU Usage Comparison

On Tuesday, Apple’s AI team released “MLX” — the new machine learning framework designed to work specifically for the Apple Silicon Chips. The design of MLX was inspired by frameworks like NumPy, PyTorch, Jax, and ArrayFire.

Photo by Sumudu Mohottige on Unsplash

Is MLX really faster than Torch on Mac?

As I own a Macbook M2 Air and regularly train ML models locally, I decided to put this hypothesis to the test by training the standard BERT transformers model on both MLX and PyTorch. The results are staggering!

Who is this blog post useful for? Mac(M1, M2, M3) owners who are looking for a faster training & inference ML framework.

How advanced is this post? Anybody previously acquainted with ML terms should be able to follow along.

Replicate my code here: https://github.com/Timothy102/mlx-bert

MLX is an array framework for machine learning on Apple silicon, brought to you by Apple machine learning research.

MLX is very Torch-like in its syntax. MLX has higher-level packages like mlx.nn and mlx.optimizers with APIs that closely follow PyTorch to simplify building more complex models.MLX has a Python API that closely follows NumPy. MLX also has a fully featured C++ API, which closely mirrors the Python API.pip install mlxImage by Author

The MLX examples repo has… Read the full blog for free on Medium.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI