Skip to content

Update Gemma3-4b E2E post-training scripts for pathways #4168

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jun 15, 2026
+8 −7
Merged

Update Gemma3-4b E2E post-training scripts for pathways #4168

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
set -ex

run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
use_pathways=${2:-false}
MODEL_NAME='gemma3-4b'

# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
Expand All @@ -31,15 +32,15 @@ python3 -m maxtext.inference.vllm_decode \
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
hbm_utilization_vllm=0.5 \
prompt='Suggest some famous landmarks in London.' \
use_chat_template=True scan_layers=false
use_chat_template=True scan_layers=false enable_single_controller=${use_pathways}
Comment thread
SurbhiJainUSC marked this conversation as resolved.

# Step 2: Run RL on the converted checkpoint
python3 -m maxtext.trainers.post_train.rl.train_rl \
base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \
load_parameters_path=${SCANNED_CKPT_PATH} \
run_name=${run_id} rl.loss_algo='grpo' scan_layers=true \
num_batches=5 batch_size=1 num_test_batches=5 \
model_name=${MODEL_NAME} enable_single_controller=True \
model_name=${MODEL_NAME} enable_single_controller=${use_pathways} \
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
rollout_tensor_parallelism=1 \
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
Expand All @@ -53,5 +54,4 @@ python3 -m maxtext.inference.vllm_decode \
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
hbm_utilization_vllm=0.5 \
prompt='Suggest some famous landmarks in London.' \
use_chat_template=True scan_layers=true

use_chat_template=True scan_layers=true enable_single_controller=${use_pathways}
7 changes: 4 additions & 3 deletions tests/end_to_end/tpu/gemma3/4b/test_gemma3_sft.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
set -ex

run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
use_pathways=${2:-false}
MODEL_NAME='gemma3-4b'

# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
Expand All @@ -31,15 +32,15 @@ python3 -m maxtext.inference.vllm_decode \
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
hbm_utilization_vllm=0.5 \
prompt="Suggest some famous landmarks in London." \
use_chat_template=True scan_layers=false
use_chat_template=True scan_layers=false enable_single_controller=${use_pathways}

# Step 2: Run SFT on the converted checkpoint
python3 -m maxtext.trainers.post_train.sft.train_sft \
base_output_directory=${BASE_OUTPUT_DIRECTORY}/sft \
load_parameters_path=${SCANNED_CKPT_PATH} \
per_device_batch_size=1 run_name=${run_id} \
steps=5 scan_layers=true \
model_name=${MODEL_NAME} enable_single_controller=True \
model_name=${MODEL_NAME} enable_single_controller=${use_pathways} \
checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False

# Step 3: Run inference on the checkpoint generated from the previous run
Expand All @@ -49,6 +50,6 @@ python3 -m maxtext.inference.vllm_decode \
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
hbm_utilization_vllm=0.5 \
prompt="Suggest some famous landmarks in London." \
use_chat_template=True scan_layers=true
use_chat_template=True scan_layers=true enable_single_controller=${use_pathways}


Loading