Skip to content

Commit cfd2cd4

Browse files
authored
Add Gemma 3 LoRA test to multipod post-training DAG (GoogleCloudPlatform#1307)
This change adds the `lora` test configuration for `gemma3-4b` in the `maxtext_e2e_tpu_post_training` DAG. It runs the end-to-end test script `tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh` to validate the Gemma 3 LoRA post-training pipeline.
1 parent e1b5af4 commit cfd2cd4

1 file changed

Lines changed: 4 additions & 0 deletions

File tree

dags/multipod/maxtext_e2e_tpu_post_training.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,10 @@ def get_workload_name(model, mode, length=6):
6767
"command": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh",
6868
"maxtext_ckpt_path": "gs://runner-maxtext-logs/gemma3-4b/multimodal_sft/{run_name}/checkpoints/5/model_params",
6969
},
70+
"lora": {
71+
"command": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh",
72+
"maxtext_ckpt_path": "gs://runner-maxtext-logs/gemma3-4b/lora/{run_name}/checkpoints/5/model_params",
73+
},
7074
"rl": {
7175
"command": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh",
7276
"maxtext_ckpt_path": "gs://runner-maxtext-logs/gemma3-4b/rl/{run_name}/checkpoints/actor/5/model_params",

0 commit comments

Comments
 (0)