CheckpointSequential support non-reentrant #86331
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/86331
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit a1fa22e: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Closes #86328 Adds `use_reentrant` argument to `checkpoint_sequential`. [ghstack-poisoned]
Closes #86328 Adds `use_reentrant` argument to `checkpoint_sequential`. [ghstack-poisoned]
albanD
left a comment
albanD
left a comment
There was a problem hiding this comment.
Not sure about some of the doc update. SGTM otherwise.
| the RNG state during each checkpoint. | ||
| Default: ``True`` | ||
| use_reentrant(bool, optional): Use checkpointing | ||
| use_reentrant(bool, optional): Use (the default) checkpointing |
There was a problem hiding this comment.
Not sure what that means? There is a Default: True in this section.
There was a problem hiding this comment.
I added this to indicate that reentrant is the default / original way of checkpointing. Usually the user would not be concerned about the implementation detail of checkpointing so clarifying which flag enables the current "default" seems valuable to me but can take this out if you prefer.
There was a problem hiding this comment.
I think this is already done by the Default: specification (note that we (try to) use this form consistently across the doc. So the user expects to find that information there already.
Closes #86328 Adds `use_reentrant` argument to `checkpoint_sequential`. [ghstack-poisoned]
Summary: Closes #86328 Adds `use_reentrant` argument to `checkpoint_sequential`. Pull Request resolved: #86331 Approved by: https://github.com/zhaojuanmao, https://github.com/albanD Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/7a411952fbb82cec38da936a7d863da49726699f Reviewed By: seemethere Differential Revision: D40167180 Pulled By: rohan-varma fbshipit-source-id: 3bf67bced911af548f4d3c6e65a91017fe03e692
Stack from ghstack (oldest at bottom):
Closes #86328
Adds
use_reentrantargument tocheckpoint_sequential.