Repository navigation
Conversation
The parameter coordinator scans ahead during prefetch and records the step at which each parameter is needed. Parameters that are already available are skipped by the prefetch, and the release path retained them until that recorded step regardless of stage3_max_reuse_distance. Since the prefetcher can scan far enough ahead to see the backward reuse of the parameters gathered in the forward pass, every forward-gathered parameter stayed resident until backward even when the user configured an aggressive release distance. Rely on the reuse-distance scan to decide which parameters to retain. It already keeps parameters reused within the configured window and releases the rest, so setting stage3_max_reuse_distance to 0 now releases gathered parameters as soon as each module's forward completes and re-gathers them during backward. Also use the future module's own leaf-module flag while scanning the trace so that parameters of leaf modules are counted correctly. Add regression tests covering aggressive release with prefetch enabled, the default retention behavior, and release without prefetching, and document the memory/communication tradeoff of a zero reuse distance. Signed-off-by: YidaWeng <84164537+YidaWeng@users.noreply.github.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: d21ce033e0
ℹ️ 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".
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
The release path now honors stage3_max_reuse_distance for forward-gathered parameters. NVMe parameters need the previous skip-based retention because releasing one also returns its swap-in buffer, and a later on-demand fetch can then fail with no buffer available. Retain skipped NVMe params while leaving the aggressive release behavior for non-offloaded parameters, and add an NVMe offload regression test. Signed-off-by: YidaWeng <84164537+YidaWeng@users.noreply.github.com>
sfc-gh-truwase
approved these changes
Nov 9, 2026
gpu-test-labeler
Bot
added
the
skip-gpu-test
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Forward-gathered parameters were not released when
stage3_max_reuse_distancewas set to a small value. The prefetcher skips parameters that are already available and records the step at which they will be needed next; the release path then retained those parameters until that step regardless of the configured reuse distance. Because the prefetcher can scan far enough ahead to see the backward reuse of parameters gathered in the forward pass, the parameters stayed resident until backward and accumulated across layers (see #7670).The reuse-distance scan already keeps parameters that are reused within the configured window and releases the rest, so the extra prefetch-skip retention is removed in favor of that scan for non-offloaded parameters.
stage3_max_reuse_distance: 0now releases parameters as soon as each module's forward completes and re-gathers them during backward, trading communication for lower forward memory. The trace scan also uses each future module's own leaf-module flag when collecting its parameters.The prefetch-skip retention is deliberately kept for NVMe-offloaded parameters: releasing one also returns its swap-in buffer, and a later on-demand fetch can then fail with no available buffer. This preserves the previous NVMe queue/buffer ordering while keeping aggressive release for parameters that stay on the accelerator.
The default (
1e9) behavior is unchanged: parameters are still kept resident through backward.Testing
tests/unit/runtime/zero/test_zero3_forward_parameter_release.pycovers aggressive release with prefetch enabled (fails before this change), default retention, release with prefetch disabled, and an NVMe-offload training regression test (skipped when async-io is unavailable).test_zero_leaf_module.py,test_zero_allreduce_fetch_params.py,test_zero_context.py, andtest_zeropp.pypass on CPU.torchrun --nproc_per_node=2.pre-commitpasses (flake8 run separately under Python 3.12).Fixes #7670