Add Apple MPS GPU support for Pytorch

[kyouma]

There is a Pytorch bug that occurs when an optimizer is created for a network which is on an MPS GPU (pytorch/pytorch#149184), thus a trick to avoid it is needed.

[kyouma]

Excuse me, I do not know how the build process is organized. The error message says:

RuntimeError: MPS backend out of memory (MPS allocated: 0 bytes, other allocations: 0 bytes, max allowed: 7.93 GB).
Tried to allocate 256 bytes on shared pool.
Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure).

Maybe the PYTORCH_MPS_HIGH_WATERMARK_RATIO parameter is enough or maybe this Pytorch forum link ("Switching runs-on from macos-latest to macos-13") contains the way to solve the build issue.

[lululxvi]

I have no experience with Apple MPS GPU. The build file is here https://github.com/lululxvi/deepxde/blob/master/.github/workflows/build.yml

[kyouma]

Hello. It seems that the GitHub Actions MPS-based environment doesn't support the necessary virtualization technology. As a result, PyTorch can detect the MPS GPU but cannot access it.

I have added a workaround to fall back to the CPU if an exception occurs. Now, the line that addresses the PyTorch optimizer bug checks for GPU availability. If the optimizer bug is fixed and that line is removed, running a few neural network layers will help verify the MPS functionality.

For future reference, please see the following links:

• https://docs.github.com/en/actions/using-github-hosted-runners/using-github-hosted-runners/about-github-hosted-runners#supported-runners-and-hardware-resources
• GitHub Action: New M1 runner available to all plans, including open source 🚀 actions/runner-images#9254
• [MNT] MPS backend test failures on MacOS sktime/pytorch-forecasting#1596
• [MNT] handle mps backend for lower versions of pytorch and fix mps failure on macOS-latest runner sktime/pytorch-forecasting#1648
• https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773/7
• https://github.blog/changelog/2024-01-30-github-actions-introducing-the-new-m1-macos-runner-available-to-open-source

Edit:

Setting PYTORCH_MPS_HIGH_WATERMARK_RATIO to 0.0, as proposed by the Pytorch exception message before this workaround, didn't help.

One more link:

• Test failure on macos-14 CI (M1 Mac, MPS torch backend): test_mps_batched_mask_to_box computational-cell-analytics/micro-sam#380