Skip to content

PyTorch support with VEC Backend for SVE512

Latest
Compare
Choose a tag to compare
@maajidkhann maajidkhann released this 05 Jul 05:39
· 1048 commits to main since this release

This release supports official OSS PyTorch with added support for VEC backend with SVE512 support. This can be leveraged by Fugaku CPU A64FX nodes for faster computation.

Settings on FX700 for software build with PyTorch: (The Python wheel file is attached in the assets section of this release)
OneDNN v3.5.0
gcc (GCC) 10.3.0
Rocky Linux 8.10
Model name: A64FX (48 core)
Pytorch tag: v2.4.0-rc2 + SVE512 A64FX changes
ACL (Not used)

Installing Dependencies:
1. Installing openblas:

Install openblas-openmp rpm package:

dnf install openblas-openmp

Installed at:
/lib64/libopenblaso.so.0

2. Installing cmake:

sudo dnf install cmake

  1. Installing python 3.11:

sudo dnf install python3.11

Building custom OneDNN:
Step 1: Clone source code:

git clone https://github.com/oneapi-src/oneDNN.git

Step 2: Checkout release branch:

cd oneDNN
git checkout <custom_branch_name>

Step 3: Build:

mkdir build
cd build
cmake ..

Building Pytorch with custom OneDNN:

git clone --recursive https://github.com/MonakaResearch/pytorch-a64fx.git
cd pytorch
git checkout vec_backend_sve512
git submodule sync
git submodule update --init --recursive

# Install Dependencies:
python -m pip install cmake ninja (if copy paste does not work type the command)

# Run this command from the PyTorch directory after cloning the source code

pip3 install -r requirements.txt

# To build Pytorch with custom OneDNN:

# Below folder contains source code of default OneDNN, we need to replace only those files which we have updated.

cd third_party/ideep/mkl-dnn/
rm -rf * (Deletes all files from here)
cp -r /OneDNN/* ./ (copy all the files from your custom OneDNN build to current directory)

# After replacing modifies files, run build command from pytorch source folder (/pytorch)

export LD_PRELOAD=/usr/local/lib64/libstdc++.so.6.0.28:/lib64/libopenblaso.so.0

NOTE:
/usr/local/lib64/libstdc++.so.6.0.28 ----> is to avoid GLIBC missing error.
/lib64/libopenblaso.so.0 -----> To include OpenBLAS in the build

Builds torch:

USE_CUDA=0 USE_MKLDNN=ON python setup.py bdist_wheel

The above command this will generate wheel file in dist folder.

Installing the wheel file:

pip3 install dist/torch-2.4.0a0+git5b6d940-cp311-cp311-linux_aarch64.whl

How to Install the torchvision compatible package w.r.t the above torch package:
The official OSS version of torch package currently available is: torch 2.3.1 (https://pypi.org/project/torch/) and its compatible torchvision version is: torchvision 0.18.1 (https://pypi.org/project/torchvision/).

However, in the release we did, we are using the latest PyTorch v2.4.0-rc2 tag + our SVE512 changes to generate the wheel file (https://github.com/MonakaResearch/pytorch-a64fx/releases/download/a64fx_rel-1/torch-2.4.0a0+git5b6d940-cp311-cp311-linux_aarch64.whl) as we wanted to make the latest code already available to the users.

so, in this case, using the OSS torchvision 0.18.1 package with our latest torch 2.4.0a0 package would be incompatible.
You would need to use the latest torchvision package which you can build yourself within few minutes.

How to Install latest torchvision package:
1.Activate your python virtual environment.

2.Uninstall the already existing default torchvision package.

pip3 uninstall torchvision

3. Now do the following steps:

git clone https://github.com/pytorch/vision.git
cd vision
python3 setup.py install

Now, you will get a message:
Finished processing dependencies for torchvision==0.20.0a0+8f9d810

This will install the latest torchvision package in your python virtual environment and this would be compatible
with our torch package installed.

NOTE:
I also have directly uploaded the compatible torchvision wheel file now in the release assets section. You can skip the above steps and directly download and install the torchvision package from here. Also refer to requirements.txt file to know the right version of the packages to be installed.