diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh index 3546a98027..e2f4966830 100644 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh @@ -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 @@ -31,7 +32,7 @@ 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 RL on the converted checkpoint python3 -m maxtext.trainers.post_train.rl.train_rl \ @@ -39,7 +40,7 @@ python3 -m maxtext.trainers.post_train.rl.train_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"]}' \ @@ -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} diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_sft.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_sft.sh index 6488c838c2..7c991a1160 100644 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_sft.sh +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_sft.sh @@ -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 @@ -31,7 +32,7 @@ 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 \ @@ -39,7 +40,7 @@ python3 -m maxtext.trainers.post_train.sft.train_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 @@ -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}