To build:
-
Clone the PyTorch repo as per instructions.
git clone --recursive https://github.com/pytorch/pytorch cd pytorch/ -
Checkout the following commit ID specified in xla/.torch_commit_id:
git checkout $(cat xla/.torch_commit_id) -
Clone the PyTorch/XLA repo:
git clone --recursive https://github.com/pytorch/xla.git -
Apply the
pytorch.patchto the currentxlacode. From within the pytorch source folder:patch -p1 < xla/pytorch.patch -
Currently PyTorch does not build with GCC 6.x, 7.x, and 8.x (various kind of ICEs). CLANG 7.x is known to be working, so install that in your VM:
sudo apt-get install clang-7 clang++-7 export CC=clang-7 CXX=clang++-7 -
Build PyTorch from source following the regular instructions.
python setup.py install -
Build the PyTorch/XLA source:
cd xla/ python setup.py install
To run the tests, follow one of the options below:
-
Run on local CPU using the XRT client:
export XLA_USE_XRT=1 XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0" export XRT_WORKERS="localservice:0;grpc://localhost:40934"Select any free TCP port you prefer instead of 40934 (totally arbitrary).
-
Run on Cloud TPU using the XRT client, use one of the following:
-
Set the XRT_TPU_CONFIG environment variable:
export XRT_TPU_CONFIG="tpu_worker;0;<IP of the TPU node>:8470" -
Create a
$HOME/.pytorch_tpu.conffile with the following content:worker: tpu_worker <IP of the TPU node>:8470
-
Note that the IP of the TPU node can change if the TPU node is reset. If PyTorch seem to hang at startup, verify that the IP of your TPU node is still the same of the one you have configured.
Then run python test/test_operations.py. Some of the tests are currently skipped.