Compilation cache doc #21819

keshavb96 · 2024-06-12T01:05:11Z

More details on persistent compilation caching. Still a WIP.

ayaka14732 · 2024-06-12T14:46:57Z

I will review this PR because I am recently working on this topic.

nouiz · 2024-06-12T16:42:41Z

docs/persistent_compilation_cache.md

+```python
+import jax
+
+jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)


Is it useful to tell that the compilation time isn't constant.
So to have a full cache, if the value isn't set to 0, you may need to run multiple time?

ayaka14732 · 2024-06-12T22:10:21Z

Besides, I am recently working on adding LRU cache eviction support for the JAX persistent compilation cache, and here is the design doc: https://docs.google.com/document/d/111YibwGXOFb_hMm-lua1u63QooAzIBEH-xfRPGmibis/edit?usp=sharing. Linked here because I think it might be useful.

ayaka14732 · 2024-06-15T07:50:29Z

@keshavb96 Can you confirm that this is ready to merge? I will review it if so.

keshavb96 · 2024-06-18T01:50:02Z

@nouiz @jaro-sevcik It looks ready to me, what do you think?

docs/persistent_compilation_cache.md

ayaka14732 · 2024-06-18T02:15:08Z

docs/persistent_compilation_cache.md

+
+### Multiple process and multiple devices (on either single or multiple nodes)
+
+In this runtime the first time a program is run (the persistent cache is cold / empty) all processes will compile, but only the process with rank 0 in the global communication group will write to the persistent cache. In subsequent runs, all processes will attempt to read from the persistent cache, so it is important for the persistent cache to be in a shared file system (eg: NFS) or remote storage (eg: GFS). If the persistent cache is local to rank 0, then all processes except rank 0 will once again compile in subsequent runs as a result of a compilation cache miss. 


There is also "Multi-Host with separate disks". For example, TPU v4-32 has 4 nodes, and there is no shared storage by default. You can choose to store the compilation cache on all of the 4 nodes.

Should we include 2 sections then? One for TPUs and one for GPUs? Can't hurt to mention both options?

docs/persistent_compilation_cache.md

ayaka14732 · 2024-06-18T02:25:17Z

docs/persistent_compilation_cache.md

+
+### Examining cache misses
+
+To merely examine and understand why there are cache misses JAX includes a configuration flag that enables the logging of all cache misses (including persistent compilation cache misses) with their explanations. This can be enabled by setting the following configuration.


This flag is for all cache misses in jax, but only tracing cache miss is implemented right now

I found it quite helpful in debugging an issue I had, maybe it's useful to mention exactly what you mentioned, i.e. that it is indeed a flag for all cache misses wherever JAX is using a cache and that if the issue they're trying to debug is a tracing cache issue that this flag is still useful?

ayaka14732 · 2024-06-18T02:29:54Z

You should split long lines of sentences in to multiple lines, similar to other docs.

ayaka14732 · 2024-06-25T20:42:13Z

ping @keshavb96

keshavb96 · 2024-06-26T02:13:39Z

@ayaka14732 sorry! I got busy working on something else and some medical emergencies, I will push the changes by tomorrow!

keshavb96 · 2024-07-01T17:52:58Z

ping @ayaka14732

ayaka14732 · 2024-07-02T16:59:42Z

docs/persistent_compilation_cache.md

@@ -11,7 +11,7 @@ The compilation cache is enabled when the
 is set. This should be done prior to the first compilation. Set the location as
 follows:

-```
+```python


Should mention that there are 3 ways to set the cache dirctory:

(1) Using environment variable

In shell, before running the script:

export JAX_COMPILATION_CACHE_DIR="/tmp/jax_cache"

Or on the top of the Python script:

import os os.environ["JAX_COMPILATION_CACHE_DIR"] = "/tmp/jax_cache"

(2) Using jax.config.update()

jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")

(3) Using set_cache_dir()

from jax.experimental.compilation_cache import compilation_cache as cc cc.set_cache_dir("/tmp/jax_cache")

The original doc mentions the existance of the function set_cache_dir(), but does not give an example of it. This is bad.

ayaka14732 · 2024-07-02T17:03:41Z

docs/persistent_compilation_cache.md

@@ -67,8 +67,135 @@ Cloud Storage (GCS) bucket. We recommend the following configuration:
 Assuming that `gs://jax-cache` is the GCS bucket, set `cache-location` as
 follows:

