Merge branch 'master' into yijiaj/mlperf-a2

GoogleCloudPlatform · Sep 30, 2024 · a57f750 · a57f750
2 parents 55921ec + 21f668e
commit a57f750
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 9 deletions.
diff --git a/dags/inference/jetstream_pytorch_inference.py b/dags/inference/jetstream_pytorch_inference.py
@@ -56,6 +56,23 @@
           "max_cache_length": 2048,
           "sharding_config": "default_shardings/llama.yaml",
       },
+      "gemma-7b": {
+          "model_name": "gemma",
+          "size": "7b",
+          "sleep_time": 120,
+          "tpu_version_cores": [(TpuVersion.V5E, 8)],
+          "checkpoint": "gs://inference-benchmarks/models/gemma-7b/pytorch/gemma-7b-merged-int8-per-channel",
+          "checkpoint_quantized": "gs://inference-benchmarks/models/gemma-7b/pytorch/gemma-7b-merged-int8-per-channel",
+          "dataset": "openorca",
+          "tokenizer": "tokenizer.model",
+          "batch_sizes": [8, 32, 64, 96, 128],
+          "request_rate": 100,
+          "num_prompts": 1000,
+          "quantize": [True],
+          "max_output_length": 1024,
+          "max_cache_length": 2048,
+          "sharding_config": "default_shardings/gemma.yaml",
+      },
       "llama2-13b": {
           "model_name": "llama-2",
           "size": "13b",

diff --git a/dags/legacy_test/tests/pytorch/r2.5/common.libsonnet b/dags/legacy_test/tests/pytorch/r2.5/common.libsonnet
@@ -24,7 +24,7 @@ local volumes = import 'templates/volumes.libsonnet';
     tpuSettings+: {
       softwareVersion: 'tpu-ubuntu2204-base',
     },
-    imageTag: 'r2.5.0-rc1_3.10',
+    imageTag: 'r2.5.0-rc5_3.10',
   },
   PyTorchTest:: common.PyTorchTest + r2_5 {
     local config = self,
@@ -99,12 +99,12 @@ local volumes = import 'templates/volumes.libsonnet';
         # Install torchvision by pinned commit in PyTorch 2.5 release branch.
         pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/test/cpu
         pip install --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@d23a6e1664d20707c11781299611436e1f0c104f"
-        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0rc1-cp310-cp310-linux_x86_64.whl
+        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0rc5-cp310-cp310-linux_x86_64.whl
         pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
         pip install pillow
         git clone --depth=1 https://github.com/pytorch/pytorch.git
         cd pytorch
-        git clone -b v2.5.0-rc1 https://github.com/pytorch/xla.git
+        git clone -b v2.5.0-rc5 https://github.com/pytorch/xla.git
       |||,
     },
     podTemplate+:: {
@@ -142,10 +142,10 @@ local volumes = import 'templates/volumes.libsonnet';
         pip uninstall -y torch torchvision
         pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/test/cpu
         pip install --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@d23a6e1664d20707c11781299611436e1f0c104f"
-        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0rc1-cp310-cp310-linux_x86_64.whl
+        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0rc5-cp310-cp310-linux_x86_64.whl
 
         mkdir -p pytorch/xla
-        git clone -b v2.5.0-rc1 https://github.com/pytorch/xla.git pytorch/xla
+        git clone -b v2.5.0-rc5 https://github.com/pytorch/xla.git pytorch/xla
 
         %s
 

diff --git a/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py b/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py
@@ -85,16 +85,16 @@ class R2_4(enum.Enum):
     TORCH_XLA_REPO_BRANCH = "-b v2.4.0-rc8"
 
   class R2_5(enum.Enum):
-    TORCH_XLA_TPU_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0rc1-cp310-cp310-linux_x86_64.whl"
-    TORCH_XLA_CUDA_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0rc1-cp310-cp310-linux_x86_64.whl"
+    TORCH_XLA_TPU_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0rc5-cp310-cp310-linux_x86_64.whl"
+    TORCH_XLA_CUDA_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0rc5-cp310-cp310-linux_x86_64.whl"
     TORCH = "torch==2.5.0"
     TORCHVISION = "torchvision==0.19.0"
     TORCHAUDIO = "torchaudio==2.5.0"
     TORCH_XLA_GPU_DOCKER = "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.5.0_3.10_cuda_12.1"
     TORCH_INDEX_CPU_URL = "https://download.pytorch.org/whl/test/cpu"
     TORCH_INDEX_CUDA_URL = "https://download.pytorch.org/whl/test/cu121"
-    TORCH_REPO_BRANCH = "-b v2.5.0-rc1"
-    TORCH_XLA_REPO_BRANCH = "-b v2.5.0-rc1"
+    TORCH_REPO_BRANCH = "-b v2.5.0-rc5"
+    TORCH_XLA_REPO_BRANCH = "-b v2.5.0-rc5"
 
 
 def get_version_mapping(test_version):