Skip to content

Commit e9e0cd2

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request #4185 from AI-Hypercomputer:yixuann-dev-qwen330
PiperOrigin-RevId: 944703230
2 parents 7c6386a + 3178ed4 commit e9e0cd2

5 files changed

Lines changed: 290 additions & 0 deletions

File tree

‎tests/end_to_end/tpu/qwen3/30b/test_qwen3.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 Qwen3-30B 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_qwen3_to_mt.sh $RUN_ID
14+
# bash test_qwen3.sh $RUN_ID
15+
16+
17+
set -ex
18+
19+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
20+
MODEL_NAME='qwen3-30b-a3b-base'
21+
22+
23+
BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
24+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
25+
SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items
26+
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
27+
28+
DATASET_PATH=gs://maxtext-dataset
29+
30+
# Step 1: Run inference on the pre-converted checkpoint
31+
python3 -m maxtext.inference.decode \
32+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
33+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
34+
per_device_batch_size=1 run_name=${run_id} \
35+
max_prefill_predict_length=8 max_target_length=16 steps=1 async_checkpointing=false \
36+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
37+
scan_layers=false prompt='I love to' attention=\'dot_product\'
38+
39+
# Step 2: Run pre-training starting from the pre-converted checkpoint
40+
python3 -m maxtext.trainers.pre_train.train \
41+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/train \
42+
dataset_path=${DATASET_PATH} tokenizer_type="huggingface" \
43+
load_parameters_path=${SCANNED_CKPT_PATH} \
44+
per_device_batch_size=0.125 run_name=${run_id} \
45+
max_target_length=64 steps=5 async_checkpointing=false \
46+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
47+
model_name=${MODEL_NAME} scan_layers=true \
48+
remat_policy=full \
49+
ici_tensor_parallelism=4 ici_fsdp_parallelism=2 weight_dtype=bfloat16 dtype=bfloat16 opt_type=sgd optimizer_memory_host_offload=true
50+
51+
# Step 3: Run inference on the checkpoint produced by the pre-training run
52+
python3 -m maxtext.inference.decode \
53+
model_name=${MODEL_NAME} tokenizer_type="huggingface" \
54+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/train/${run_id}/checkpoints/4/items \
55+
per_device_batch_size=1 run_name=${run_id} \
56+
max_prefill_predict_length=8 max_target_length=16 steps=4 async_checkpointing=false \
57+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
58+
scan_layers=true prompt='I love to' attention=\'dot_product\'

‎tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
#!/bin/bash
2+
3+
# Validates the Qwen3 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_qwen3_to_mt.sh $RUN_ID
14+
# bash test_qwen3_rl.sh $RUN_ID
15+
16+
set -ex
17+
18+
export VLLM_WORKER_MULTIPROC_METHOD=spawn
19+
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=upb
20+
export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"
21+
export VLLM_ENABLE_V1_MULTIPROCESSING=0
22+
export JAX_RANDOM_WEIGHTS=1
23+
export SKIP_JAX_PRECOMPILE=1
24+
export NEW_MODEL_DESIGN=1
25+
export TPU_MIN_LOG_LEVEL=0
26+
export TF_CPP_MIN_LOG_LEVEL=0
27+
export TPU_STDERR_LOG_LEVEL=0
28+
29+
# Force Python to use "spawn" instead of "fork" for multiprocessing to prevent gRPC socket corruption
30+
cat << 'EOF' > usercustomize.py
31+
import multiprocessing
32+
try:
33+
multiprocessing.set_start_method("spawn", force=True)
34+
except Exception:
35+
pass
36+
EOF
37+
export PYTHONPATH=${PYTHONPATH:-.}:$(pwd)
38+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
39+
MODEL_NAME='qwen3-30b-a3b-base'
40+
41+
BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
42+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
43+
SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items
44+
45+
# Step 1: Run inference on the pre-converted checkpoint
46+
python3 -m maxtext.inference.vllm_decode \
47+
model_name=${MODEL_NAME} \
48+
load_parameters_path=${SCANNED_CKPT_PATH} \
49+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
50+
hbm_utilization_vllm=0.85 \
51+
prompt='Suggest some famous landmarks in London.' \
52+
max_target_length=256 max_num_batched_tokens=256 \
53+
ici_tensor_parallelism=8 \
54+
allow_split_physical_axes=True prefuse_moe_weights=True \
55+
use_chat_template=True scan_layers=True enable_single_controller=True
56+
57+
# Step 2: Run RL starting from the pre-converted checkpoint
58+
python3 -m maxtext.trainers.post_train.rl.train_rl \
59+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \
60+
load_parameters_path=${SCANNED_CKPT_PATH} \
61+
run_name=${run_id} rl.loss_algo='grpo' scan_layers=True \
62+
num_batches=5 batch_size=4 train_micro_batch_size=1 num_test_batches=5 \
63+
model_name=${MODEL_NAME} enable_single_controller=True \
64+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
65+
rollout_tensor_parallelism=8 \
66+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
67+
vllm_additional_config='{"maxtext_config": {"model_name": "'${MODEL_NAME}'", "allow_split_physical_axes": true, "scan_layers": false, "prefuse_moe_weights": True}}' \
68+
remat_policy=full hbm_utilization_vllm=0.55 use_pathways=True \
69+
chips_per_vm=8 ici_tensor_parallelism=4\
70+
max_target_length=512 weight_dtype=bfloat16 dtype=bfloat16 opt_type=sgd \
71+
enable_tunix_perf_metrics=True rl.num_generations=16 \
72+
debug.rl=False rl.reshard_chunk_size=1
73+
74+
# Step 3: Run inference on the checkpoint produced by the RL run
75+
python3 -m maxtext.inference.vllm_decode \
76+
model_name=${MODEL_NAME} \
77+
load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/4/items \
78+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
79+
hbm_utilization_vllm=0.85 \
80+
prompt='Suggest some famous landmarks in London.' \
81+
max_target_length=256 max_num_batched_tokens=256 \
82+
ici_tensor_parallelism=8 \
83+
allow_split_physical_axes=True prefuse_moe_weights=True \
84+
use_chat_template=True scan_layers=True enable_single_controller=True