-```
+```python
 import jax

 jax.config.update("jax_compilation_cache_dir", "gs://jax-cache")


Should mention that gs:// does not work automatically. It requires etils[epath] to be installed.

Should mention that etils[epath] supports many cloud providers, while GCS is one common choice.

ayaka14732 · 2024-07-02T17:06:42Z

docs/persistent_compilation_cache.md

@@ -27,7 +27,7 @@ is an alternate way of setting `cache-location`.

 `cache-location` can be a directory on the local filesystem. For example:

-```
+```python
 import jax

 jax.config.update("jax_compilation_cache_dir", "/tmp/jax-cache")


"the cache does not have an eviction mechanism implemented" is no longer true. There is indeed an LRU cache eviction mechanism for local filesystems, after #21394

ayaka14732 · 2024-07-02T17:09:59Z

docs/persistent_compilation_cache.md

+When the signature for a function created using the parameters above matches 
+that of a compiled function in the persistent cache the function will not be compiled, 
+but will just be read and deserialized from the persistent cache. 


This pararaph is redundant, because this is just how a cache works. Can just mention that the cache key consists of these items.

ayaka14732 · 2024-07-02T17:12:39Z

docs/persistent_compilation_cache.md

+
+## Different Runtimes
+
+Below we outline some observed behavior of the persistent compilation cache 


Why are these called "observed behavior"? Aren't the behaviors defined in the code?

ayaka14732 · 2024-07-02T17:13:07Z

docs/persistent_compilation_cache.md

+### Single node with single process and single device
+
+This is the simplest runtime and the only process that compiles and writes to the compilation cache is the singular process. 
+The number of devices does not matter to the cache key in this setup, only the type of device does. 
+
+### Single node with single process and multiple devices
+
+Once again the only process that compiles and writes to the compilation cache is the singular proess. 
+The difference between this setup and the previous is that now the number of devices matters in addition to the type of device. 


We should only focus on the number of nodes here. Since there is only one node, only one binary will be produced. Therefore, these two points are essentially the same. The number of devices does not matter here. It is a part of the cache key, which has already been mentioned above.

ayaka14732 · 2024-07-02T17:19:27Z

docs/persistent_compilation_cache.md

+Once again the only process that compiles and writes to the compilation cache is the singular proess. 
+The difference between this setup and the previous is that now the number of devices matters in addition to the type of device. 
+
+### Multiple process and multiple devices (on either single or multiple nodes)


There is a jax flag jax_share_binary_between_hosts and an issue #18819 which might be related to this topic.

ayaka14732 · 2024-07-02T17:25:22Z

docs/persistent_compilation_cache.md

+It can be helpful to examine what exactly is happening with the persistent compilation cache for debugging. 
+While there is no singular canonical way of debugging and examining what's happening in the compilation cache, 
+here are a few suggestions on how to begin.
+


Should mention that users can enable the logging of related source files by placing

import os os.environ["JAX_DEBUG_LOG_MODULES"] = "jax._src.compiler,jax._src.lru_cache"

on the top of the script.

ayaka14732 · 2024-07-02T17:28:12Z

docs/persistent_compilation_cache.md

+## Logging cache activity
+
+It can be helpful to examine what exactly is happening with the persistent compilation cache for debugging. 
+While there is no singular canonical way of debugging and examining what's happening in the compilation cache, 


