install torch torch_xla prerelease 2.1 wheel

tests/pytorch/r2.1/common.libsonnet

@@ -96,10 +96,13 @@ local volumes = import 'templates/volumes.libsonnet';
         sudo apt install -y libopenblas-base
         # for huggingface tests
         sudo apt install -y libsndfile-dev
-        pip install --user \
-          https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.1-cp310-cp310-linux_x86_64.whl \
-          'torch_xla[tpuvm] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.1-cp310-cp310-linux_x86_64.whl'
-        pip3 install --user --pre --no-deps torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        # TODO change back to torch2.1 once pytorch released torch2.1
+        # pip install --user \
+        #   https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.1-cp310-cp310-linux_x86_64.whl \
+        #   'torch_xla[tpuvm] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.1-cp310-cp310-linux_x86_64.whl'
+        pip3 install --user --pre --no-deps torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230825-cp310-cp310-linux_x86_64.whl
+        pip install torch_xla[tpuvm]
         pip3 install pillow
         git clone --depth=1 -b release/2.1 https://github.com/pytorch/pytorch.git
         cd pytorch
@@ -129,10 +132,13 @@ local volumes = import 'templates/volumes.libsonnet';
         sudo apt install -y libopenblas-base
         # for huggingface tests
         sudo apt install -y libsndfile-dev
-        pip install --user \
-          https://storage.googleapis.com/pytorch-xla-releases/wheels/xrt/tpuvm/torch-2.1-cp310-cp310-linux_x86_64.whl \
-          'torch_xla[tpuvm] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/xrt/tpuvm/torch_xla-2.1-cp310-cp310-linux_x86_64.whl'
-        pip3 install --user --pre --no-deps torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        # TODO change back to torch2.1 once pytorch released torch2.1
+        # pip install --user \
+        #   https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.1-cp310-cp310-linux_x86_64.whl \
+        #   'torch_xla[tpuvm] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.1-cp310-cp310-linux_x86_64.whl'
+        pip3 install --user --pre --no-deps torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230825-cp310-cp310-linux_x86_64.whl
+        pip install torch_xla[tpuvm]
         pip3 install pillow
         git clone --depth=1 -b release/2.1 https://github.com/pytorch/pytorch.git
         cd pytorch

tests/pytorch/r2.1/hf-mae.libsonnet

@@ -122,18 +122,6 @@ local utils = import 'templates/utils.libsonnet';
       },
     },
   },
-  // TODO: Remove after 2.1 release cut
-  local xrt = self.xrt,
-  xrt:: common.XrtTpuVmMixin {
-    tpuSettings+: {
-      tpuVmExports+: |||
-        export XLA_USE_BF16=$(XLA_USE_BF16)
-      |||,
-      tpuVmExtraSetup: |||
-        echo 'export XLA_USE_BF16=1' >> ~/.bash_profile
-      |||,
-    },
-  },
   local v2_8 = self.v2_8,
   v2_8:: {
     accelerator: tpus.v2_8,
@@ -147,7 +135,5 @@ local utils = import 'templates/utils.libsonnet';
     accelerator: tpus.v4_8,
   },
   configs: [
-    hf_mae + v3_8 + hf_vit_mae + timeouts.Hours(2) + xrt,
-    hf_mae + v4_8 + hf_vit_mae + timeouts.Hours(2) + xrt,
   ],
 }

tests/pytorch/r2.1/llama2-model.libsonnet

@@ -58,9 +58,13 @@ local utils = import 'templates/utils.libsonnet';
         pip3 install numpy
         sudo apt-get install numactl -y
         sudo apt-get install libopenblas-dev -y
-        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp310-cp310-linux_x86_64.whl
-        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl
-        pip3 install torch_xla[tpuvm]
+        # TODO change back to torch2.1 once pytorch released torch2.1 from current used pre-release wheel
+        # pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp310-cp310-linux_x86_64.whl
+        # pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl
+        # pip3 install torch_xla[tpuvm]
+        pip3 install --user --pre --no-deps torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230825-cp310-cp310-linux_x86_64.whl
+        pip install torch_xla[tpuvm]
 
         # install tokenizer model
         wget https://storage.googleapis.com/tpu-pytorch/lsiyuan-experiment/llama/spiece.model
@@ -107,9 +111,13 @@ local utils = import 'templates/utils.libsonnet';
         pip3 install numpy
         sudo apt-get install numactl -y
         sudo apt-get install libopenblas-dev -y
-        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp310-cp310-linux_x86_64.whl
-        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl
-        pip3 install torch_xla[tpuvm]
+        # TODO change back to torch2.1 once pytorch released torch2.1
+        # pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp310-cp310-linux_x86_64.whl
+        # pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl
+        # pip3 install torch_xla[tpuvm]
+        pip3 install --user --pre --no-deps torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230825-cp310-cp310-linux_x86_64.whl
+        pip install torch_xla[tpuvm]
 
         # install tokenizer model
         wget https://storage.googleapis.com/tpu-pytorch/lsiyuan-experiment/llama/spiece.model

tests/pytorch/r2.1/sd-model.libsonnet

@@ -62,11 +62,15 @@ local utils = import 'templates/utils.libsonnet';
         # mv quantize.py ~/.local/lib/python3.8/site-packages/taming/modules/vqvae/
 
         # taming-transformers and CLIP override existing torch and torchvision so we need to reinstall
-        pip uninstall -y torch torchvision
-        pip install --user \
-          https://storage.googleapis.com/pytorch-xla-releases/wheels/xrt/tpuvm/torch-nightly-cp310-cp310-linux_x86_64.whl \
-          'torch_xla[tpuvm] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/xrt/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl'
-        pip3 install --user --pre --no-deps torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        # TODO change back to torch2.1 once pytorch released torch2.1
+        # pip uninstall -y torch torchvision
+        # pip install --user \
+        #   https://storage.googleapis.com/pytorch-xla-releases/wheels/xrt/tpuvm/torch-nightly-cp310-cp310-linux_x86_64.whl \
+        #   'torch_xla[tpuvm] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/xrt/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl'
+        # pip3 install --user --pre --no-deps torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        pip3 install --user --pre --no-deps torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230825-cp310-cp310-linux_x86_64.whl
+        pip install torch_xla[tpuvm]
 
         # Setup data
         wget -nv https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz