Skip to content

Commit 6bf8e6f

Browse files
chiajunglienchiajunglien
committed
feat(testing): add end-to-end training and inference tests for llama3.1-70b
Introduces E2E test configurations and scripts for llama3.1-70b model, covering both pre-training and post-training validation pipelines.
1 parent 2ee9fe1 commit 6bf8e6f

5 files changed

Lines changed: 306 additions & 0 deletions

File tree

‎tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b.sh‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
#!/bin/bash
2+
3+
# Validates the Llama3.1-70b pre-training pipeline using a pre-converted MaxText checkpoint.
4+
5+
# The flow of this script is as follows:
6+
# 1. Run inference on the pre-converted checkpoint.
7+
# 2. Run pre-training starting from the pre-converted checkpoint.
8+
# 3. Run inference on the checkpoint produced by the pre-training run.
9+
10+
# Usage:
11+
# export HF_TOKEN=<your Hugging Face access token>
12+
# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
13+
# bash test_llama3.1_70b_to_mt.sh $RUN_ID
14+
# bash test_llama3.1_70b.sh $RUN_ID
15+
16+
17+
set -ex
18+
19+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
20+
MODEL_NAME='llama3.1-70b'
21+
22+
# To convert the multimodal model, make sure the use_multimodal is set to be true
23+
USE_MULTIMODAL=false
24+
25+
# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
26+
BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
27+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
28+
29+
# Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
30+
DATASET_PATH=gs://maxtext-dataset
31+
32+
# Step 1: Run inference on the original checkpoint converted from Hugging Face
33+
if [ ${USE_MULTIMODAL} == true ]; then
34+
python3 -m maxtext.inference.decode \
35+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
36+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
37+
per_device_batch_size=1 run_name=${run_id} \
38+
max_prefill_predict_length=272 max_target_length=300 steps=1 async_checkpointing=false \
39+
scan_layers=false use_multimodal=true \
40+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
41+
prompt=\'Describe\ image\ \<start_of_image\>\' image_path=\'tests/assets/test_image.jpg\' attention=\'dot_product\'
42+
else
43+
python3 -m maxtext.inference.decode \
44+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
45+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
46+
per_device_batch_size=1 run_name=${run_id} \
47+
max_prefill_predict_length=8 max_target_length=16 steps=1 async_checkpointing=false \
48+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
49+
scan_layers=false prompt='I love to' attention=\'dot_product\'
50+
fi
51+
52+
# Step 2: Run Pre-training on the converted checkpoint
53+
# We can also run training by using the scanned converted checkpoint
54+
# Note that scanned checkpoint helps with efficient training
55+
python3 -m maxtext.trainers.pre_train.train \
56+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/train \
57+
dataset_path=${DATASET_PATH} tokenizer_type="huggingface" \
58+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
59+
per_device_batch_size=1 run_name=${run_id} \
60+
max_target_length=2048 steps=5 async_checkpointing=false \
61+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
62+
model_name=${MODEL_NAME} scan_layers=false use_multimodal=${USE_MULTIMODAL} \
63+
ici_tensor_parallelism=4
64+
65+
# Step 3: Run inference on the checkpoint generated from the previous run
66+
if [ ${USE_MULTIMODAL} == true ]; then
67+
python3 -m maxtext.inference.decode \
68+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
69+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/train/${run_id}/checkpoints/4/items \
70+
per_device_batch_size=1 run_name=${run_id} \
71+
max_prefill_predict_length=272 max_target_length=300 steps=1 async_checkpointing=false \
72+
scan_layers=false use_multimodal=true \
73+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
74+
prompt=\'Describe\ image\ \<start_of_image\>\' image_path=\'tests/assets/test_image.jpg\' attention=\'dot_product\'
75+
else
76+
python3 -m maxtext.inference.decode \
77+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
78+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/train/${run_id}/checkpoints/4/items \
79+
per_device_batch_size=1 run_name=${run_id} \
80+
max_prefill_predict_length=8 max_target_length=16 steps=1 async_checkpointing=false \
81+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
82+
scan_layers=false prompt='I love to' attention=\'dot_product\'
83+
fi

