diff --git a/dataflow/gemma-flex-template/requirements-test.txt b/dataflow/gemma-flex-template/requirements-test.txt
index 95aeed12fb..6138ec5ac4 100644
--- a/dataflow/gemma-flex-template/requirements-test.txt
+++ b/dataflow/gemma-flex-template/requirements-test.txt
@@ -1,4 +1,4 @@
-google-cloud-aiplatform==1.62.0
+google-cloud-aiplatform==1.158.0
 google-cloud-dataflow-client==0.8.14
 google-cloud-pubsub==2.28.0
 google-cloud-storage==2.18.2
diff --git a/dataflow/gemma/requirements-test.txt b/dataflow/gemma/requirements-test.txt
index 238d774fdd..45a4425267 100644
--- a/dataflow/gemma/requirements-test.txt
+++ b/dataflow/gemma/requirements-test.txt
@@ -1,4 +1,4 @@
-google-cloud-aiplatform==1.49.0
+google-cloud-aiplatform==1.158.0
 google-cloud-dataflow-client==0.8.10
 google-cloud-storage==2.16.0
 pytest==9.0.3; python_version >= "3.10"
diff --git a/dataflow/run-inference/requirements-test.txt b/dataflow/run-inference/requirements-test.txt
index 03d1c64ee5..a88ea43490 100644
--- a/dataflow/run-inference/requirements-test.txt
+++ b/dataflow/run-inference/requirements-test.txt
@@ -1,4 +1,4 @@
-google-cloud-aiplatform==1.57.0
+google-cloud-aiplatform==1.158.0
 google-cloud-dataflow-client==0.8.14
 google-cloud-storage==2.10.0
 pytest==9.0.3; python_version >= "3.10"
diff --git a/gemma2/requirements-test.txt b/gemma2/requirements-test.txt
index c9e154ba44..b3b2b3ca4c 100644
--- a/gemma2/requirements-test.txt
+++ b/gemma2/requirements-test.txt
@@ -1 +1 @@
-pytest==9.0.3; python_version >= "3.10"
+pytest==9.0.3
diff --git a/gemma2/requirements.txt b/gemma2/requirements.txt
index c46c2a82fd..67a5814ea6 100644
--- a/gemma2/requirements.txt
+++ b/gemma2/requirements.txt
@@ -1,2 +1,2 @@
-google-cloud-aiplatform[full]==1.157.0
-protobuf==5.29.5
+google-cloud-aiplatform==1.158.0
+protobuf==6.33.6