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

Author: SurbhiJainUSC

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,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 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"]}' \
@@ -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}

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,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
@@ -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}