‎tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
#!/bin/bash
2+
3+
# Validates the Llama3.1-70b RL pipeline using a pre-converted MaxText checkpoint.
4+
5+
# The flow of this script is as follows:
6+
# 1. Run inference on the pre-converted checkpoint.
7+
# 2. Run RL starting from the pre-converted checkpoint.
8+
# 3. Run inference on the checkpoint produced by the RL run.
9+
10+
# Usage:
11+
# export HF_TOKEN=<your Hugging Face access token>
12+
# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
13+
# bash test_llama3.1_70b_to_mt.sh $RUN_ID
14+
# bash test_llama3.1_70b_rl.sh $RUN_ID
15+
16+
17+
set -ex
18+
19+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
20+
use_pathways=${2:-false}
21+
MODEL_NAME='llama3.1-70b'
22+
23+
# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
24+
BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
25+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
26+
SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items
27+
28+
# Step 1: Run inference on the original checkpoint converted from Hugging Face
29+
python3 -m maxtext.inference.vllm_decode \
30+
model_name=${MODEL_NAME} \
31+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
32+
tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
33+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
34+
hbm_utilization_vllm=0.85 \
35+
prompt='Suggest some famous landmarks in London.' \
36+
use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
37+
ici_tensor_parallelism=8
38+
39+
# Step 2: Run RL on the converted checkpoint
40+
python3 -m maxtext.trainers.post_train.rl.train_rl \
41+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \
42+
load_parameters_path=${SCANNED_CKPT_PATH} \
43+
run_name=${run_id} rl.loss_algo='grpo' scan_layers=true \
44+
num_batches=5 batch_size=1 num_test_batches=5 \
45+
model_name=${MODEL_NAME} tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
46+
enable_single_controller=${use_pathways} \
47+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
48+
rollout_tensor_parallelism=4 \
49+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
50+
vllm_additional_config='{"maxtext_config": {"model_name": "llama3.1-70b", "log_config": "false"}}'
51+
52+
53+
# Step 3: Run inference on the checkpoint generated from the previous run
54+
python3 -m maxtext.inference.vllm_decode \
55+
model_name=${MODEL_NAME} \
56+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/5/model_params \
57+
tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
58+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
59+
hbm_utilization_vllm=0.85 \
60+
prompt='Suggest some famous landmarks in London.' \
61+
use_chat_template=True scan_layers=true enable_single_controller=${use_pathways} \
62+
ici_tensor_parallelism=8

‎tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_sft.sh‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
#!/bin/bash
2+
3+
# Validates the Llama3.1-70b SFT pipeline using a pre-converted MaxText checkpoint.
4+
5+
# The flow of this script is as follows:
6+
# 1. Run inference on the pre-converted checkpoint.
7+
# 2. Run SFT starting from the pre-converted checkpoint.
8+
# 3. Run inference on the checkpoint produced by the SFT run.
9+
10+
# Usage:
11+
# export HF_TOKEN=<your Hugging Face access token>
12+
# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
13+
# bash test_llama3.1_70b_to_mt.sh $RUN_ID
14+
# bash test_llama3.1_70b_sft.sh $RUN_ID
15+
16+
17+
set -ex
18+
19+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
20+
use_pathways=${2:-false}
21+
MODEL_NAME='llama3.1-70b'
22+
23+
# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
24+
BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
25+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
26+
SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items
27+
28+
# Step 1: Run inference on the original checkpoint converted from Hugging Face
29+
python3 -m maxtext.inference.vllm_decode \
30+
model_name=${MODEL_NAME} \
31+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
32+
tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
33+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
34+
hbm_utilization_vllm=0.85 \
35+
prompt="Suggest some famous landmarks in London." \
36+
use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
37+
ici_tensor_parallelism=8
38+
39+
# Step 2: Run SFT on the converted checkpoint
40+
python3 -m maxtext.trainers.post_train.sft.train_sft \
41+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/sft \
42+
load_parameters_path=${SCANNED_CKPT_PATH} \
43+
per_device_batch_size=1 run_name=${run_id} \
44+
steps=5 scan_layers=true \
45+
model_name=${MODEL_NAME} tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
46+
enable_single_controller=${use_pathways} \
47+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False
48+
49+
# Step 3: Run inference on the checkpoint generated from the previous run
50+
python3 -m maxtext.inference.vllm_decode \
51+
model_name=${MODEL_NAME} \
52+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/sft/${run_id}/checkpoints/5/model_params \
53+
tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
54+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
55+
hbm_utilization_vllm=0.85 \
56+
prompt="Suggest some famous landmarks in London." \
57+
use_chat_template=True scan_layers=true enable_single_controller=${use_pathways} \
58+
ici_tensor_parallelism=8

