Skip to content

Commit 3e5afc9

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request AI-Hypercomputer#4278 from CIeNET-International:emma/e2e-training
PiperOrigin-RevId: 949618256
2 parents a375072 + edc8b72 commit 3e5afc9

6 files changed

Lines changed: 287 additions & 172 deletions

File tree

‎tests/end_to_end/tpu/gemma4/26b/convert_gemma4.sh‎

Lines changed: 0 additions & 82 deletions
This file was deleted.

‎tests/end_to_end/tpu/gemma4/26b/test_gemma4.sh‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
#!/bin/bash
2+
3+
# Validates the Gemma4-26B 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_gemma4_to_mt.sh $RUN_ID
14+
# bash test_gemma4.sh $RUN_ID
15+
16+
17+
set -ex
18+
19+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
20+
MODEL_NAME='gemma4-26b'
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=8192 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+
64+
# Step 3: Run inference on the checkpoint generated from the previous run
65+
if [ ${USE_MULTIMODAL} == true ]; then
66+
python3 -m maxtext.inference.decode \
67+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
68+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/train/${run_id}/checkpoints/4/items \
69+
per_device_batch_size=1 run_name=${run_id} \
70+
max_prefill_predict_length=272 max_target_length=300 steps=1 async_checkpointing=false \
71+
scan_layers=false use_multimodal=true \
72+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
73+
prompt=\'Describe\ image\ \<start_of_image\>\' image_path=\'tests/assets/test_image.jpg\' attention=\'dot_product\'
74+
else
75+
python3 -m maxtext.inference.decode \
76+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
77+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/train/${run_id}/checkpoints/4/items \
78+
per_device_batch_size=1 run_name=${run_id} \
79+
max_prefill_predict_length=8 max_target_length=16 steps=1 async_checkpointing=false \
80+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
81+
scan_layers=false prompt='I love to' attention=\'dot_product\'
82+
fi

‎tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.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 Gemma4-26B 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_gemma4_to_mt.sh $RUN_ID
14+
# bash test_gemma4_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='gemma4-26b'
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+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
33+
hbm_utilization_vllm=0.85 \
34+
prompt="Suggest some famous landmarks in London." \
35+
use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
36+
prefuse_moe_weights=True ici_tensor_parallelism=8
37+
38+
# Step 2: Run RL on the converted checkpoint
39+
python3 -m maxtext.trainers.post_train.rl.train_rl \
40+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \
41+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
42+
run_name=${run_id} rl.loss_algo='grpo' scan_layers=false \
43+
num_batches=5 batch_size=16 num_test_batches=5 \
44+
model_name=${MODEL_NAME} enable_single_controller=${use_pathways} \
45+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
46+
rollout_tensor_parallelism=4 \
47+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
48+
vllm_additional_config='{"maxtext_config": {"model_name": "gemma4-26b", "log_config": "false", "prefuse_moe_weights": "true"}}'
49+
50+
# Step 3: Run inference on the checkpoint generated from the previous run
51+
python3 -m maxtext.inference.vllm_decode \
52+
model_name=${MODEL_NAME} \
53+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/5/model_params \
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=false enable_single_controller=${use_pathways} \
58+
prefuse_moe_weights=True ici_tensor_parallelism=8

‎tests/end_to_end/tpu/gemma4/26b/test_gemma4_sft.sh‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#!/bin/bash
2+
3+
# Validates the Gemma4-26B 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_gemma4_to_mt.sh $RUN_ID
14+
# bash test_gemma4_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='gemma4-26b'
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+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
33+
hbm_utilization_vllm=0.85 \
34+
prompt="Suggest some famous landmarks in London." \
35+
use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
36+
prefuse_moe_weights=True ici_tensor_parallelism=8
37+
38+
# Step 2: Run SFT on the converted checkpoint
39+
python3 -m maxtext.trainers.post_train.sft.train_sft \
40+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/sft \
41+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
42+
per_device_batch_size=1 run_name=${run_id} \
43+
steps=5 scan_layers=false \
44+
model_name=${MODEL_NAME} enable_single_controller=${use_pathways} \
45+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False
46+
47+
# Step 3: Run inference on the checkpoint generated from the previous run
48+
python3 -m maxtext.inference.vllm_decode \
49+
model_name=${MODEL_NAME} \
50+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/sft/${run_id}/checkpoints/5/model_params \
51+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
52+
hbm_utilization_vllm=0.85 \
53+
prompt="Suggest some famous landmarks in London." \
54+
use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
55+
prefuse_moe_weights=True ici_tensor_parallelism=8

0 commit comments

Comments
 (0)