Cache AoT compilation result #927
Merged
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
ruomingp
reviewed
Jan 16, 2025
axlearn/common/trainer.py
Outdated
| # excessive recompilations. Note: this could introduce overhead to training due to | ||
| # pre-compilation checks (such as sharding check) that increases the step time for some | ||
| # models. Note that this cache is always disabled at steps when xsc is enabled. | ||
| disable_python_train_step_cache: Optional[bool] = None |
There was a problem hiding this comment.
Let's avoid negative boolean fields, as enable_python_train_step_cache=True is more readable than disable_python_train_step_cache=False.
Suggested change
| disable_python_train_step_cache: Optional[bool] = None | |
| enable_python_train_step_cache: Optional[bool] = None |
axlearn/common/trainer.py
Outdated
Comment on lines
1012
to
1013
| self._compiled_train_step = compiled_jit_train_step_fn | ||
| return self._compiled_train_step |
There was a problem hiding this comment.
Should we leave self._compiled_train_step unchanged when with_xsc is true?
Suggested change
| self._compiled_train_step = compiled_jit_train_step_fn | |
| return self._compiled_train_step | |
| return compiled_jit_train_step_fn |
ruomingp
approved these changes
Jan 16, 2025
apghml
reviewed
Jan 16, 2025
ruomingp
reviewed
Jan 16, 2025
axlearn/common/trainer.py
Outdated
Comment on lines
282
to
283
| if cfg.cache_python_train_step is True: | ||
| raise ValueError("cache_python_train_step cannot be True when xsc is enabled.") |
There was a problem hiding this comment.
Actually cache_python_train_step is still useful even with XSC, since the non-XSC steps will be cached and only the XSC step needs to be recompiled. The XSC step does not run often, so maybe this is OK?
markblee
approved these changes
Jan 16, 2025
axlearn/common/trainer.py
Outdated
| # pre-compilation checks (such as sharding check) that increases the step time for some | ||
| # models. Note that this cache is always disabled at steps when xsc is enabled. | ||
| # Defaults to None which is interpreted as True. | ||
| cache_python_train_step: Optional[bool] = None |
There was a problem hiding this comment.
nit --
Suggested change
| cache_python_train_step: Optional[bool] = None | |
| cache_compiled_train_step: Optional[bool] = None |
markblee
approved these changes
Jan 16, 2025
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.
This improves the step time of some models on v6e by 2x. See the comments of
cache_compiled_train_step.