‎tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_hf.sh‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
#!/bin/bash
2+
3+
# Converts a MaxText checkpoint to a Hugging Face model checkpoint.
4+
5+
# Usage:
6+
# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
7+
# bash test_llama3.1_70b_to_hf.sh $RUN_ID $CHECKPOINT_PATH $USE_MULTIMODAL $SCAN_LAYERS
8+
9+
set -ex
10+
11+
run_id=$1
12+
CKPT_PATH=$2
13+
USE_MULTIMODAL=${3:-false}
14+
SCAN_LAYERS=${4:-false}
15+
16+
MODEL_NAME='llama3.1-70b'
17+
BASE_OUTPUT_DIRECTORY="gs://runner-maxtext-logs/${MODEL_NAME}"
18+
19+
if [ "${SCAN_LAYERS,,}" = "true" ]; then
20+
scan_status="scanned"
21+
else
22+
scan_status="unscanned"
23+
fi
24+
25+
python3 -m maxtext.checkpoint_conversion.to_huggingface \
26+
model_name=${MODEL_NAME} \
27+
tokenizer_type="huggingface" \
28+
load_parameters_path=${CKPT_PATH} \
29+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/${scan_status}/${run_id} \
30+
use_multimodal=${USE_MULTIMODAL} \
31+
scan_layers=$SCAN_LAYERS

‎tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_mt.sh‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
#!/bin/bash
2+
3+
# Converts Llama3.1-70b HuggingFace checkpoint to MaxText format and validates logit correctness.
4+
5+
# The flow of this script is as follows:
6+
# 1. Install PyTorch (CPU) required for checkpoint conversion.
7+
# 2. Convert the HuggingFace checkpoint to MaxText format in both unscanned and scanned formats.
8+
# 3. Run a forward pass logits check to verify the converted checkpoint matches the original HF model.
9+
10+
# Usage:
11+
# export HF_TOKEN=<your Hugging Face access token>
12+
# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
13+
# bash test_llama3.1_70b_to_mt.sh $RUN_ID - to convert the checkpoint and run logit check for non-multimodal version
14+
# bash test_llama3.1_70b_to_mt.sh $RUN_ID true - to convert the checkpoint and run logit check for multimodal version
15+
16+
17+
set -ex
18+
19+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
20+
MODEL_NAME='llama3.1-70b'
21+
HF_GOLDEN_MODEL='meta-llama/Llama-3.1-70B'
22+
23+
# To convert the multimodal model, make sure the use_multimodal is set to be true
24+
USE_MULTIMODAL=${2:-false}
25+
26+
# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you want to store scanned and unscanned checkpoints
27+
BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext
28+
29+
# Step 1: Install torch
30+
python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
31+
32+
# Step 2: Convert the checkpoint from Hugging Face to make it compatible with MaxText
33+
34+
# Step 2.a: Convert to unscanned checkpoint (for inference)
35+
python3 -m maxtext.checkpoint_conversion.to_maxtext \
36+
model_name=${MODEL_NAME} \
37+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id} \
38+
use_multimodal=${USE_MULTIMODAL} \
39+
scan_layers=false --lazy_load_tensors=True \
40+
hardware=cpu skip_jax_distributed_system=True \
41+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False
42+
43+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id}/0/items
44+
echo "Unscanned checkpoint path: ${UNSCANNED_CKPT_PATH}"
45+
46+
# Step 2.b: Convert to scanned checkpoint (for training)
47+
python3 -m maxtext.checkpoint_conversion.to_maxtext \
48+
model_name=${MODEL_NAME} \
49+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id} \
50+
use_multimodal=${USE_MULTIMODAL} \
51+
scan_layers=true --lazy_load_tensors=True \
52+
hardware=cpu skip_jax_distributed_system=True \
53+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False
54+
55+
SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id}/0/items
56+
echo "Scanned checkpoint path: ${SCANNED_CKPT_PATH}"
57+
58+
# Step 3: Test whether the forward pass logits match the original HF model
59+
# to get higher precision (eg. float32) run on CPU with `JAX_PLATFORMS=cpu`
60+
# ToDo: improve forward_pass_logit_checker to test multi-modal prompt
61+
if [ "${USE_MULTIMODAL}" = "false" ]; then
62+
python3 -m tests.utils.forward_pass_logit_checker \
63+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
64+
model_name=${MODEL_NAME} \
65+
use_multimodal=${USE_MULTIMODAL} \
66+
scan_layers=false \
67+
weight_dtype=bfloat16 \
68+
--hf_model_path=${HF_GOLDEN_MODEL} \
69+
--max_kl_div=0.03 \
70+
--run_hf_model=true \
71+
hardware=cpu skip_jax_distributed_system=True
72+
fi

0 commit comments

Comments
 (0)