Z3: Support for activation ckpt with frozen params#8148
Conversation
sfc-gh-truwase
changed the title
![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: 8a77470bf0
ℹ️ 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".
29c461f to
e72bfd0
Compare
|
Hi @sfc-gh-truwase, Great PR, this one is much simpler than my previous attempt. I tested this PR against several ZeRO-3 activation-checkpoint lifecycle cases. The following two checks failed. In both failed cases, a recomputed frozen parameter remained gathered beyond its final use. The difference is whether that parameter has a backward consumer:
These failures show that the weights remain gathered beyond their final use at the observed point in backward. Actually, my previous attempt only solves the first one. |
|
@tohtana thanks for the reviewing, extended testing, and reporting limitations of the PR. For the two failing cases,
|
|
@tohtana, I have used AI to integrate your UTs as XFAILs. Do you think we still need explicit asserts in zero3 code? |
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
- Release no-grad-input leftovers each microbatch via a coordinator sweep in stage3 backward_epilogue (params whose post-backward hook never fires). - Key active-backward submodules by ds_id (dict, not LIFO deque) so a multi-tensor z3-leaf module releases regardless of hook ordering. - Extend user-backward unit tests; drop the two prior xfail markers. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
eb7a3f8 to
335e93d
Compare
The _dump_params / _dump_param_ids fetch traces were emitting on every forward/backward step; make them opt-in like the other coordinator debug logs. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com>
2 participants
Build on #8130