Latest changes

AI-Hypercomputer · Dec 5, 2024 · 4f321d7 · 4f321d7
1 parent a75d084
commit 4f321d7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/MaxText/configs/base.yml b/MaxText/configs/base.yml
@@ -467,7 +467,7 @@ enable_jax_profiler: False
 jax_profiler_port: 9999
 
 # Checkpoint Structured logging
-enable_checkpoint_cloud_logger: False
+enable_checkpoint_cloud_logger: True
 
 # Single-controller
 enable_single_controller: True

diff --git a/benchmarks/maxtext_xpk_runner.py b/benchmarks/maxtext_xpk_runner.py
@@ -310,6 +310,7 @@ def build_user_command(
       f' {libtpu_flags}'
       # f' echo {model.tuning_params["sa_block_q"]}-q-dq-{model.tuning_params["sa_block_q_dq"]}-q-dkv-{model.tuning_params["sa_block_q_dkv"]} &&'
       # f' echo {model.tuning_params["ici_fsdp_parallelism"]} {model.tuning_params["ici_tensor_parallelism"]} &&'
+      f' export ENABLE_PATHWAYS_PERSISTENCE=1 &&'
       f' export JAX_PLATFORMS={jax_platforms} &&'
       # f' export JAX_DEBUG_NANS=True &&'
       # f' export TPU_MEGACORE=megachip_tccontrol &&'
@@ -326,7 +327,7 @@ def build_user_command(
       ' vertex_tensorboard_project="" vertex_tensorboard_region=""'
       f' run_name="{pathways_prefix}{run_name}"'
       f' enable_checkpointing=True'
-      f' checkpoint_period=40'
+      f' checkpoint_period=80'
       f' async_checkpointing=True'
   )