Synchronize ZeRO destroy before unpinning offload buffers - #8327
Conversation
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cbd91d19cb
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| hook.remove() | ||
| print_rank_0("Removed grad acc hooks", force=False) | ||
| self.ipg_buckets.clear() | ||
| if get_accelerator().is_available(): |
There was a problem hiding this comment.
This is a non-merge commit, but its commit message has no Signed-off-by trailer, so it does not satisfy the repository's mandatory commit requirement and may fail DCO/CI validation. Recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L6-L10
Useful? React with 👍 / 👎.
1 participant
Problem
offload_states(pin_memory=True, non_blocking=True)can queue device-to-host copies into native pinned buffers. The ZeRO destroy paths unpinned those buffers without first ensuring the queued copies had completed, allowing native storage to be released while still in use.Approach
Synchronize the accelerator in the shared ZeRO-1/2 destroy path and the ZeRO-3 destroy path immediately before unpinning optimizer-owned offload buffers. Teardown remains state-free because destroy is terminal. The regression covers ZeRO stages 1, 2, and 3 and repeats destroy to exercise destructor-style cleanup.
Testing
pre-commit run --all-files