MPS Accelerator

[justusschock]

What does this PR do?

Adds a prototype of M1 GPU support via PyTorch's MPS backend (part of #13102 , renaming existing accelerators will be done separately)

To give it a try, install PL from this branch and change the accelerator to 'mps'

Remaining ToDos:

• Extend Tests
• Docs
• Getting tests to pass on all platforms
• Move _MPS_AVAILABLE to accelerator file
• Integrate MPS into GPU debug prints

Does your PR introduce any breaking changes? If yes, please list them.

Nope

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

cc @Borda @akihironitta @rohitgr7 @justusschock

[johnnynunez]

I'm testing this! And works fine! It is cool to use. The problem now is the pytorch backend that most of features hadn't been implemented.
pytorch/pytorch#77764

[justusschock]

@johnnynunez Thanks for testing it. You can still use the cpu fallback with PYTORCH_ENABLE_MPS_FALLBACK=1 python your_script.py . That will move the computation for unsupported ops to cpu, but luckily due to the shared memory concept it should be very cheap moving the actual tensors

[johnnynunez]

@johnnynunez Thanks for testing it. You can still use the cpu fallback with PYTORCH_ENABLE_MPS_FALLBACK=1 python your_script.py . That will move the computation for unsupported ops to cpu, but luckily due to the shared memory concept it should be very cheap moving the actual tensors

thank you! I will test it. I have M1 Max with 32gb

still is in beta:
/Users/johnny/Projects/DS_TFM/src/utils/models/metrics.py:85: UserWarning: The operator 'aten::index.Tensor' is not currently supported on the MPS backend and will fall back to
run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
output, target = output[:, :, ..., dimensions_idces], target[:, :, ..., dimensions_idces]
loc("mps_multiply"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/b6051351-c030-11ec-96e9-3e7866fcf3a1/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":219:0)): error: input types 'tensor<64x1x1x1x2xf32>' and 'tensor<64x1x1x1x2xi1>' are not broadcast compatible
LLVM ERROR: Failed to infer result type(s).

[Borda]

@justusschock mind checking the collision? (in git there are much smaller, don't be scared by the GH message) I think we are ready to merge it... :)

[justusschock]

let me see if I can get to it tomorrow :)

[rasbt]

Just saw this finally got merged. Wohoo! Great work @justusschock & everyone! 🎉

[saikatkumardey]

@johnnynunez were you able to fix the issue? I have a similar error while trying to use AveragePrecision() on Mac M1 Pro.

By debugging it further I see that this line generates the following error.

loc("mps_subtract"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/b6051351-c030-11ec-96e9-3e7866fcf3a1/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":219:0)): error: input types 'tensor<128xi64>' and 'tensor<128xf32>' are not broadcast compatible

[awaelchli]

@saikatkumardey Thanks, since you are mentioning AveragePrecision, I suggest to open an issue directly on https://github.com/Lightning-AI/metrics. Please note, not all functions in PyTorch are implemented for MPS.