MPS (Mac M1) device support

[carmocca]

🚀 Feature

https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/
Docs: https://pytorch.org/docs/master/notes/mps.html

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @Borda @akihironitta @rohitgr7 @justusschock

[carmocca]

I heard @awaelchli just got an M1 😈

[awaelchli]

@justusschock also has one if I remember correctly 😃

[justusschock]

I'll investigate on the weekend

[carmocca]

After a bit of offline discussion, we thought about this API:

• The MPSAccelerator is introduced.
• The GPUAccelerator is deprecated and renamed in favor of CUDAAccelerator
• 2 new accelerator options: Trainer(accelerator="cuda"|"mps") to explicitly choose one of the above.
• Trainer(accelerator="gpu") now chooses from the 2 accelerators above based on the available hardware. (side note: this is how we would also support ROCm).
• Trainer(accelerator="auto") is different in that it also considers IPUs, TPUs, HPUs...

This shouldn't be a problem because a machine cannot have both accelerators available.

[eubinean]

Just when I thought I wanted to raise an issue! I'm so rooting for this.

[ananthsub]

I'm glad we did the accelerator refactor to make supporting new features like this much easier :)
GPUs is ambiguous, whereas the torch.device type is not: #10410 (comment)