add v5p config for llama2-7b, llama2-70b, gpt3-175b (GoogleCloudPlatform#725)

uncomment models for v5e

add v5p-128

minor change

resolve format

fix variable name

fix variable name

Author: notabee

dags/common/model_configs.py

@@ -25,9 +25,18 @@ class MaxTextV5eModelConfigs(enum.Enum):
   DEFAULT_128B = "default_128b_v5e_256"
   GPT3_175B = "gpt_3_175b_v5e_256"
   LLAMA2_7B = "llama2_7b_v5e_256"
+  LLAMA2_13B = "llama2_13b_v5e_256"
   LLAMA2_70B = "llama2_70b_v5e_256"
 
 
+class MaxTextV5pModelConfigs(enum.Enum):
+  # Refers to model configs in https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/maxtext_v5p_model_configs.py
+  GPT3_175B = "gpt_3_175b_v5p_128"
+  GPT3_175B_SC = "gpt_3_175b_v5p_128_sc"
+  LLAMA2_7B = "llama2_7b_v5p_128"
+  LLAMA2_70B = "llama2_70b_v5p_128"
+
+
 class MaxTextTrilliumModelConfigs(enum.Enum):
   # Refers to model configs in https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/maxtext_trillium_model_configs.py
   GPT3_175B = "gpt_3_175b"

dags/common/test_owner.py

@@ -50,6 +50,7 @@ class Team(enum.Enum):
 MOHIT_K = "Mohit K."
 ANISHA_M = "Anisha M."
 YUWEI_Y = "Yuwei Y."
+RISHABH_B = "Rishabh B."
 
 # Multi-tier Checkpointing
 ABHINAV_S = "ABHINAV S."

dags/common/vm_resource.py

@@ -249,6 +249,13 @@ class XpkClusters:
       project=Project.CLOUD_TPU_MULTIPOD_DEV.value,
       zone=Zone.EUROPE_WEST4_B.value,
   )
+  TPU_V5P_128_CLUSTER = XpkClusterConfig(
+      name="v5p-128-bodaborg-europe-west4-b",
+      device_version=TpuVersion.V5P,
+      core_count=128,
+      project=Project.CLOUD_TPU_MULTIPOD_DEV.value,
+      zone=Zone.EUROPE_WEST4_B.value,
+  )
   TPU_V5E_256_CLUSTER = XpkClusterConfig(
       name="v5e-256-bodaborg-europe-west4",
       device_version=TpuVersion.V5E,

dags/multipod/maxtext_v5p_configs_perf.py

@@ -0,0 +1,87 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+A DAG to run perf tests for MaxText model configs on V5p.
+"""
+import datetime
+from airflow import models
+from airflow.utils.task_group import TaskGroup
+from dags import composer_env
+from dags.common import test_owner
+from dags.common.vm_resource import Project, XpkClusters, DockerImage
+from dags.common.model_configs import MaxTextV5pModelConfigs
+from dags.multipod.configs import maxtext_sweep_gke_config
+from dags.multipod.configs.common import SetupMode
+from xlml.apis import metric_config
+
+# Run once a day at 4 am UTC (8 pm PST / 9 pm PDT)
+SCHEDULED_TIME = "0 4 * * *" if composer_env.is_prod_env() else None
+DOCKER_IMAGES = [
+    (SetupMode.STABLE, DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK),
+    (SetupMode.NIGHTLY, DockerImage.MAXTEXT_TPU_JAX_NIGHTLY),
+]
+QUANTIZATION_SWEEP = {"M_QUANTIZATION": ["", "int8"]}
+BASE_OUTPUT_DIRECTORY = "gs://runner-maxtext-logs"
+
+with models.DAG(
+    dag_id="maxtext_v5p_configs_perf",
+    schedule=SCHEDULED_TIME,
+    tags=[
+        "multipod_team",
+        "maxtext",
+        "stable",
+        "nightly",
+        "mlscale_perfx",
+    ],
+    start_date=datetime.datetime(2024, 2, 19),
+    catchup=False,
+) as dag:
+  quarantine_task_group = TaskGroup(
+      group_id="Quarantine", dag=dag, prefix_group_id=False
+  )
+  for mode, image in DOCKER_IMAGES:
+    for model in MaxTextV5pModelConfigs:
+      base_run_model_cmds = [
+          "bash preflight.sh",
+          f"python3 -m benchmarks.benchmark_runner on-device --base_output_directory={BASE_OUTPUT_DIRECTORY} --model_name={model.value} --libtpu_type=maxtext-docker --num_steps=15",
+      ]
+      maxtext_sweep_gke_test = (
+          maxtext_sweep_gke_config.get_maxtext_sweep_gke_config(
+              test_owner=test_owner.RISHABH_B,
+              dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
+              composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
+              dataset_name=metric_config.DatasetOption.XLML_DATASET,
+              cluster=XpkClusters.TPU_V5P_128_CLUSTER,
+              time_out_in_min=360,
+              base_output_directory=BASE_OUTPUT_DIRECTORY,
+              num_slices=[1, 2],
+              docker_image=image.value,
+              run_name_prefix=f"maxtext-{model.name.lower()}-{mode.value}",
+              base_run_model_cmds=base_run_model_cmds,
+              sweep_params=QUANTIZATION_SWEEP,
+          )
+      )
+
+      chain_num = 4
+      prev = maxtext_sweep_gke_test[0].run_with_name_gen_and_quarantine(
+          quarantine_task_group
+      )
+      for i in range(1, len(maxtext_sweep_gke_test)):
+        curr = maxtext_sweep_gke_test[i].run_with_name_gen_and_quarantine(
+            quarantine_task_group
+        )
+        if i % chain_num != 0:
+          prev >> curr
+        prev = curr