Conversation
…kpoint & zero3 & single-model-multi-branch network
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 96e8807dad
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| (output_in, ) = inputs | ||
| ctx.module = module | ||
| if output_in.requires_grad: | ||
| if output_in.requires_grad and torch._C._current_graph_task_id() == -1: # hooks only in really forward phase |
There was a problem hiding this comment.
Preserve the callback for reentrant checkpoint recomputation
When use_reentrant=True and the checkpoint input requires gradients, the initial no-grad pass increments this counter with graph task ID -1, but the recomputation runs inside backward and fails this new condition. Its ctx therefore never receives post_backward_function; nevertheless, PostBackwardFunctionModule.backward decrements the earlier counter to zero and accesses that missing attribute, aborting the ZeRO-3 backward. Assign the callback independently of the guarded increment and cover the reentrant training loop on an actual device.
AGENTS.md reference: AGENTS.md:L35-L36
Useful? React with 👍 / 👎.
| if "ds_grads_remaining" not in module.__dict__: | ||
| module.ds_grads_remaining = 0 |
There was a problem hiding this comment.
Add the required Signed-off-by trailer
This is a non-merge commit, but its commit message has no Signed-off-by trailer, so it violates the repository's mandatory commit requirement and may fail DCO enforcement. Recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
96e8807 to
9b4ffb6
Compare
To fix #8481