MPS-Ready, ARM64 Docker Image
MBAnslow opened this issue · comments
🚀 The feature, motivation and pitch
It'd be very helpful to release an ARM64 pytorch docker image for running pytorch models with docker on M1 chips natively using the MPS backend.
I test and debug prototypes based on pytorch locally during development. This was totally manageable using previous intel-based macs, but now that my pytorch docker images are running in AMD64 emulation mode, it's so slow that I can't even debug and test my prototypes locally anymore.
I have tried out using the MPS backend in a virtual environment outside of docker and it's impressively fast! It's actually fast enough to plausibly run prototypes locally for demo purposes and massively speed up testing without throttling the model inference computations which is a real game changer. However, this isn't implemented in an official pytorch image yet. It'd be really great if it were.
Alternatives
No response
Additional context
No response
Are you able to get pip to find torch in a docker container?
I have been experiencing issues with getting python3.7 pip and m1 docker to play nicely together. I keep getting this error message "Could not find a version that satisfies the requirement torch==1.4.0"
I was able to install pytorch outside of the container, however.
Thanks @MBAnslow for starting this thread.
but now that my pytorch docker images are running in AMD64 emulation mode, it's so slow that I can't even debug and test my prototypes locally anymore.
I am not too familiar with the Docker workflow. Can you please provide more details about what your setup is and what is needed to provide PyTorch docker image ?
https://docs.docker.com/desktop/mac/apple-silicon/
These docker builds go through emulation or there is native support available?
I am experiencing the same issue. Currently, it looks like there are only linux/amd64 images on the pytorch dockerhub. To get this to run on Mac M1, I need to use the --platform linux/amd64 to force AMD64 emulation. I think the solution would be to also build and publish a linux/arm64 image to dockerhub. There is also the option of a multiarch image which docker seems to expect everyone to transition to.
In summary, running Intel-based containers on Arm-based machines should be regarded as “best effort” only. We recommend running arm64 containers on Apple silicon machines whenever possible, and encouraging container authors to produce arm64, or multi-arch, versions of their containers. We expect this issue to become less common over time, as more and more images are rebuilt supporting multiple architectures.
From https://docs.docker.com/desktop/mac/apple-silicon/#known-issues
I believe this issue combines 2 steps, which are currently missing in pytorch, but are really needed:
- Make pytorch docker images multiarch - this is crucial and needed for anything that builds on top of pytorch images (many apps). Requested here #80764
- Support Apple's MPS (Apple GPUs) in pytorch docker image. Image is based on ubuntu, so currently MPS is not supported there (at least I was not able to make it work), as the official pytorch guide says "MPS is supported only on latest Mac OS". So even if the latest Mac OS is hosting the pytorch docker image, the image itself can't use the power of Apple chips.
Both are very important because Apple is not selling Intel CPUs anymore, they were replaced with Apple arm chips --> Docker pytorch images (and their derivatives) on Mac are extremely slow (practically unusable), because it has to emulate amd64 architecture on arm64 + it can't use the GPU cores (via MPS).
GPUs other than Nvidia GPUs are not supported in Docker, not even through OpenCL. nvidia-docker is all we have regarding GPUs in Docker. There is of course ROCm support, but that's only for AMD's Frontier supercomputer and extremely expensive accelerators. Perhaps you could set up a macOS virtual machine capable of running Metal with some other kind of software?
Actually, DirectX supports non-Nvidia GPUs in WSL, so Intel and AMD GPUs might work in Docker.
Nvm, it seems Nvidia also has a monopoly over Docker + WSL: https://docs.docker.com/desktop/windows/wsl/#gpu-support
Hmm, I understand the need to regular linux-arm64 docker, but I'm quite confused how it can be integrated with MPS, which is a feature of Mac OS.
any updates?
MPS (Metal Performance Shaders) is MacOS specific feature and as such not available on Linux-ARM64 (even if runs on Apple Silicon)
Closing the issue for now, but please do not hesitate to file a new one, once MPS support on Linux would become available.
P.S. Please note, that Linux/aarch64 docker images with PyTorch are available, but those are good only for CPU workflows.
https://rosenzweig.io/blog/first-conformant-m1-gpu-driver.html
There are now MPS drivers for Asahi Linux
Following.
(yes, I know there is a button to Subscribe to the Issue, but I also want to show there is still interest in making Pytorch GPU in Docker for M1 work).