This sentence sounds passive. I think that there is no "singular canonical way" of debugging anything. Should just remove this sentence.

ayaka14732 · 2024-07-02T17:29:43Z

docs/persistent_compilation_cache.md

+  - The HLO of the function contains a pointer to the custom_partitioning callback, and leads to different cache keys for the same computation across runs. 
+  - In this situation, caching still proceeds, but a different key is produced every time, making the cache ineffective. 
+
+### Working around custom_partitioning


nit: quote `custom_partitioning`

ayaka14732

Let's get this merged first!

keshavb96 added 4 commits June 4, 2024 19:41

persistent compilation cache details

095d946

addressed comments

6f9da97

minor edits

476cd1f

configure caching time threshold

b765547

ayaka14732 self-assigned this Jun 12, 2024

ayaka14732 self-requested a review June 12, 2024 13:37

nouiz reviewed Jun 12, 2024

View reviewed changes

keshavb96 marked this pull request as ready for review June 12, 2024 19:06

shard_map workaround to custom_partitioning

f2759d8

ayaka14732 mentioned this pull request Jun 18, 2024

Docs for persistent compilation cache are superseded #21908

Closed

ayaka14732 requested changes Jun 18, 2024

View reviewed changes

minor edits

cf0b8fd

ayaka14732 reviewed Jul 2, 2024

View reviewed changes

ayaka14732 approved these changes Jul 4, 2024

View reviewed changes

google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Jul 4, 2024

kokoro-team removed the kokoro:force-run label Jul 4, 2024

copybara-service bot merged commit 1e14157 into google:main Jul 4, 2024
16 checks passed

ayaka14732 mentioned this pull request Jul 4, 2024

Update persistent compilation cache doc #22271

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Compilation cache doc #21819

Compilation cache doc #21819

keshavb96 commented Jun 12, 2024 •

edited

Loading

ayaka14732 commented Jun 12, 2024

nouiz Jun 12, 2024

ayaka14732 commented Jun 12, 2024

ayaka14732 commented Jun 15, 2024

keshavb96 commented Jun 18, 2024

ayaka14732 Jun 18, 2024

keshavb96 Jun 26, 2024

ayaka14732 Jun 18, 2024

keshavb96 Jun 26, 2024

ayaka14732 commented Jun 18, 2024

ayaka14732 commented Jun 25, 2024

keshavb96 commented Jun 26, 2024

keshavb96 commented Jul 1, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 Jul 2, 2024

ayaka14732 left a comment


		### Multiple process and multiple devices (on either single or multiple nodes)

		In this runtime the first time a program is run (the persistent cache is cold / empty) all processes will compile, but only the process with rank 0 in the global communication group will write to the persistent cache. In subsequent runs, all processes will attempt to read from the persistent cache, so it is important for the persistent cache to be in a shared file system (eg: NFS) or remote storage (eg: GFS). If the persistent cache is local to rank 0, then all processes except rank 0 will once again compile in subsequent runs as a result of a compilation cache miss.


		### Examining cache misses

		To merely examine and understand why there are cache misses JAX includes a configuration flag that enables the logging of all cache misses (including persistent compilation cache misses) with their explanations. This can be enabled by setting the following configuration.


		## Different Runtimes

		Below we outline some observed behavior of the persistent compilation cache

Compilation cache doc #21819

Compilation cache doc #21819

Conversation

keshavb96 commented Jun 12, 2024 • edited Loading

ayaka14732 commented Jun 12, 2024

Choose a reason for hiding this comment

ayaka14732 commented Jun 12, 2024

ayaka14732 commented Jun 15, 2024

keshavb96 commented Jun 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ayaka14732 commented Jun 18, 2024

ayaka14732 commented Jun 25, 2024

keshavb96 commented Jun 26, 2024

keshavb96 commented Jul 1, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ayaka14732 left a comment

Choose a reason for hiding this comment

keshavb96 commented Jun 12, 2024 •

edited

Loading