Update MoE test config (#493)

GoogleCloudPlatform · Nov 28, 2024 · aea739e · aea739e
1 parent 0ccefe3
commit aea739e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py b/dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py
@@ -57,7 +57,7 @@
           {
               "script_name": "tpu/mixtral/8x7b/2_test_mixtral",
               "cluster": XpkClusters.TPU_V4_128_CLUSTER,
-              "time_out_in_min": 60,
+              "time_out_in_min": 90,
           },
       ],
       "mixtral-8x22b": [
@@ -68,7 +68,7 @@
           },
           {
               "script_name": "tpu/mixtral/8x22b/2_test_mixtral",
-              "cluster": XpkClusters.TPU_V5E_256_CLUSTER,
+              "cluster": XpkClusters.TPU_V4_128_CLUSTER,
               "time_out_in_min": 60,
           },
       ],