‎tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#!/bin/bash
2+
3+
# Validates the Qwen3 SFT pipeline using a pre-converted MaxText checkpoint.
4+
5+
6+
# The flow of this script is as follows:
7+
# 1. Run inference on the pre-converted checkpoint.
8+
# 2. Run SFT starting from the pre-converted checkpoint.
9+
# 3. Run inference on the checkpoint produced by the SFT run.
10+
11+
# Usage:
12+
# export HF_TOKEN=<your Hugging Face access token>
13+
# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
14+
# bash test_qwen3_to_mt.sh $RUN_ID
15+
# bash test_qwen3_sft.sh $RUN_ID
16+
17+
18+
set -ex
19+
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
20+
export NEW_MODEL_DESIGN=1
21+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
22+
MODEL_NAME='qwen3-30b-a3b-base'
23+
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 pre-converted checkpoint
29+
python3 -m maxtext.inference.vllm_decode \
30+
model_name=${MODEL_NAME} \
31+
load_parameters_path=${SCANNED_CKPT_PATH} \
32+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
33+
hbm_utilization_vllm=0.85 \
34+
use_chat_template=True scan_layers=true enable_single_controller=True \
35+
ici_tensor_parallelism=4 prompt="Suggest some famous landmarks in London."
36+
37+
# Step 2: Run SFT starting from the pre-converted checkpoint
38+
python3 -m maxtext.trainers.post_train.sft.train_sft \
39+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/sft \
40+
load_parameters_path=${SCANNED_CKPT_PATH} \
41+
per_device_batch_size=0.125 run_name=${run_id} \
42+
steps=5 scan_layers=true \
43+
model_name=${MODEL_NAME} \
44+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
45+
remat_policy=full \
46+
ici_tensor_parallelism=1 ici_fsdp_parallelism=1 ici_expert_parallelism=8 \
47+
enable_single_controller=True max_target_length=16 weight_dtype=bfloat16 dtype=bfloat16 opt_type=sgd
48+
49+
# Step 3: Run inference on the checkpoint produced by the SFT 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+
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
54+
hbm_utilization_vllm=0.85 \
55+
use_chat_template=True scan_layers=true enable_single_controller=True \
56+
ici_tensor_parallelism=4 prompt="Suggest some famous landmarks in London."

‎tests/end_to_end/tpu/qwen3/30b/test_qwen3_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_qwen3_to_hf.sh $RUN_ID $CHECKPOINT_PATH $SCAN_LAYERS
8+
9+
10+
set -ex
11+
12+
run_id=$1
13+
CKPT_PATH=$2
14+
SCAN_LAYERS=${3:-false}
15+
16+
MODEL_NAME='qwen3-30b-a3b-base'
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+
# Convert the checkpoint to Hugging Face format
26+
python3 -m maxtext.checkpoint_conversion.to_huggingface \
27+
model_name=${MODEL_NAME} \
28+
tokenizer_type="huggingface" \
29+
load_parameters_path=${CKPT_PATH} \
30+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/${scan_status}/${run_id} \
31+
scan_layers=$SCAN_LAYERS

‎tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
#!/bin/bash
2+
3+
# Converts Qwen3-30B 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_qwen3_to_mt.sh $RUN_ID
14+
15+
set -ex
16+
17+
run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
18+
MODEL_NAME='qwen3-30b-a3b-base'
19+
HF_GOLDEN_MODEL='Qwen/Qwen3-30B-A3B-Base'
20+
21+
22+
BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext
23+
24+
# Step 1: Install torch
25+
python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
26+
27+
# Step 2: Convert the checkpoint from Hugging Face
28+
python3 -m maxtext.checkpoint_conversion.to_maxtext \
29+
model_name=${MODEL_NAME} \
30+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id} \
31+
scan_layers=false \
32+
hardware=cpu skip_jax_distributed_system=True \
33+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
34+
--lazy_load_tensors=False \
35+
--eager_load_method='safetensors'
36+
37+
UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id}/0/items
38+
echo "Unscanned checkpoint path: ${UNSCANNED_CKPT_PATH}"
39+
40+
# Convert to scanned format
41+
python3 -m maxtext.checkpoint_conversion.to_maxtext \
42+
model_name=${MODEL_NAME} \
43+
base_output_directory=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id} \
44+
scan_layers=true \
45+
hardware=cpu skip_jax_distributed_system=True \
46+
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
47+
--lazy_load_tensors=False \
48+
--eager_load_method='safetensors'
49+
50+
SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id}/0/items
51+
echo "Scanned checkpoint path: ${SCANNED_CKPT_PATH}"
52+
53+
# Step 3: Run forward pass logits check
54+
python3 -m tests.utils.forward_pass_logit_checker \
55+
load_parameters_path=${UNSCANNED_CKPT_PATH} \
56+
model_name=${MODEL_NAME} \
57+
scan_layers=false \
58+
--hf_model_path=${HF_GOLDEN_MODEL} \
59+
--max_kl_div=0.03 \
60+
--run_hf_model=true \
61+
hardware=cpu skip_jax_distributed_system=True

0 commit comments

Comments
 (0)