Support train deep learning models on Apple Silicon GPU?

[overheat]

Feature description

Support train deep learning models on Apple Silicon GPU? M1 and M2

Feature motivation

Faster

(Optional) Suggest a Solution

[antimora]

It's supported but there is a bug in tch-rs which prevents from automatically installing torch C++ libs (see the issue #180). However, we checked in a work around script. Please see my comment on how to run it: #180 (comment). This is available on main branch currently.

Let us know if you're having trouble setting it up. I develop and test on Mac M1 pro and I created this workaround script. So I can confirm Burn with tch backend works.

Also to note: we enabled the accelerate framework support (for MacOS and iOS) for burn-ndarray (see #183). There is a "blas-accelerate" flag in burn-ndarray tat you can use to build.

[overheat]

@antimora Thanks.

When I try this script, it shows:

✗ python scripts/fix-tch-build-arm64.py                                 
Installing/Upgrading torch via pip install ...
Requirement already satisfied: torch in /opt/homebrew/lib/python3.11/site-packages (1.13.0a0+git49444c3)
Requirement already satisfied: typing_extensions in /opt/homebrew/lib/python3.11/site-packages (from torch) (4.5.0)
Requirement already satisfied: tomli in /opt/homebrew/lib/python3.11/site-packages (2.0.1)
Requirement already satisfied: tomli-w in /opt/homebrew/lib/python3.11/site-packages (1.0.0)
Updating config.toml with torch library paths ... 
Traceback (most recent call last):
  File "/Volumes/Acrucial/repos/ai/burn/scripts/fix-tch-build-arm64.py", line 71, in <module>
    main()
  File "/Volumes/Acrucial/repos/ai/burn/scripts/fix-tch-build-arm64.py", line 67, in main
    update_toml_config()
  File "/Volumes/Acrucial/repos/ai/burn/scripts/fix-tch-build-arm64.py", line 26, in update_toml_config
    import tomli
ModuleNotFoundError: No module named 'tomli'

But I have tomli

python3 -m pip install tomli
Requirement already satisfied: tomli in /opt/homebrew/lib/python3.11/site-packages (2.0.1)

BTW, conda is actived. I think that is ok. Right?

[overheat]

BTW, conda is actived. I think that is ok. Right?

Actrully, it does.
I need this:

conda install tomli
conda install tomli-w

Inspired by this article

virtual environment in which you can install custom Python modules. But the virtual environment is initially empty—even if you’ve already installed tomli on your computer!

[overheat]

On my Macbook pro 13' 2020(M1), example mnist one iteration takes about 2 minutes with ndarray-blas-accelerate features. In contract, about 4 minutes with ndarray features.

@antimora How about on your Mac M1 pro for example mnist?

[antimora]

@overheat yes, with accelerate it takes about this long. This is because NDArray is not optimized as torch yet. @nathanielsimard made a few optimizations so you'd want to update the main.

However, currently burn-tch backend with CPU and GPU on Mac is fast. Try with "tch-gpu" or "tch-cpu" feature flags. You'll notice for some reason CPU is faster for this MNIST example, but only with the current model. It appears the model is small enough, it completes under one minute on CPU. Probably the overhead copying data to GPU is more. When I updated with a bigger model, GPU was faster. You can view activity monitor on your mac to see how much GPU is busy, like this:

[antimora]

BTW, conda is actived. I think that is ok. Right?

Actrully, it does. I need this:

conda install tomli
conda install tomli-w

Inspired by this article

virtual environment in which you can install custom Python modules. But the virtual environment is initially empty—even if you’ve already installed tomli on your computer!

Thanks for reporting. Probably there were some differences in the installation. I use pip and I didn't test with conda. Python packages installations is always difficult for different environments compared to rust =D

[antimora]

@overheat Were you able to train your models on Mac M1/M2 GPU? Can we consider this issue closed?

[antimora]

@nathanielsimard we should close this as duplicate #180. I don't have ability to close this ticket.