Preserve ZeRO-3 release output metadata in DeepCompile - #8508
Conversation
![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: b02a6338e0
ℹ️ 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".
| name=f"release_ds_param_{release_node.target}_{node.name}_{ds_id}", | ||
| meta=_make_node_meta(node, ds_id, False)) | ||
| new_node.meta["val"] = None | ||
| new_node.meta["val"] = node.meta["val"] |
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 DCO requirement; recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
b02a633 to
340a236
Compare
1 participant
Problem
The ZeRO-3 compile pass sets
meta["val"]toNoneon insertedrelease_paramnodes, even though the operator returns its input Tensor unchanged.The activation-offload pass checks
meta["val"]to identify floating-point tensors among the outputs saved for backward. Setting it toNonecauses the pass to skip tensors returned byrelease_param.Correcting this metadata exposes a second issue: forward offload records the release node’s name, while backward reload looks up the original saved-output name. These names differ, so the backward pass does not insert the corresponding reload.
Approach
Preserve the input’s exact
meta["val"]object on each release node, retaining its Tensor properties and alias relationships.Also preserve the original saved-output name and use it to match forward offload with backward reload. When undoing an offload, restore the connection to the actual input of the offload node.
Parameter-release side effects, execution order, and graph dependencies remain unchanged.