Skip to content

Commit 4c8e672

Browse files
camiloCienetcamiloCienet
committed
Fix: Fixed Jetstream_benchmarking_servig DAG
- Point v5e tests to new reservation in CIENET-CMCS project - Point v6e tests to new reservation in tpu-prod-env-automated
1 parent d4c530a commit 4c8e672

8 files changed

Lines changed: 209 additions & 166 deletions

‎dags/common/vm_resource.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,8 @@ class Zone(enum.Enum):
124124
US_EAST5_C = "us-east5-c"
125125
# reserved v5e in tpu-prod-env-multipod
126126
US_WEST4_B = "us-west4-b"
127+
US_WEST4_A = "us-west4-a"
128+
US_EAST4_B = "us-east4-b"
127129
# reserved v5e in cloud-tpu-inference-test
128130
US_WEST1_C = "us-west1-c"
129131
# reserved a3+ cluster in supercomputer-testing

‎dags/inference/configs/jetstream_benchmark_serving_gce_config.py‎

Lines changed: 40 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value
2626
RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value
2727
GCS_SUBFOLDER_PREFIX = test_owner.Team.INFERENCE.value
28+
VENV_DIR = "venv-312"
2829

2930

3031
def get_config(
@@ -70,11 +71,14 @@ def get_config(
7071
# Make the PATH change permanent for subsequent sessions (optional, but good practice)
7172
"echo 'export PATH=\"$HOME/.local/bin:$PATH\"' >> ~/.bashrc",
7273
"source ~/.bashrc",
73-
"uv venv --python 3.12 venv-312 --seed",
74-
"source venv-312/bin/activate",
74+
f"uv venv --python 3.12 {VENV_DIR} --seed --clear",
75+
f"source {VENV_DIR}/bin/activate",
7576
"pip install uv",
7677
"uv pip install maxtext --resolution=lowest",
7778
"install_maxtext_github_deps",
79+
"uv pip install rouge-score",
80+
"sudo apt-get install -y git-lfs",
81+
"git lfs install",
7882
)
7983

8084
set_up_cmds += setup_maxtext_cmds
@@ -113,15 +117,19 @@ def get_config(
113117

114118
# Let gcs path be directly used, else use maxtext/assets dir
115119
if not model_configs["tokenizer"].startswith("gs://"):
116-
tokenizer_path = f"assets/{model_configs['tokenizer']}"
117-
full_tokenizer_path = f"maxtext/assets/{model_configs['tokenizer']}"
120+
tokenizer_path = (
121+
f"src/maxtext/assets/tokenizers/{model_configs['tokenizer']}"
122+
)
123+
full_tokenizer_path = (
124+
f"maxtext/src/maxtext/assets/tokenizers/{model_configs['tokenizer']}"
125+
)
118126
else:
119127
tokenizer_path = model_configs["tokenizer"]
120128
full_tokenizer_path = model_configs["tokenizer"]
121129

122130
run_model_cmds = (
123131
# Start virtual environment
124-
"source .env/bin/activate",
132+
f"source {VENV_DIR}/bin/activate",
125133
"wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json > /dev/null 2>&1",
126134
# Get commit hash of the maxtext and jetstream repos
127135
f"export METADATA_DICT='{json.dumps(additional_metadata_dict)}'",
@@ -131,58 +139,37 @@ def get_config(
131139
'export METADATA_DICT=$(jq -c \'. + { "jetstream_commit_hash": $newVal}\' --arg newVal ${JETSTREAM_COMMIT_HASH} <<<"$METADATA_DICT")',
132140
### Benchmark
133141
"cd maxtext",
134-
# Configure flags
135-
f"export MODEL_NAME={model_configs['model_name']}",
136-
f"export TOKENIZER_PATH={tokenizer_path}",
137-
f"export WEIGHT_DTYPE={model_configs['weight_dtype']}",
138-
f"export SCAN_LAYERS={model_configs['scan_layers']}",
139-
f"export MAX_PREFILL_PREDICT_LENGTH={model_configs['max_prefill_predict_length']}",
140-
f"export MAX_TARGET_LENGTH={model_configs['max_target_length']}",
141-
f"export ATTENTION={model_configs['attention']}",
142-
f"export ICI_FSDP_PARALLELISM={model_configs['ici_fsdp_parallelism']}",
143-
f"export ICI_AUTOREGRESSIVE_PARALLELISM={model_configs['ici_autoregressive_parallelism']}",
144-
f"export ICI_TENSOR_PARALLELISM={model_configs['ici_tensor_parallelism']}",
145-
f"export UNSCANNED_CKPT_PATH={model_configs['checkpoint']}",
146-
"export LOAD_PARAMETERS_PATH=${UNSCANNED_CKPT_PATH}",
147-
f"export QUANTIZATION={model_configs['quantization']}",
148-
f"export QUANTIZE_KVCACHE={model_configs['quantize_kvcache']}",
149-
f"export KV_QUANT_DTYPE={model_configs['kv_quant_dtype']}",
150-
f"export PER_DEVICE_BATCH_SIZE={model_configs['per_device_batch_size']}",
151-
f"export PREFILL_CACHE_AXIS_ORDER={model_configs['prefill_cache_axis_order']}",
152-
f"export AR_CACHE_AXIS_ORDER={model_configs['ar_cache_axis_order']}",
153-
f"export COMPUTE_AXIS_ORDER={model_configs['compute_axis_order']}",
154-
f"export RESHAPE_Q={model_configs['reshape_q']}",
155-
f"export KV_QUANT_AXIS={model_configs['kv_quant_axis']}",
156142
# Start JetStream MaxText server in the background
157-
"""python3 -m MaxText.maxengine_server \
158-
src/maxtext/configs/inference/inference_jetstream.yml \
159-
model_name=${MODEL_NAME} \
160-
tokenizer_path=${TOKENIZER_PATH} \
161-
weight_dtype=${WEIGHT_DTYPE} \
162-
scan_layers=${SCAN_LAYERS} \
163-
max_prefill_predict_length=${MAX_PREFILL_PREDICT_LENGTH} \
164-
max_target_length=${MAX_TARGET_LENGTH} \
165-
attention=${ATTENTION} \
166-
ici_fsdp_parallelism=${ICI_FSDP_PARALLELISM} \
167-
ici_autoregressive_parallelism=${ICI_AUTOREGRESSIVE_PARALLELISM} \
168-
ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \
169-
load_parameters_path=${LOAD_PARAMETERS_PATH} \
170-
quantization=${QUANTIZATION} \
171-
quantize_kvcache=${QUANTIZE_KVCACHE} \\"""
172-
+ (
173-
"""kv_quant_dtype=${KV_QUANT_DTYPE} \\"""
174-
if model_configs["kv_quant_dtype"]
175-
else ""
176-
)
177-
+ """per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
178-
prefill_cache_axis_order=${PREFILL_CACHE_AXIS_ORDER} \
179-
ar_cache_axis_order=${AR_CACHE_AXIS_ORDER} \
180-
compute_axis_order=${COMPUTE_AXIS_ORDER} \
181-
reshape_q=${RESHAPE_Q} \
182-
kv_quant_axis=${KV_QUANT_AXIS} &""",
143+
f"""python3 -m MaxText.maxengine_server \\
144+
src/maxtext/configs/inference/inference_jetstream.yml \\
145+
model_name='{model_configs['model_name']}' \\
146+
tokenizer_path='{tokenizer_path}' \\
147+
weight_dtype='{model_configs['weight_dtype']}' \\
148+
scan_layers='{model_configs['scan_layers']}' \\
149+
max_prefill_predict_length='{model_configs['max_prefill_predict_length']}' \\
150+
max_target_length='{model_configs['max_target_length']}' \\
151+
attention='{model_configs['attention']}' \\
152+
ici_fsdp_parallelism='{model_configs['ici_fsdp_parallelism']}' \\
153+
ici_autoregressive_parallelism='{model_configs['ici_autoregressive_parallelism']}' \\
154+
ici_tensor_parallelism='{model_configs['ici_tensor_parallelism']}' \\
155+
load_parameters_path='{model_configs['checkpoint']}' \\
156+
quantization='"{model_configs.get('quantization') or ""}"' \\
157+
quantize_kvcache='{model_configs['quantize_kvcache']}' \\
158+
kv_quant_dtype='"{model_configs.get('kv_quant_dtype') or ""}"' \\
159+
per_device_batch_size='{model_configs['per_device_batch_size']}' \\
160+
prefill_cache_axis_order='{model_configs['prefill_cache_axis_order']}' \\
161+
ar_cache_axis_order='{model_configs['ar_cache_axis_order']}' \\
162+
compute_axis_order='{model_configs['compute_axis_order']}' \\
163+
reshape_q='{model_configs['reshape_q']}' \\
164+
kv_quant_axis='"{model_configs.get('kv_quant_axis') or ""}"' &""",
183165
"cd ..",
184166
# Give server time to start
185167
f"sleep {model_configs['sleep_time']}",
168+
"cd JetStream",
169+
# Since we change to the Jetstream dir as root, we need to change the ownership of the dir to the user running the script
170+
"sudo chown -R $USER:$USER $HOME/JetStream",
171+
"git lfs pull",
172+
"cd ..",
186173
# Run benchmark, run eval, save benchmark and eval results, and save predictions to /tmp/request-outputs.json
187174
f"""python JetStream/benchmarks/benchmark_serving.py \
188175
--tokenizer {full_tokenizer_path} \
@@ -233,7 +220,6 @@ def get_config(
233220
json_lines=metric_config.JSONLinesConfig("metric_report.jsonl"),
234221
use_runtime_generated_gcs_folder=True,
235222
)
236-
237223
return task.run_queued_resource_test(
238224
task_test_config=job_test_config,
239225
task_gcp_config=job_gcp_config,

‎dags/inference/jetstream_inference_e2e.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,7 @@
1616

1717
import datetime
1818
from airflow import models
19-
from dags.common.vm_resource import TpuVersion
20-
from dags.inference.maxtext_model_config_generator import generate_model_configs
19+
from dags.inference.maxtext_model_config_generator import generate_model_configs, TpuConfig
2120

2221
"""A JetStream inference E2E test (JAX nightly, no schedule) DAG.
2322
@@ -72,7 +71,7 @@
7271
"jetstream_branch": "",
7372
"sleep_time": 360,
7473
"time_out_in_min": 60,
75-
"tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
74+
"tpu_version_cores": [(TpuConfig.V5E, 8), (TpuConfig.TRILLIUM, 8)],
7675
"model_name": LLAMA2_7B,
7776
"tokenizer": "tokenizer.llama2",
7877
"weight_dtype": "bfloat16",
@@ -100,7 +99,7 @@
10099
"jetstream_branch": "",
101100
"sleep_time": 360,
102101
"time_out_in_min": 60,
103-
"tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
102+
"tpu_version_cores": [(TpuConfig.V5E, 8), (TpuConfig.TRILLIUM, 8)],
104103
"model_name": GEMMA_7B,
105104
"tokenizer": "tokenizer.gemma",
106105
"weight_dtype": "bfloat16",

‎dags/inference/maxtext_inference.py‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@
1818
import numpy as np
1919
from airflow import models
2020
from dags import composer_env
21-
from dags.common.vm_resource import TpuVersion
22-
from dags.inference.maxtext_model_config_generator import generate_model_configs
21+
22+
from dags.inference.maxtext_model_config_generator import generate_model_configs, TpuConfig
2323

2424
USER_PREFIX = ""
2525

@@ -117,7 +117,10 @@
117117
"jetstream_branch": jetstream_branch,
118118
"sleep_time": 360,
119119
"time_out_in_min": 120,
120-
"tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
120+
"tpu_version_cores": [
121+
(TpuConfig.temp_V5E, 8),
122+
(TpuConfig.temp_TRILLIUM, 8),
123+
],
121124
"model_name": LLAMA2_7B,
122125
"tokenizer": "tokenizer.llama2",
123126
"weight_dtype": "bfloat16",
@@ -155,7 +158,10 @@
155158
"jetstream_branch": jetstream_branch,
156159
"sleep_time": 360,
157160
"time_out_in_min": 120,
158-
"tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
161+
"tpu_version_cores": [
162+
(TpuConfig.temp_V5E, 8),
163+
(TpuConfig.temp_TRILLIUM, 8),
164+
],
159165
"model_name": LLAMA2_13B,
160166
"tokenizer": "tokenizer.llama2",
161167
"weight_dtype": "bfloat16",
@@ -193,7 +199,10 @@
193199
"jetstream_branch": jetstream_branch,
194200
"sleep_time": 360,
195201
"time_out_in_min": 240,
196-
"tpu_version_cores": [(TpuVersion.V5P, 8), (TpuVersion.TRILLIUM, 8)],
202+
"tpu_version_cores": [
203+
(TpuConfig.V5P, 8),
204+
(TpuConfig.temp_TRILLIUM, 8),
205+
],
197206
"model_name": LLAMA2_70B,
198207
"tokenizer": "tokenizer.llama2",
199208
"weight_dtype": "bfloat16",
@@ -230,7 +239,10 @@
230239
"jetstream_branch": jetstream_branch,
231240
"sleep_time": 360,
232241
"time_out_in_min": 120,
233-
"tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
242+
"tpu_version_cores": [
243+
(TpuConfig.temp_V5E, 8),
244+
(TpuConfig.temp_TRILLIUM, 8),
245+
],
234246
"model_name": GEMMA_7B,
235247
"tokenizer": "tokenizer.gemma",
236248
"weight_dtype": "bfloat16",
@@ -269,7 +281,10 @@
269281
"jetstream_branch": jetstream_branch,
270282
"sleep_time": 240,
271283
"time_out_in_min": 240,
272-
"tpu_version_cores": [(TpuVersion.V5P, 8), (TpuVersion.TRILLIUM, 8)],
284+
"tpu_version_cores": [
285+
(TpuConfig.V5P, 8),
286+
(TpuConfig.temp_TRILLIUM, 8),
287+
],
273288
"model_name": MIXTRAL_8_7B,
274289
"tokenizer": "gs://maxtext-external/mixtral-8x7B-v0.1-Instruct/tokenizer.mistral",
275290
"weight_dtype": "bfloat16",
@@ -576,7 +591,6 @@
576591
},
577592
}
578593
)
579-
580594
# run_configs = [
581595
# f"{LLAMA2_7B}-{BASE_MODE}-{W_BF16_KV_BF16}",
582596
# f"{LLAMA2_7B}-{BASE_MODE}-{W_INT8_KV_INT8}",

‎dags/inference/maxtext_inference_microbenchmark.py‎

Lines changed: 11 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@
1919
import itertools
2020
import numpy
2121
from airflow import models
22-
from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
22+
2323
from dags.inference.configs import maxtext_inference_microbenchmark_gce_config
24+
from dags.inference.maxtext_model_config_generator import TpuConfig
2425
from dags.multipod.configs.common import SetupMode
2526

2627
USER_PREFIX = ""
@@ -110,6 +111,8 @@ def generate_model_configs(
110111
tpu_version,
111112
tpu_cores,
112113
):
114+
config = tpu_version.value
115+
113116
model_configs = {}
114117
model_configs["model_config_name"] = model_config_name
115118

@@ -177,32 +180,18 @@ def generate_model_configs(
177180
test_name = f"{test_name_prefix}-{test_run_tag}"
178181
model_configs["run_name"] = test_run_tag
179182

180-
if tpu_version == TpuVersion.V5E:
181-
# v5e benchmarks
182-
project_name = Project.TPU_PROD_ENV_AUTOMATED.value
183-
zone = Zone.US_EAST1_C.value
184-
network = V5_NETWORKS
185-
subnetwork = V5E_SUBNETWORKS
186-
runtime_version = RuntimeVersion.V2_ALPHA_TPUV5_LITE.value
187-
if tpu_version == TpuVersion.TRILLIUM:
188-
project_name = Project.TPU_PROD_ENV_AUTOMATED.value
189-
zone = Zone.US_EAST5_A.value
190-
network = V6E_GCE_NETWORK
191-
subnetwork = V6E_GCE_SUBNETWORK
192-
runtime_version = RuntimeVersion.V2_ALPHA_TPUV6.value
193-
194183
maxtext_kv_cache_layout_optimization = (
195184
maxtext_inference_microbenchmark_gce_config.config(
196-
tpu_version=tpu_version,
185+
tpu_version=config.tpu_version_name,
197186
tpu_cores=tpu_cores,
198-
tpu_zone=zone,
187+
tpu_zone=config.zone,
199188
time_out_in_min=sweep_model_configs["time_out_in_min"],
200189
test_name=test_name,
201190
test_mode=SetupMode.STABLE,
202-
project_name=project_name,
203-
runtime_version=runtime_version,
204-
network=network,
205-
subnetwork=subnetwork,
191+
project_name=config.project_name,
192+
runtime_version=config.runtime_version,
193+
network=config.network,
194+
subnetwork=config.subnetwork,
206195
is_tpu_reserved=True,
207196
model_configs=model_configs,
208197
maxtext_branch=model_configs["maxtext_branch"],
@@ -247,7 +236,7 @@ def generate_model_configs(
247236
if not MAXTEXT_BRANCH
248237
else f"-b {MAXTEXT_BRANCH}",
249238
"sleep_time": 60,
250-
"tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
239+
"tpu_version_cores": [(TpuConfig.V5E, 8), (TpuConfig.TRILLIUM, 8)],
251240
"model_name": LLAMA2_7B,
252241
"tokenizer": "tokenizer.llama2",
253242
"weight_dtype": "bfloat16",

0 commit comments

Comments
 (0)