add axis env state to cache keys, fixes #9187 #9188
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
mattjj
force-pushed
the
cache-sensitivity-to-axis-env-state
branch
from
a34acf6 to
a63993f
Compare
1 task
mattjj
force-pushed
the
cache-sensitivity-to-axis-env-state
branch
2 times, most recently
from
dfa06d2 to
5400420
Compare
mattjj
commented
Jan 13, 2022
mattjj
force-pushed
the
cache-sensitivity-to-axis-env-state
branch
2 times, most recently
from
f0f597f to
240079f
Compare
hawkinsp
reviewed
Jan 14, 2022
mattjj
force-pushed
the
cache-sensitivity-to-axis-env-state
branch
from
240079f to
a7bb498
Compare
mattjj
force-pushed
the
cache-sensitivity-to-axis-env-state
branch
from
a7bb498 to
25bfd6f
Compare
hawkinsp
approved these changes
Jan 18, 2022
| "Promise that a function or with-suite does not depend implicitly on axis env" | ||
| s = thread_local_state.trace_state | ||
| prev_axis_env, s.axis_env = s.axis_env, [] | ||
| jax_config.update_thread_local_jit_state(axis_env_state=()) |
There was a problem hiding this comment.
I wonder if it would be possible to put a dummy value that caused an error if you were to actually try to use it here?
There was a problem hiding this comment.
Such things are possible...
The current approach would also lead to an error, just with a potentially misleading error message (unbound axis name).
There was a problem hiding this comment.
I added a comment to this effect but didn't add any fancy error mechanisms. Gotta leave some work for the future, right?
google-ml-butler
bot
added
kokoro:force-run
pull ready
mattjj
force-pushed
the
cache-sensitivity-to-axis-env-state
branch
from
15e89d5 to
6bac182
Compare
|
Subsumed by 11298. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
fixes #9187
This change could cause some new and even unnecessary cache misses. That is, while the cache hits in #9187 were buggy and should be cache misses, to correct them we may end up being overly defensive.
The reason for defensiveness is that in general, even for a fixed sequence of input abstract values to a jitted function, whenever the named axis environment (which represents the mapped-over names bound by
vmap,pmap, and/orxmap) is different from one we've seen before with those input abstract values, we need to re-trace the Python callable. That's because the Python callable could bind a collective likepsumoraxis_index, and the result of such a collective depends on the axis environment. (When a mapped axis has no name, i.e. when its name iscore.no_axis_name, this error can't happen. So for caching purposes unnamed axes can be ignored.)But some Python callables may not involve collectives at all! It'd be a waste to re-trace those. One such instance arose in this PR: the
_multi_slicefunction, which is used inpmapdispatch, was getting re-traced (and re-compiled) even though it binds no collective primitives. So I introduced a low-level way to, in effect, promise that no collectives are bound: usecore.stash_axis_envto temporarily empty out the axis environment and thus effectively ignore its value.Still, this issue could arise in user code... is the only option to provide such a "promise" API to users?
TODO:
core.stash_axis_env(tried this and ran into a weird__doc__issue...)