MetalML: machine learning on Apple Silicon
Bringing Metal GPU acceleration to familiar scikit-learn workflows.
MetalML started with a practical idea: keep the familiar scikit-learn workflow while making supported numerical operations available on a Mac’s GPU. Instead of introducing a separate training framework, it exposes estimators with the fit, predict and transform methods that Python users already know.
The project combines native Metal compute kernels and Metal Performance Shaders with scikit-learn compatibility. It is an experimental library, with acceleration chosen separately for each supported operation.
A familiar workflow, a different backend
The entry point is the Python package. Existing workflows can use MetalML estimators in pipelines, cross-validation and grid search. Public estimators that the library does not implement are forwarded to the installed scikit-learn. Importantly, importing MetalML does not modify existing scikit-learn estimators or patch their internal factories.
python -m pip install metalml
import numpy as np
import metalml as ml
rng = np.random.default_rng(7)
X = rng.normal(size=(5000, 32)).astype(np.float32)
y = X[:, 0] * 3 - X[:, 1]
model = ml.linear_model.Ridge(alpha=1.0)
model.fit(X, y)
print(model.predict(X[:5]))
print(model.operation_backends_)What actually runs on the GPU?
Acceleration depends on the estimator and the operation. Some algorithms can use the GPU during fitting; others accelerate prediction while retaining CPU training. A single pipeline may therefore contain both CPU and GPU steps. The implementation uses tiled kernels, fused operations and reusable buffers rather than moving every part of the workflow onto the GPU.
- Scaling and normalisation can accelerate statistics and element-wise operations.
- Ridge uses GPU matrix products alongside a CPU linear-system solve.
- Eligible KMeans configurations can accelerate Lloyd iterations and distance calculations.
- Trees, forests and kernel SVMs can accelerate prediction while fitting remains on the CPU.
Precision and fallback are part of the design
GPU operations use dense float32 arrays. Sparse inputs and precision-preserving float64 or integer inputs generally use the scikit-learn path. This matters because NumPy often creates float64 arrays by default. Converting to float32 should be a deliberate numerical choice, rather than something that happens invisibly. GPU results can differ slightly from CPU results, particularly near classification boundaries or tied distances.
Inspect the work instead of assuming a speedup
Backend diagnostics make the execution path visible. The operation_backends_ attribute records which backend was used, while fallback_reasons_ helps explain CPU execution. Configuration contexts can select CPU or Metal execution for comparisons. Estimator boundaries return NumPy arrays, so a complete pipeline is not kept resident on the GPU.
with ml.config_context(backend="cpu"):
baseline = ml.linear_model.Ridge().fit(X, y)
print(ml.backend_info())
print(model.operation_backends_)
print(model.fallback_reasons_)Current status and next steps
Apple Silicon is the target platform, but the repository explicitly notes that physical Apple Silicon validation and performance tuning are still pending. GPU execution has been verified on an Intel Mac with a Radeon Pro 5300M. The automatic dispatch rules are based on that hardware and are not calibrated for Apple Silicon. Small workloads may be faster on the CPU, so performance should be measured on the actual device and dataset. The project is released under Apache 2.0, with further work around hardware validation, additional GPU training paths and reducing data movement between pipeline steps.