Arm64 Support For Pytorch?

[jxtps]

I noticed that org.bytedeco CUDA now contains cuda-11.6-8.3-1.5.7-linux-arm64-redist.jar, so it sounds like you've compiled CUDA to run on arm64?

Any chance of also compiling pytorch for arm64 (ideally both for CPU and GPU)?

That would enable using e.g. AWS G5g instances: https://aws.amazon.com/ec2/instance-types/g5g/

Thanks!

[saudet]

I'm sure there's some way to get that working, we just need to spend some time on it. As usual, contributions are welcome!

[jxtps]

Ok, so there's a smattering of resources out there referencing various builds for various ARM architectures:

1. https://discuss.pytorch.org/t/crosscompilation-for-arm64/9885/9
   2. Leads to: https://gist.github.com/soumith/45f61db3379db67414814ad9d5f2a427 - a gist, "Install PyTorch on NVIDIA AGX Xavier using JetPack 4.1"
   3. Leads to https://forums.developer.nvidia.com/t/pytorch-for-jetson-version-1-10-now-available/72048 - NVIDIA Jetson (an ARM computer) pytorch pip wheels with Build From Source instructions.
2. https://nmilosev.svbtle.com/compling-arm-stuff-without-an-arm-board-build-pytorch-for-the-raspberry-pi - Build PyTorch for the Raspberry Pi
3. https://github.com/KumaTea/pytorch-aarch64 - PyTorch, vision, audio, text and csprng wheels (whl) and docker images for aarch64 / ARMv8 / ARM64 devices
4. https://mathinf.eu/pytorch/arm64/ - PyTorch packages for ARM64. Jetson Build with CUDA. At the Developers Day, PyTorch announced the availability of official ARM64 packages for the nightly (development version) builds.
5. https://schweter.eu/2017/11/pytorch-arm64.html - PyTorch on ARM64 (no CUDA)
6. https://qengineering.eu/install-pytorch-on-raspberry-pi-4.html - This page will guide you through the installation of PyTorch 1.11.0, or an earlier version, TorchVision, LibTorch and Caffe2 on a Raspberry Pi 4 with a 64-bit operating system.

[saudet]

There's no official guide, so we have to spend a bit of time modifying the build script...

[saudet]

@frankfliu @zachgk @lanking520 How important are ARM builds for Amazon?

@jxtps See discussion deepjavalibrary/djl#1550 for some context

[lanking520]

@saudet not sure for Amazon, but we do have some customers in DJL are looking for support on that.

[jxtps]

It looks like PyTorch is publishing ARM wheels and have been for some time: pytorch/pytorch#53357

If you search for "aarch64" on https://download.pytorch.org/whl/torch_stable.html it looks like recent versions are all there.

Could whatever build scripts they're using to produce those wheels help here?

[lanking520]

DJL has already published AARCH64 for PyTorch: http://docs.djl.ai/engines/pytorch/pytorch-engine/index.html#for-aarch64-build

[jxtps]

Does that include GPU support, e.g. for use with the AWS G5g instance type?

[saudet]

Could whatever build scripts they're using to produce those wheels help here?

If it can cross compile, we can integrate that right away.

[jxtps]

Nevermind, now that I look closer at https://download.pytorch.org/whl/torch_stable.html they don't have anything cross-compiled for cuda (cu113, cu100 etc), it's all just for cpu. I take it that's not helpful here?

[saudet]

The CUDA builds also compile the libraries for CPU, so we'll need to get that working first anyway.

[jxtps]

Ok, the main project for the build scripts appears to be: https://github.com/pytorch/builder

Of particular relevance:

• https://github.com/pytorch/builder/tree/main/conda - files to build conda packages of pytorch, torchvision and other dependencies and repos
• https://github.com/pytorch/builder/tree/main/manywheel - the scripts to build linux wheels

There are multiple scripts in there. From what I can tell build_all_docker.sh -> build_docker.sh -> [invoke docker]. The various Dockerfiles then take a myriad of actions and perform the actual builds using various scripts that I haven't deciphered.

[saudet]

Thanks! I'm not seeing anything about cross compiling for linux-arm64 though. They seem to be building everything on ARM hardware, just like DJL does, but unfortunately, I don't have the resources of Amazon and Facebook, like these guys do, to do builds on ARM hardware.

[lanking520]

You can try to use Mac OS docker build to build ARM64, or use ARM simulator to initiate docker build

[saudet]

Yeah, we could try QEMU and what not, someone would have to try that stuff out...