add jax persistent cache thresholds

google · Aug 20, 2024 · a2e2893 · a2e2893
1 parent 15966fa
commit a2e2893
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/MaxText/configs/base.yml b/MaxText/configs/base.yml
@@ -171,6 +171,10 @@ local_checkpoint_period: 0
 # Jax cache directory
 jax_cache_dir: "~/jax_cache"
 
+# Jax persistent cache thresholds
+jax_persistent_cache_min_entry_size_bytes: 0
+jax_persistent_cache_min_compile_time_secs: 1
+
 # Hardware
 hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu', 'gpu_multiprocess' and 'cpu'
 

diff --git a/MaxText/pyconfig.py b/MaxText/pyconfig.py
@@ -278,6 +278,12 @@ def __init__(self, argv: list[str], **kwargs):
     if raw_keys["jax_cache_dir"]:
       compilation_cache.set_cache_dir(os.path.expanduser(raw_keys["jax_cache_dir"]))
 
+    if raw_keys["jax_persistent_cache_min_entry_size_bytes"]:
+      jax.config.update("jax_persistent_cache_min_entry_size_bytes", os.path.expanduser(raw_keys["jax_persistent_cache_min_entry_size_bytes"]))
+
+    if raw_keys["jax_persistent_cache_min_compile_time_secs"]:
+      jax.config.update("jax_persistent_cache_min_compile_time_secs", os.path.expanduser(raw_keys["jax_persistent_cache_min_compile_time_secs"]))
+
     if raw_keys["model_name"] == "gpt3-175b":
       _HyperParameters.configure_gpt3_task(raw_keys)
     _HyperParameters.user_init(raw_keys)