Avoid re-computing computation hashes#8976
Merged
rpsilva-aws merged 1 commit intopytorch:masterfrom Apr 15, 2025
Merged
Conversation
9f7b5d1 to
7de630d
Compare
tengyifei
approved these changes
Apr 15, 2025
Collaborator
|
Looks like you need to update some tests that check for these kind of logs. |
7de630d to
1b06c82
Compare
1b06c82 to
fd585f1
Compare
jeffhataws
reviewed
Apr 15, 2025
jeffhataws
approved these changes
Apr 15, 2025
rpsilva-aws
added a commit
to rpsilva-aws/xla
that referenced
this pull request
Apr 15, 2025
rpsilva-aws
added a commit
to rpsilva-aws/xla
that referenced
this pull request
Apr 16, 2025
zpcore
pushed a commit
that referenced
this pull request
Apr 16, 2025
jeffhataws
pushed a commit
that referenced
this pull request
Apr 18, 2025
This file contains hidden or 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
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.
Currently, we are recomputing the hash of the underlying computation for every hash lookup, as a mere logging in two places. For small models where tracing is not negligible, this can have a small impact - particularly since we deserialize the protobuf deterministically (requiring the ordering of unordered dictionary/map entries). The logging was unchanged, but the underlying deserialization logic is relatively slower, in order to guarantee deterministic hashes for user computations. C++'s evaluates stream operators eagerly, so the impact is there with or without the logging levels.
This is only observed if the model is tracing bound. We recently saw ~5% throughput impact for small BERT models.
Note that this is only used to provide an unique hash string for which a hash key maps to. The actual hash of the protobuf is only meaningful for UserComputation computations, where it is factored in the hash key. In all other cases, it is unnecessary and serves as an unique (debug) identifier, and the user can still verify the mapping for any given graph hash key when enabling
post_compilation_analysis.We see this during hash lookup, which is evaluated every time. We also see it in
Compile, though it is there only for the very first computation (across all instances). The user can still access the computation proto hash by enablingPT_XLA_DEBUG.e.g. for BERT HF pretraining (20 steps) - 48 metrics with 27 samples each, the collective tracing of each hash computation metric is as follows: