I Tested Apple’s New MLX Framework Against Torch on M2 Air
Last Updated on December 30, 2023 by Editorial Team
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
Towards AI Academy
We Build Enterprise-Grade AI. We'll Teach You to Master It Too.
15 engineers. 100,000+ students. Towards AI Academy teaches what actually survives production.
Start free — no commitment:
→ 6-Day Agentic AI Engineering Email Guide — one practical lesson per day
→ Agents Architecture Cheatsheet — 3 years of architecture decisions in 6 pages
Our courses:
→ AI Engineering Certification — 90+ lessons from project selection to deployed product. The most comprehensive practical LLM course out there.
→ Agent Engineering Course — Hands on with production agent architectures, memory, routing, and eval frameworks — built from real enterprise engagements.
→ AI for Work — Understand, evaluate, and apply AI for complex work tasks.
Note: Article content contains the views of the contributing authors and not Towards AI.
Related posts
Recent Posts
