Is it possible to enable the GPU on an Apple M3 Max?

I’ve been using livebook for a while and been playing a bit with some of the numerics libraries. Recently I bought a mac apple silicon with Apple M3 Max GPU. I wonder if is possible to enable the GPU for machine learning. I have managed to train a model but when I check the GPU and CPU history only the CPU seems to be used. Maybe is not possible but better to be sure.

Thanks for the great work with Livebook!

I tried this a while back and found It’s possible using Torchx and setting the device to mps.

However some APIs are not supported yet and resulted in errors.

For example, taking the sample at nx/exla/examples/regression.exs at main · elixir-nx/nx · GitHub at changing the default backend:

Nx.default_backend({Torchx.Backend, device: :mps})

give this error when running it:

** (ArgumentError) cannot perform operation across devices mps and cpu
    (torchx 0.6.4) lib/torchx.ex:467: anonymous fn/2 in Torchx.prepare_tensors_list!/2

I’d be interested to hear if anyone has managed to get it to work.

1 Like

Thanks for your answer. I’ve been using the recommendation from the documentation:
{:exla, “~> 0.2”}

But did not have any effect. I can give a try to Torchx, just in case. Thanks!

@hiramegl yeah, EXLA does not support Metal at the moment. The underlying project (XLA) does not have official support for Metal, there is a separate plugin from Apple, but it is closed-source and already outdated, so we wait for an official support. As for Torchx, the mps backend is somewhat experimental, so depending on what computation you run, some operations may not be supported, but worth a shot.


Ok, thank you @jonatanklosko! i’m digging into TorchX and mps. We’ll see if I have some luck :slight_smile: