Repository navigation
Conversation
PartitionedParameterCoordinator.__params_to_release was decorated with @functools.lru_cache(maxsize=None) on an instance method, which caused three issues: 1. The class-level unbounded cache keyed on (self, submodule_to_release, step_id) held strong references to every PartitionedParameterCoordinator instance and its nn.Module submodules forever, preventing garbage collection after engine.destroy() and leaking partitioned parameter storage across engine lifecycles. 2. _clear_trace_structures() cleared the recorded submodule trace without invalidating the class-level cache, so re-recorded traces (e.g. dynamic control flow or MoE routing changes) continued serving stale parameter release sets from the first trace. 3. release_sub_module() mutated the cached set in-place via params_to_release.update(recompute_params) during backward passes with activation checkpointing, permanently adding recompute parameters to the cached release set for subsequent forward passes at the same (submodule, step_id). Replace the class-level @functools.lru_cache(maxsize=None) with a per-instance __params_to_release_cache dictionary keyed by (submodule_to_release.ds_id, step_id) storing FrozenSet[int], clear it in _clear_trace_structures(), and use a non-mutating set union (params_to_release | recompute_params) in release_sub_module(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera <parthvadera@gmail.com>
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. |
![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: 9a55c4fcde
ℹ️ 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".
| setattr(coordinator, PREFIX + "trace_mode", ZeRoTraceMode.COMPLETE) | ||
| setattr(coordinator, PREFIX + "submodule_order", tuple(order)) | ||
| setattr(coordinator, PREFIX + "step_id", step_id) |
There was a problem hiding this comment.
The retrace and recompute regressions manufacture a completed trace by writing name-mangled coordinator fields (and later hand-edit parameter status), so they never exercise trace_prologue, trace re-recording, or the activation-checkpoint hooks that populate recompute state; a break in the real ZeRO-3 flow would therefore still pass these tests, while an equivalent implementation that renames these fields would fail. Replace these cases with a minimal engine-level dynamic-control-flow/checkpoint run and assert an observable outcome, as the repository requires contract-based integration coverage for training-loop behavior.
AGENTS.md reference: AGENTS.md:L76-L81
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Addressed in 96e7b48b42: updated the test suite to drive ZeRO-3 retrace and recompute parameter-release caches through full engine forward/backward passes rather than mutating internal state directly.
Replace the synthetic coordinator state setup in test_zero_params_to_release_cache.py with engine-level ZeRO-3 runs that exercise trace_prologue invalidation, trace re-recording under dynamic control flow, and activation-checkpoint backward recompute hooks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera <parthvadera@gmail.com>
…#9034) Keep the __params_to_release signature and initial complete-trace guard aligned with upstream master (deepspeedai#9034) while caching FrozenSet[int] per (ds_id, step_id) and clearing any wrapped lru_cache on trace invalidation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera <parthvadera@gmail.com>
gpu-test-labeler
Bot
added
the
skip-gpu-test
Summary
PartitionedParameterCoordinator.__params_to_release(deepspeed/runtime/zero/partitioned_param_coordinator.py) was decorated with@functools.lru_cache(maxsize=None)on an instance method. Because@functools.lru_cacheattaches a single unbounded cache dictionary to the function object on the class and keys entries by(self, submodule_to_release, step_id), this caused three related bugs in ZeRO-3:__params_to_releaseduringZeRoTraceMode.COMPLETEstored strong references toself(PartitionedParameterCoordinator) andsubmodule_to_release(torch.nn.Module) in the class-levellru_cache, preventing garbage collection of the coordinator, submodules, and their ZeRO-3 parameter partitions even afterengine.destroy()and deleting all user references._invalidate_trace()/_clear_trace_structures()cleared__submodule_orderand re-recorded a new execution trace (for example under dynamic control flow or MoE routing changes), the class-levellru_cachewas never cleared. If the same(self, submodule_to_release, step_id)occurred in the new trace,__params_to_releasereturned the stale release set computed from the old__submodule_order, either releasing parameters still needed withinmax_reuse_distance_in_numelor retaining parameters that should be partitioned.release_sub_module(), the backward path executedparams_to_release.update(recompute_params)directly on the mutablesetreturned by the cache, permanently mutating the cached release set for subsequent forward passes at the same(submodule, step_id).Fix
@functools.lru_cache(maxsize=None)on__params_to_releasewith a per-instanceself.__params_to_release_cache: Dict[Tuple[int, int], FrozenSet[int]]keyed by(submodule_to_release.ds_id, step_id).self.__params_to_release_cachein_clear_trace_structures()alongside__submodule_order,__step_id_module_fetched_for, and__param_queue.FrozenSet[int]from__params_to_releaseand replaceparams_to_release.update(recompute_params)withparams_to_release = params_to_release | recompute_paramsinrelease_sub_module().Sibling Execution Paths Checked
DeepSpeedZeroOptimizer) and ZeRO Stage 2 (DeepSpeedZeroOptimizer/stage_1_and_2.py): Parameters are not partitioned across submodules and do not usePartitionedParameterCoordinatoror submodule release traces.DeepSpeedZeroOptimizer_Stage3,DeepSpeedZeRoOffload,PartitionedParameterCoordinator): All sharePartitionedParameterCoordinator.release_sub_moduleand_clear_trace_structures, covered by this fix.Validation & Benchmarks
gloosingle-rankDistributedTest)ds_recompute_parameters(activation checkpointing recompute path)_clear_trace_structures()) followed by trace re-recordingLinear(1024, 1024)model,gc.collect()after each lifecycle):@functools.lru_cache(maxsize=None)): 20/20PartitionedParameterCoordinatorinstances and 80/80Linearsubmodules retained in memory (__params_to_release.cache_info().currsize == 80; RSS grew from 249.6 MB to 888.0 MB, +638.4 MB leaked).__params_to_release_cache+FrozenSet[int]): 0/20 coordinators and 0/80 submodules retained aftergc.collect()(0.0 MB leaked).TestZero3CoordinatorCacheLifecycletests (test_release_sub_module_follows_retraced_order,test_backward_recompute_release_does_not_mutate_cached_forward_release, andtest_engine_coordinator_released_after_complete_trace) pass.Duplicate Check & Acknowledgments
__params_to_release,lru_cache,_clear_trace_structures, andds_recompute_parameters; no existing issue or PR covers the@functools.lru_cache(maxsize=None)retention, trace-invalidation staleness, orparams_to_release.update(recompute_params)cache mutation.PartitionedParameterCoordinatorand__params_to_release(#1901), @inkcherry (Heyang Qin) for theds_recompute_parametersrelease logic during activation checkpointing, and @tohtana (Masahiro Tanaka) and @hwchen2017 (Hongwei Chen) for recent ZeRO-3 trace invalidation and persistent-parameter handling inpartitioned_param_coordinator.py.Disclosure: Drafted and validated with assistance from GitHub Copilot CLI (Claude Sonnet 5.5).